[Security assistant] Fix invoke_assistant_* telemetry (elastic#198594)

mikecote · Oct 31, 2024 · 8ab30f5 · 8ab30f5
1 parent 0ad421b
commit 8ab30f5
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 26 deletions.
diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.test.ts
@@ -290,6 +290,7 @@ describe('chatCompleteRoute', () => {
                 actionTypeId: '.gen-ai',
                 model: 'gpt-4',
                 assistantStreamingEnabled: false,
+                isEnabledKnowledgeBase: false,
               });
             }),
           };

diff --git a/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts
@@ -25,6 +25,9 @@ import { buildResponse } from '../../lib/build_response';
 import {
   appendAssistantMessageToConversation,
   createConversationWithUserInput,
+  DEFAULT_PLUGIN_NAME,
+  getIsKnowledgeBaseInstalled,
+  getPluginNameFromRequest,
   langChainExecute,
   performChecks,
 } from '../helpers';
@@ -63,9 +66,9 @@ export const chatCompleteRoute = (
         const assistantResponse = buildResponse(response);
         let telemetry;
         let actionTypeId;
+        const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
+        const logger: Logger = ctx.elasticAssistant.logger;
         try {
-          const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
-          const logger: Logger = ctx.elasticAssistant.logger;
           telemetry = ctx.elasticAssistant.telemetry;
           const inference = ctx.elasticAssistant.inference;
 
@@ -219,13 +222,27 @@ export const chatCompleteRoute = (
           });
         } catch (err) {
           const error = transformError(err as Error);
+          const pluginName = getPluginNameFromRequest({
+            request,
+            defaultPluginName: DEFAULT_PLUGIN_NAME,
+            logger,
+          });
+          const v2KnowledgeBaseEnabled =
+            ctx.elasticAssistant.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault;
+          const kbDataClient =
+            (await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient({
+              v2KnowledgeBaseEnabled,
+            })) ?? undefined;
+          const isKnowledgeBaseInstalled = await getIsKnowledgeBaseInstalled(kbDataClient);
+
           telemetry?.reportEvent(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
             actionTypeId: actionTypeId ?? '',
             model: request.body.model,
             errorMessage: error.message,
             // TODO rm actionTypeId check when llmClass for bedrock streaming is implemented
             // tracked here: https://github.com/elastic/security-team/issues/7363
             assistantStreamingEnabled: request.body.isStream ?? false,
+            isEnabledKnowledgeBase: isKnowledgeBaseInstalled,
           });
           return assistantResponse.error({
             body: error.message,

diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts
@@ -35,7 +35,7 @@ import { FindResponse } from '../ai_assistant_data_clients/find';
 import { EsPromptsSchema } from '../ai_assistant_data_clients/prompts/types';
 import { AIAssistantDataClient } from '../ai_assistant_data_clients';
 import { MINIMUM_AI_ASSISTANT_LICENSE } from '../../common/constants';
-import { ESQL_DOCS_LOADED_QUERY } from './knowledge_base/constants';
+import { SECURITY_LABS_RESOURCE, SECURITY_LABS_LOADED_QUERY } from './knowledge_base/constants';
 import { buildResponse, getLlmType } from './utils';
 import {
   AgentExecutorParams,
@@ -436,15 +436,13 @@ export const langChainExecute = async ({
     executorParams
   );
 
-  const { esqlExists, isModelDeployed } = await getIsKnowledgeBaseEnabled(kbDataClient);
+  const isKnowledgeBaseInstalled = await getIsKnowledgeBaseInstalled(kbDataClient);
 
   telemetry.reportEvent(INVOKE_ASSISTANT_SUCCESS_EVENT.eventType, {
     actionTypeId,
     model: request.body.model,
-    // TODO rm actionTypeId check when llmClass for bedrock streaming is implemented
-    // tracked here: https://github.com/elastic/security-team/issues/7363
-    assistantStreamingEnabled: isStream && actionTypeId === '.gen-ai',
-    isEnabledKnowledgeBase: isModelDeployed && esqlExists,
+    assistantStreamingEnabled: isStream,
+    isEnabledKnowledgeBase: isKnowledgeBaseInstalled,
   });
   return response.ok<StreamResponseWithHeaders['body'] | StaticReturnType['body']>(result);
 };
@@ -671,23 +669,20 @@ export const isV2KnowledgeBaseEnabled = ({
  * Telemetry function to determine whether knowledge base has been installed
  * @param kbDataClient
  */
-export const getIsKnowledgeBaseEnabled = async (
+export const getIsKnowledgeBaseInstalled = async (
   kbDataClient?: AIAssistantKnowledgeBaseDataClient | null
-): Promise<{
-  esqlExists: boolean;
-  isModelDeployed: boolean;
-}> => {
-  let esqlExists = false;
+): Promise<boolean> => {
+  let securityLabsDocsExist = false;
   let isModelDeployed = false;
   if (kbDataClient != null) {
     try {
       isModelDeployed = await kbDataClient.isModelDeployed();
       if (isModelDeployed) {
-        esqlExists =
+        securityLabsDocsExist =
           (
             await kbDataClient.getKnowledgeBaseDocumentEntries({
-              query: ESQL_DOCS_LOADED_QUERY,
-              required: true,
+              kbResource: SECURITY_LABS_RESOURCE,
+              query: SECURITY_LABS_LOADED_QUERY,
             })
           ).length > 0;
       }
@@ -696,8 +691,5 @@ export const getIsKnowledgeBaseEnabled = async (
     }
   }
 
-  return {
-    esqlExists,
-    isModelDeployed,
-  };
+  return isModelDeployed && securityLabsDocsExist;
 };
diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/constants.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/constants.ts
@@ -13,3 +13,5 @@ export const ESQL_DOCS_LOADED_QUERY =
   'You can chain processing commands, separated by a pipe character: `|`.';
 export const SECURITY_LABS_RESOURCE = 'security_labs';
 export const USER_RESOURCE = 'user';
+// Query for determining if Security Labs docs have been loaded. Intended for use with Telemetry
+export const SECURITY_LABS_LOADED_QUERY = 'What is Elastic Security Labs';
diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts
@@ -51,7 +51,7 @@ jest.mock('./helpers', () => {
 
   return {
     ...original,
-    getIsKnowledgeBaseEnabled: jest.fn(),
+    getIsKnowledgeBaseInstalled: jest.fn(),
     appendAssistantMessageToConversation: jest.fn(),
     langChainExecute: jest.fn(),
     getPluginNameFromRequest: jest.fn(),

diff --git a/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts
@@ -24,7 +24,7 @@ import { ElasticAssistantRequestHandlerContext, GetElser } from '../types';
 import {
   appendAssistantMessageToConversation,
   DEFAULT_PLUGIN_NAME,
-  getIsKnowledgeBaseEnabled,
+  getIsKnowledgeBaseInstalled,
   getPluginNameFromRequest,
   getSystemPromptFromUserConversation,
   langChainExecute,
@@ -170,14 +170,13 @@ export const postActionsConnectorExecuteRoute = (
             (await assistantContext.getAIAssistantKnowledgeBaseDataClient({
               v2KnowledgeBaseEnabled,
             })) ?? undefined;
-          const isEnabledKnowledgeBase = await getIsKnowledgeBaseEnabled(kbDataClient);
-
+          const isKnowledgeBaseInstalled = await getIsKnowledgeBaseInstalled(kbDataClient);
           telemetry.reportEvent(INVOKE_ASSISTANT_ERROR_EVENT.eventType, {
             actionTypeId: request.body.actionTypeId,
             model: request.body.model,
             errorMessage: error.message,
             assistantStreamingEnabled: request.body.subAction !== 'invokeAI',
-            isEnabledKnowledgeBase,
+            isEnabledKnowledgeBase: isKnowledgeBaseInstalled,
           });
 
           return resp.error({