From 8ab30f5fcd73822d36565920dec3bc6ed663b153 Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Thu, 31 Oct 2024 15:34:22 -0600 Subject: [PATCH] [Security assistant] Fix `invoke_assistant_*` telemetry (#198594) --- .../routes/chat/chat_complete_route.test.ts | 1 + .../server/routes/chat/chat_complete_route.ts | 21 +++++++++++-- .../server/routes/helpers.ts | 30 +++++++------------ .../server/routes/knowledge_base/constants.ts | 2 ++ .../post_actions_connector_execute.test.ts | 2 +- .../routes/post_actions_connector_execute.ts | 7 ++--- 6 files changed, 37 insertions(+), 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 index 4aca370aa700f..f03a3394cdaac 100644 --- 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 index f2365f0320967..c6eb81dd86ebd 100644 --- 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,6 +222,19 @@ 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, @@ -226,6 +242,7 @@ export const chatCompleteRoute = ( // 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 index bf12a99a0afbb..d25ed5fc77f10 100644 --- 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(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 => { + 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 index 8bf17027e751e..052b2cac57609 100644 --- 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 index 998790f332c45..a7abac27dac6f 100644 --- 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 index acf2dd32a060b..bb217f7f5aa3a 100644 --- 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({