Fix KB API

elastic · Feb 16, 2025 · fddf4fa · fddf4fa
1 parent e9206f9
commit fddf4fa
Show file tree

Hide file tree

Showing 24 changed files with 122 additions and 103 deletions.
diff --git a/...curity/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/...curity/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts
@@ -443,10 +443,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
     const created =
       docsCreated.length > 0
         ? await this.findDocuments<EsKnowledgeBaseEntrySchema>({
-            page: 1,
-            perPage: 10000,
-            filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
-          })
+          page: 1,
+          perPage: 10000,
+          filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
+        })
         : undefined;
     // Intentionally no telemetry here - this path only used to install security docs
     // Plans to make this function private in a different PR so no user entry ever is created in this path

diff --git a/...curity/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts b/...curity/plugins/elastic_assistant/server/routes/anonymization_fields/bulk_actions_route.ts
@@ -165,7 +165,7 @@ export const bulkActionAnonymizationFieldsRoute = (
         try {
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
           // Perform license and authenticated user checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,
@@ -218,10 +218,10 @@ export const bulkActionAnonymizationFieldsRoute = (
           const created =
             docsCreated.length > 0
               ? await dataClient?.findDocuments<EsAnonymizationFieldsSchema>({
-                  page: 1,
-                  perPage: 1000,
-                  filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
-                })
+                page: 1,
+                perPage: 1000,
+                filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
+              })
               : undefined;
 
           return buildBulkResponse(response, {

diff --git a/...tions/security/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts b/...tions/security/plugins/elastic_assistant/server/routes/anonymization_fields/find_route.ts
@@ -58,7 +58,7 @@ export const findAnonymizationFieldsRoute = (
           const { query } = request;
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
           // Perform license and authenticated user checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,

diff --git a/...ck/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/...ck/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts
@@ -78,7 +78,7 @@ export const chatCompleteRoute = (
             (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false;
 
           // Perform license and authenticated user checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,

diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/get_evaluate.ts
@@ -49,7 +49,7 @@ export const getEvaluateRoute = (router: IRouter<ElasticAssistantRequestHandlerC
         const logger = assistantContext.logger.get('evaluate');
 
         // Perform license, authenticated user and evaluation FF checks
-        const checkResponse = performChecks({
+        const checkResponse = await performChecks({
           capability: 'assistantModelEvaluation',
           context: ctx,
           request,

diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
@@ -103,7 +103,7 @@ export const postEvaluateRoute = (
         const savedObjectsClient = ctx.elasticAssistant.savedObjectsClient;
 
         // Perform license, authenticated user and evaluation FF checks
-        const checkResponse = performChecks({
+        const checkResponse = await performChecks({
           capability: 'assistantModelEvaluation',
           context: ctx,
           request,
@@ -275,20 +275,20 @@ export const postEvaluateRoute = (
               // Skeleton request from route to pass to the agents
               // params will be passed to the actions executor
               const skeletonRequest: KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody> =
-                {
-                  ...request,
-                  body: {
-                    alertsIndexPattern: '',
-                    allow: [],
-                    allowReplacement: [],
-                    subAction: 'invokeAI',
-                    // The actionTypeId is irrelevant when used with the invokeAI subaction
-                    actionTypeId: '.gen-ai',
-                    replacements: {},
-                    size: DEFAULT_SIZE,
-                    conversationId: '',
-                  },
-                };
+              {
+                ...request,
+                body: {
+                  alertsIndexPattern: '',
+                  allow: [],
+                  allowReplacement: [],
+                  subAction: 'invokeAI',
+                  // The actionTypeId is irrelevant when used with the invokeAI subaction
+                  actionTypeId: '.gen-ai',
+                  replacements: {},
+                  size: DEFAULT_SIZE,
+                  conversationId: '',
+                },
+              };
 
               const contentReferencesEnabled =
                 assistantContext.getRegisteredFeatures(
@@ -362,19 +362,19 @@ export const postEvaluateRoute = (
               const agentRunnable =
                 isOpenAI || llmType === 'inference'
                   ? await createOpenAIToolsAgent({
-                      llm,
-                      tools,
-                      prompt: formatPrompt(defaultSystemPrompt),
-                      streamRunnable: false,
-                    })
+                    llm,
+                    tools,
+                    prompt: formatPrompt(defaultSystemPrompt),
+                    streamRunnable: false,
+                  })
                   : llmType && ['bedrock', 'gemini'].includes(llmType)
-                  ? createToolCallingAgent({
+                    ? createToolCallingAgent({
                       llm,
                       tools,
                       prompt: formatPrompt(defaultSystemPrompt),
                       streamRunnable: false,
                     })
-                  : await createStructuredChatAgent({
+                    : await createStructuredChatAgent({
                       llm,
                       tools,
                       prompt: formatPromptStructured(defaultSystemPrompt),

diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts
@@ -433,19 +433,19 @@ interface PerformChecksParams {
 
 type PerformChecks =
   | {
-      isSuccess: true;
-      currentUser: AuthenticatedUser;
-    }
+    isSuccess: true;
+    currentUser: AuthenticatedUser;
+  }
   | {
-      isSuccess: false;
-      response: IKibanaResponse;
-    };
-export const performChecks = ({
+    isSuccess: false;
+    response: IKibanaResponse;
+  };
+export const performChecks = async ({
   capability,
   context,
   request,
   response,
-}: PerformChecksParams): PerformChecks => {
+}: PerformChecksParams): Promise<PerformChecks> => {
   const assistantResponse = buildResponse(response);
 
   if (!hasAIAssistantLicense(context.licensing.license)) {
@@ -459,7 +459,7 @@ export const performChecks = ({
     };
   }
 
-  const currentUser = context.elasticAssistant.getCurrentUser();
+  const currentUser = await context.elasticAssistant.getCurrentUser();
 
   if (currentUser == null) {
     return {

diff --git a/...rity/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts b/...rity/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts
@@ -200,7 +200,7 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug
           const logger = ctx.elasticAssistant.logger;
 
           // Perform license, authenticated user and FF checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,
@@ -311,10 +311,10 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug
           const created =
             docsCreated.length > 0
               ? await kbDataClient?.findDocuments<EsKnowledgeBaseEntrySchema>({
-                  page: 1,
-                  perPage: 100,
-                  filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
-                })
+                page: 1,
+                perPage: 100,
+                filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
+              })
               : undefined;
 
           return buildBulkResponse(

diff --git a/...s/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts b/...s/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts
@@ -48,7 +48,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
           const logger = ctx.elasticAssistant.logger;
 
           // Perform license, authenticated user and FF checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,

diff --git a/...s/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/delete_route.ts b/...s/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/delete_route.ts
@@ -52,7 +52,7 @@ export const deleteKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
           const logger = ctx.elasticAssistant.logger;
 
           // Perform license, authenticated user and FF checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,

diff --git a/...ons/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts b/...ons/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts
@@ -59,7 +59,7 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
 
           // Perform license, authenticated user and FF checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,
@@ -117,8 +117,8 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout
                 (
                   (result?.data.aggregations?.global_aggs as estypes.AggregationsGlobalAggregate)
                     ?.kb_resource_aggregation as {
-                    buckets: estypes.AggregationsBuckets;
-                  }
+                      buckets: estypes.AggregationsBuckets;
+                    }
                 )?.buckets,
                 ['key', kbResource]
               ) as {
@@ -132,22 +132,22 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout
                 entry == null
                   ? []
                   : [
-                      {
-                        id: bucketId,
-                        createdAt: entry.created_at,
-                        createdBy: entry.created_by,
-                        updatedAt: entry.updated_at,
-                        updatedBy: entry.updated_by,
-                        users: [],
-                        name,
-                        namespace: entry.namespace,
-                        type: DocumentEntryType.value,
-                        kbResource,
-                        source: '',
-                        required,
-                        text: `${entryCount}`,
-                      },
-                    ];
+                    {
+                      id: bucketId,
+                      createdAt: entry.created_at,
+                      createdBy: entry.created_by,
+                      updatedAt: entry.updated_at,
+                      updatedBy: entry.updated_by,
+                      users: [],
+                      name,
+                      namespace: entry.namespace,
+                      type: DocumentEntryType.value,
+                      kbResource,
+                      source: '',
+                      required,
+                      text: `${entryCount}`,
+                    },
+                  ];
               return entries;
             })
             .flat();

diff --git a/...s/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/update_route.ts b/...s/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/update_route.ts
@@ -50,7 +50,7 @@ export const updateKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout
           const logger = ctx.elasticAssistant.logger;
 
           // Perform license, authenticated user and FF checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,

diff --git a/...utions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/...utions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts
@@ -70,7 +70,7 @@ export const postActionsConnectorExecuteRoute = (
         let onLlmResponse;
 
         try {
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,

diff --git a/.../solutions/security/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts b/.../solutions/security/plugins/elastic_assistant/server/routes/prompts/bulk_actions_route.ts
@@ -159,7 +159,7 @@ export const bulkPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L
         try {
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
           // Perform license and authenticated user checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,
@@ -213,10 +213,10 @@ export const bulkPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L
           const created =
             docsCreated.length > 0
               ? await dataClient?.findDocuments<EsPromptsSchema>({
-                  page: 1,
-                  perPage: 100,
-                  filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
-                })
+                page: 1,
+                perPage: 100,
+                filter: docsCreated.map((c) => `_id:${c}`).join(' OR '),
+              })
               : undefined;
 
           return buildBulkResponse(response, {

diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/prompts/find_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/prompts/find_route.ts
@@ -47,7 +47,7 @@ export const findPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L
           const { query } = request;
           const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
           // Perform license and authenticated user checks
-          const checkResponse = performChecks({
+          const checkResponse = await performChecks({
             context: ctx,
             request,
             response,
@@ -64,8 +64,8 @@ export const findPromptsRoute = (router: ElasticAssistantPluginRouter, logger: L
             sortOrder: query.sort_order,
             filter: query.filter
               ? `${decodeURIComponent(
-                  query.filter
-                )} and not (prompt_type: "system" and is_default: true)`
+                query.filter
+              )} and not (prompt_type: "system" and is_default: true)`
               : 'not (prompt_type: "system" and is_default: true)',
             fields: query.fields,
           });