From 8a7e8c8c1cb0b2986093155700029e61e9643b6b Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Fri, 13 Sep 2024 17:17:08 -0600 Subject: [PATCH] [Security GenAI] Re-add telemetry for `isEnabledKnowledgeBase` (#192785) (cherry picked from commit cabaf7a0774cda8775c059c36b7db7679ca19d8a) --- .../chat_send/use_chat_send.test.tsx | 21 +++++++--- .../assistant/chat_send/use_chat_send.tsx | 11 +++++ .../impl/assistant_context/types.tsx | 1 + .../setup_knowledge_base_button.tsx | 2 +- .../lib/telemetry/event_based_telemetry.ts | 14 +++++++ .../server/routes/helpers.ts | 41 ++++++++++++++++++- .../routes/post_actions_connector_execute.ts | 17 ++++++++ .../telemetry/events/ai_assistant/index.ts | 6 +++ .../telemetry/events/ai_assistant/types.ts | 1 + 9 files changed, 106 insertions(+), 8 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx index 97d95f641e2b5c..0de7adc484fc1d 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.test.tsx @@ -61,10 +61,11 @@ describe('use chat send', () => { }); it('handleOnChatCleared clears the conversation', async () => { (clearConversation as jest.Mock).mockReturnValueOnce(testProps.currentConversation); - const { result } = renderHook(() => useChatSend(testProps), { + const { result, waitForNextUpdate } = renderHook(() => useChatSend(testProps), { wrapper: TestProviders, }); - await act(async () => { + await waitForNextUpdate(); + act(() => { result.current.handleOnChatCleared(); }); expect(clearConversation).toHaveBeenCalled(); @@ -95,7 +96,7 @@ describe('use chat send', () => { }); }); it('handleRegenerateResponse removes the last message of the conversation, resends the convo to GenAI, and appends the message received', async () => { - const { result } = renderHook( + const { result, waitForNextUpdate } = renderHook( () => useChatSend({ ...testProps, currentConversation: { ...welcomeConvo, id: 'welcome-id' } }), { @@ -103,7 +104,10 @@ describe('use chat send', () => { } ); - result.current.handleRegenerateResponse(); + await waitForNextUpdate(); + act(() => { + result.current.handleRegenerateResponse(); + }); expect(removeLastMessage).toHaveBeenCalledWith('welcome-id'); await waitFor(() => { @@ -114,10 +118,13 @@ describe('use chat send', () => { }); it('sends telemetry events for both user and assistant', async () => { const promptText = 'prompt text'; - const { result } = renderHook(() => useChatSend(testProps), { + const { result, waitForNextUpdate } = renderHook(() => useChatSend(testProps), { wrapper: TestProviders, }); - result.current.handleChatSend(promptText); + await waitForNextUpdate(); + act(() => { + result.current.handleChatSend(promptText); + }); await waitFor(() => { expect(reportAssistantMessageSent).toHaveBeenNthCalledWith(1, { @@ -126,6 +133,7 @@ describe('use chat send', () => { actionTypeId: '.gen-ai', model: undefined, provider: 'OpenAI', + isEnabledKnowledgeBase: false, }); expect(reportAssistantMessageSent).toHaveBeenNthCalledWith(2, { conversationId: testProps.currentConversation?.title, @@ -133,6 +141,7 @@ describe('use chat send', () => { actionTypeId: '.gen-ai', model: undefined, provider: 'OpenAI', + isEnabledKnowledgeBase: false, }); }); }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx index 9b671bca64a2ee..95de527ff7a3a0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx @@ -9,6 +9,8 @@ import React, { useCallback, useState } from 'react'; import { HttpSetup } from '@kbn/core-http-browser'; import { i18n } from '@kbn/i18n'; import { Replacements } from '@kbn/elastic-assistant-common'; +import { useKnowledgeBaseStatus } from '../api/knowledge_base/use_knowledge_base_status'; +import { ESQL_RESOURCE } from '../../knowledge_base/setup_knowledge_base_button'; import { DataStreamApis } from '../use_data_stream_apis'; import { NEW_CHAT } from '../conversations/conversation_sidepanel/translations'; import type { ClientMessage } from '../../assistant_context/types'; @@ -56,6 +58,12 @@ export const useChatSend = ({ const { isLoading, sendMessage, abortStream } = useSendMessage(); const { clearConversation, removeLastMessage } = useConversation(); + const { data: kbStatus } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); + const isSetupComplete = + kbStatus?.elser_exists && + kbStatus?.index_exists && + kbStatus?.pipeline_exists && + kbStatus?.esql_exists; // Handles sending latest user prompt to API const handleSendMessage = useCallback( @@ -116,6 +124,7 @@ export const useChatSend = ({ actionTypeId: currentConversation.apiConfig.actionTypeId, model: currentConversation.apiConfig.model, provider: currentConversation.apiConfig.provider, + isEnabledKnowledgeBase: isSetupComplete ?? false, }); const responseMessage: ClientMessage = getMessageFromRawResponse(rawResponse); @@ -131,12 +140,14 @@ export const useChatSend = ({ actionTypeId: currentConversation.apiConfig.actionTypeId, model: currentConversation.apiConfig.model, provider: currentConversation.apiConfig.provider, + isEnabledKnowledgeBase: isSetupComplete ?? false, }); }, [ assistantTelemetry, currentConversation, http, + isSetupComplete, selectedPromptContexts, sendMessage, setCurrentConversation, diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx index c49ab38452dc0d..dad5ef04e0c188 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx @@ -52,6 +52,7 @@ export interface AssistantTelemetry { actionTypeId: string; model?: string; provider?: string; + isEnabledKnowledgeBase: boolean; }) => void; reportAssistantQuickPrompt: (params: { conversationId: string; promptTitle: string }) => void; reportAssistantSettingToggled: (params: { assistantStreamingEnabled?: boolean }) => void; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx index f9566ff8e89bcf..008597c1e82436 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/setup_knowledge_base_button.tsx @@ -13,7 +13,7 @@ import { useAssistantContext } from '../..'; import { useSetupKnowledgeBase } from '../assistant/api/knowledge_base/use_setup_knowledge_base'; import { useKnowledgeBaseStatus } from '../assistant/api/knowledge_base/use_knowledge_base_status'; -const ESQL_RESOURCE = 'esql'; +export const ESQL_RESOURCE = 'esql'; /** * Self-contained component that renders a button to set up the knowledge base. diff --git a/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts b/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts index b27105abf8180b..4ba95896d70583 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts @@ -73,6 +73,7 @@ export const KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT: EventTypeOpts<{ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{ assistantStreamingEnabled: boolean; actionTypeId: string; + isEnabledKnowledgeBase: boolean; model?: string; }> = { eventType: 'invoke_assistant_success', @@ -96,12 +97,19 @@ export const INVOKE_ASSISTANT_SUCCESS_EVENT: EventTypeOpts<{ optional: true, }, }, + isEnabledKnowledgeBase: { + type: 'boolean', + _meta: { + description: 'Is knowledge base enabled', + }, + }, }, }; export const INVOKE_ASSISTANT_ERROR_EVENT: EventTypeOpts<{ errorMessage: string; assistantStreamingEnabled: boolean; + isEnabledKnowledgeBase: boolean; actionTypeId: string; model?: string; }> = { @@ -132,6 +140,12 @@ export const INVOKE_ASSISTANT_ERROR_EVENT: EventTypeOpts<{ optional: true, }, }, + isEnabledKnowledgeBase: { + type: 'boolean', + _meta: { + description: 'Is knowledge base enabled', + }, + }, }, }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts index cdfa3ad26293f1..d457f9c88bf690 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -28,11 +28,12 @@ import { AwaitedProperties, PublicMethodsOf } from '@kbn/utility-types'; import { ActionsClient } from '@kbn/actions-plugin/server'; import { AssistantFeatureKey } from '@kbn/elastic-assistant-common/impl/capabilities'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { AIAssistantKnowledgeBaseDataClient } from '../ai_assistant_data_clients/knowledge_base'; 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_RESOURCE } from './knowledge_base/constants'; +import { ESQL_DOCS_LOADED_QUERY, ESQL_RESOURCE } from './knowledge_base/constants'; import { buildResponse, getLlmType } from './utils'; import { AgentExecutorParams, @@ -442,12 +443,15 @@ export const langChainExecute = async ({ executorParams ); + const { esqlExists, isModelDeployed } = await getIsKnowledgeBaseEnabled(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, }); return response.ok(result); }; @@ -652,3 +656,38 @@ export const isV2KnowledgeBaseEnabled = ({ }); return context.elasticAssistant.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; }; + +/** + * Telemetry function to determine whether knowledge base has been installed + * @param kbDataClient + */ +export const getIsKnowledgeBaseEnabled = async ( + kbDataClient?: AIAssistantKnowledgeBaseDataClient | null +): Promise<{ + esqlExists: boolean; + isModelDeployed: boolean; +}> => { + let esqlExists = false; + let isModelDeployed = false; + if (kbDataClient != null) { + try { + isModelDeployed = await kbDataClient.isModelDeployed(); + if (isModelDeployed) { + esqlExists = + ( + await kbDataClient.getKnowledgeBaseDocumentEntries({ + query: ESQL_DOCS_LOADED_QUERY, + required: true, + }) + ).length > 0; + } + } catch (e) { + /* if telemetry related requests fail, fallback to default values */ + } + } + + return { + esqlExists, + isModelDeployed, + }; +}; 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 38df73b7d25b6d..0988d9e5f89739 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 @@ -23,6 +23,9 @@ import { buildResponse } from '../lib/build_response'; import { ElasticAssistantRequestHandlerContext, GetElser } from '../types'; import { appendAssistantMessageToConversation, + DEFAULT_PLUGIN_NAME, + getIsKnowledgeBaseEnabled, + getPluginNameFromRequest, getSystemPromptFromUserConversation, langChainExecute, } from './helpers'; @@ -144,11 +147,25 @@ export const postActionsConnectorExecuteRoute = ( if (onLlmResponse) { await onLlmResponse(error.message, {}, true); } + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const v2KnowledgeBaseEnabled = + assistantContext.getRegisteredFeatures(pluginName).assistantKnowledgeBaseByDefault; + const kbDataClient = + (await assistantContext.getAIAssistantKnowledgeBaseDataClient( + v2KnowledgeBaseEnabled + )) ?? undefined; + const isEnabledKnowledgeBase = await getIsKnowledgeBaseEnabled(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, }); return resp.error({ diff --git a/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/index.ts b/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/index.ts index 482a38a983f4e6..117d6216ed2abf 100644 --- a/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/index.ts +++ b/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/index.ts @@ -66,6 +66,12 @@ export const assistantMessageSentEvent: TelemetryEvent = { optional: true, }, }, + isEnabledKnowledgeBase: { + type: 'boolean', + _meta: { + description: 'Is knowledge base enabled', + }, + }, }, }; diff --git a/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/types.ts b/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/types.ts index ee879d516885d3..2dd6bf6215dbf0 100644 --- a/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/types.ts +++ b/x-pack/plugins/security_solution/public/common/lib/telemetry/events/ai_assistant/types.ts @@ -19,6 +19,7 @@ export interface ReportAssistantMessageSentParams { actionTypeId: string; provider?: string; model?: string; + isEnabledKnowledgeBase: boolean; } export interface ReportAssistantQuickPromptParams {