[8.x] [Security GenAI] Re-add telemetry for `isEnabledKnowledgeB…

…ase` (#192785) (#192937) # Backport This will backport the following commits from `main` to `8.x`: - [[Security GenAI] Re-add telemetry for `isEnabledKnowledgeBase` (#192785)](#192785)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Steph Milovic <[email protected]>
elastic · Sep 14, 2024 · 0c66ca8 · 0c66ca8
1 parent eef194b
commit 0c66ca8
Show file tree

Hide file tree

Showing 9 changed files with 106 additions and 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
@@ -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,15 +96,18 @@ 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' } }),
       {
         wrapper: TestProviders,
       }
     );
 
-    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,13 +133,15 @@ describe('use chat send', () => {
         actionTypeId: '.gen-ai',
         model: undefined,
         provider: 'OpenAI',
+        isEnabledKnowledgeBase: false,
       });
       expect(reportAssistantMessageSent).toHaveBeenNthCalledWith(2, {
         conversationId: testProps.currentConversation?.title,
         role: 'assistant',
         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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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<StreamResponseWithHeaders['body'] | StaticReturnType['body']>(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
@@ -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
@@ -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
@@ -19,6 +19,7 @@ export interface ReportAssistantMessageSentParams {
   actionTypeId: string;
   provider?: string;
   model?: string;
+  isEnabledKnowledgeBase: boolean;
 }
 
 export interface ReportAssistantQuickPromptParams {