[8.x] [AI Assistant] Use semantic_text for internal knowledge base (e…

…lastic#186499) (elastic#200243) # Backport This will backport the following commits from `main` to `8.x`: - [[AI Assistant] Use `semantic_text` for internal knowledge base (elastic#186499)](elastic#186499)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)
jloleysens · Nov 14, 2024 · 665fadc · 665fadc
1 parent 9a9c344
commit 665fadc
Show file tree

Hide file tree

Showing 36 changed files with 997 additions and 420 deletions.
diff --git a/...k/packages/kbn-ai-assistant/src/chat/welcome_message_knowledge_base_setup_error_panel.tsx b/...k/packages/kbn-ai-assistant/src/chat/welcome_message_knowledge_base_setup_error_panel.tsx
@@ -37,7 +37,9 @@ export function WelcomeMessageKnowledgeBaseSetupErrorPanel({
 }) {
   const { http } = useKibana().services;
 
-  const modelName = knowledgeBase.status.value?.model_name;
+  const modelId = knowledgeBase.status.value?.endpoint?.service_settings?.model_id;
+  const deploymentState = knowledgeBase.status.value?.model_stats?.deployment_state;
+  const allocationState = knowledgeBase.status.value?.model_stats?.allocation_state;
 
   return (
     <div
@@ -56,48 +58,42 @@ export function WelcomeMessageKnowledgeBaseSetupErrorPanel({
 
           <EuiDescriptionListDescription>
             <ul>
-              {!knowledgeBase.status.value?.deployment_state ? (
+              {!deploymentState ? (
                 <li>
                   <EuiIcon type="alert" color="subdued" />{' '}
                   <FormattedMessage
                     id="xpack.aiAssistant.welcomeMessage.modelIsNotDeployedLabel"
-                    defaultMessage="Model {modelName} is not deployed"
+                    defaultMessage="Model {modelId} is not deployed"
                     values={{
-                      modelName: <EuiCode>{modelName}</EuiCode>,
+                      modelId: <EuiCode>{modelId}</EuiCode>,
                     }}
                   />
                 </li>
               ) : null}
 
-              {knowledgeBase.status.value?.deployment_state &&
-              knowledgeBase.status.value.deployment_state !== 'started' ? (
+              {deploymentState && deploymentState !== 'started' ? (
                 <li>
                   <EuiIcon type="alert" color="subdued" />{' '}
                   <FormattedMessage
                     id="xpack.aiAssistant.welcomeMessage.modelIsNotStartedLabel"
-                    defaultMessage="Deployment state of {modelName} is {deploymentState}"
+                    defaultMessage="Deployment state of {modelId} is {deploymentState}"
                     values={{
-                      modelName: <EuiCode>{modelName}</EuiCode>,
-                      deploymentState: (
-                        <EuiCode>{knowledgeBase.status.value?.deployment_state}</EuiCode>
-                      ),
+                      modelId: <EuiCode>{modelId}</EuiCode>,
+                      deploymentState: <EuiCode>{deploymentState}</EuiCode>,
                     }}
                   />
                 </li>
               ) : null}
 
-              {knowledgeBase.status.value?.allocation_state &&
-              knowledgeBase.status.value.allocation_state !== 'fully_allocated' ? (
+              {allocationState && allocationState !== 'fully_allocated' ? (
                 <li>
                   <EuiIcon type="alert" color="subdued" />{' '}
                   <FormattedMessage
                     id="xpack.aiAssistant.welcomeMessage.modelIsNotFullyAllocatedLabel"
-                    defaultMessage="Allocation state of {modelName} is {allocationState}"
+                    defaultMessage="Allocation state of {modelId} is {allocationState}"
                     values={{
-                      modelName: <EuiCode>{modelName}</EuiCode>,
-                      allocationState: (
-                        <EuiCode>{knowledgeBase.status.value?.allocation_state}</EuiCode>
-                      ),
+                      modelId: <EuiCode>{modelId}</EuiCode>,
+                      allocationState: <EuiCode>{allocationState}</EuiCode>,
                     }}
                   />
                 </li>
@@ -114,9 +110,9 @@ export function WelcomeMessageKnowledgeBaseSetupErrorPanel({
           <FormattedMessage
             id="xpack.aiAssistant.welcomeMessage.div.checkTrainedModelsToLabel"
             defaultMessage="
-                {retryInstallingLink} or check {trainedModelsLink} to ensure {modelName} is deployed and running."
+                {retryInstallingLink} or check {trainedModelsLink} to ensure {modelId} is deployed and running."
             values={{
-              modelName,
+              modelId,
               retryInstallingLink: (
                 <EuiLink
                   data-test-subj="observabilityAiAssistantWelcomeMessageKnowledgeBaseSetupErrorPanelRetryInstallingLink"

diff --git a/x-pack/packages/kbn-ai-assistant/src/hooks/use_knowledge_base.tsx b/x-pack/packages/kbn-ai-assistant/src/hooks/use_knowledge_base.tsx
@@ -5,27 +5,17 @@
  * 2.0.
  */
 import { i18n } from '@kbn/i18n';
-import type {
-  MlDeploymentAllocationState,
-  MlDeploymentState,
-} from '@elastic/elasticsearch/lib/api/types';
 import { useMemo, useState } from 'react';
 import {
   type AbortableAsyncState,
   useAbortableAsync,
+  APIReturnType,
 } from '@kbn/observability-ai-assistant-plugin/public';
 import { useKibana } from './use_kibana';
 import { useAIAssistantAppService } from './use_ai_assistant_app_service';
 
 export interface UseKnowledgeBaseResult {
-  status: AbortableAsyncState<{
-    ready: boolean;
-    enabled: boolean;
-    error?: any;
-    deployment_state?: MlDeploymentState;
-    allocation_state?: MlDeploymentAllocationState;
-    model_name?: string;
-  }>;
+  status: AbortableAsyncState<APIReturnType<'GET /internal/observability_ai_assistant/kb/status'>>;
   isInstalling: boolean;
   installError?: Error;
   install: () => Promise<void>;

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/config.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/config.ts
@@ -9,7 +9,7 @@ import { schema, type TypeOf } from '@kbn/config-schema';
 
 export const config = schema.object({
   enabled: schema.boolean({ defaultValue: true }),
-  modelId: schema.maybe(schema.string()),
+  modelId: schema.maybe(schema.string()), // TODO: Remove
   scope: schema.maybe(schema.oneOf([schema.literal('observability'), schema.literal('search')])),
   enableKnowledgeBase: schema.boolean({ defaultValue: true }),
 });

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/context.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/context.ts
@@ -24,8 +24,8 @@ export function registerContextFunction({
   client,
   functions,
   resources,
-  isKnowledgeBaseAvailable,
-}: FunctionRegistrationParameters & { isKnowledgeBaseAvailable: boolean }) {
+  isKnowledgeBaseReady,
+}: FunctionRegistrationParameters & { isKnowledgeBaseReady: boolean }) {
   functions.registerFunction(
     {
       name: CONTEXT_FUNCTION_NAME,
@@ -54,7 +54,7 @@ export function registerContextFunction({
           ...(dataWithinTokenLimit.length ? { data_on_screen: dataWithinTokenLimit } : {}),
         };
 
-        if (!isKnowledgeBaseAvailable) {
+        if (!isKnowledgeBaseReady) {
           return { content };
         }
 

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/index.ts
@@ -87,7 +87,7 @@ export const registerFunctions: RegistrationCallback = async ({
     );
   }
 
-  const { ready: isReady } = await client.getKnowledgeBaseStatus();
+  const { ready: isKnowledgeBaseReady } = await client.getKnowledgeBaseStatus();
 
   functions.registerInstruction(({ availableFunctionNames }) => {
     const instructions: string[] = [];
@@ -109,7 +109,7 @@ export const registerFunctions: RegistrationCallback = async ({
         Data that is compact enough automatically gets included in the response for the "${CONTEXT_FUNCTION_NAME}" function.`);
     }
 
-    if (isReady) {
+    if (isKnowledgeBaseReady) {
       if (availableFunctionNames.includes(SUMMARIZE_FUNCTION_NAME)) {
         instructions.push(`You can use the "${SUMMARIZE_FUNCTION_NAME}" function to store new information you have learned in a knowledge database.
           Only use this function when the user asks for it.
@@ -129,11 +129,11 @@ export const registerFunctions: RegistrationCallback = async ({
     return instructions.map((instruction) => dedent(instruction));
   });
 
-  if (isReady) {
+  if (isKnowledgeBaseReady) {
     registerSummarizationFunction(registrationParameters);
   }
 
-  registerContextFunction({ ...registrationParameters, isKnowledgeBaseAvailable: isReady });
+  registerContextFunction({ ...registrationParameters, isKnowledgeBaseReady });
 
   registerElasticsearchFunction(registrationParameters);
   const request = registrationParameters.resources.request;

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/plugin.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/plugin.ts
@@ -36,6 +36,7 @@ import { registerFunctions } from './functions';
 import { recallRankingEvent } from './analytics/recall_ranking';
 import { initLangtrace } from './service/client/instrumentation/init_langtrace';
 import { aiAssistantCapabilities } from '../common/capabilities';
+import { registerMigrateKnowledgeBaseEntriesTask } from './service/task_manager_definitions/register_migrate_knowledge_base_entries_task';
 
 export class ObservabilityAIAssistantPlugin
   implements
@@ -114,7 +115,8 @@ export class ObservabilityAIAssistantPlugin
     }) as ObservabilityAIAssistantRouteHandlerResources['plugins'];
 
     // Using once to make sure the same model ID is used during service init and Knowledge base setup
-    const getModelId = once(async () => {
+    const getSearchConnectorModelId = once(async () => {
+      // TODO: Remove this once the modelId is removed from the config
       const configModelId = this.config.modelId;
       if (configModelId) {
         return configModelId;
@@ -156,11 +158,18 @@ export class ObservabilityAIAssistantPlugin
     const service = (this.service = new ObservabilityAIAssistantService({
       logger: this.logger.get('service'),
       core,
-      taskManager: plugins.taskManager,
-      getModelId,
+      getSearchConnectorModelId,
       enableKnowledgeBase: this.config.enableKnowledgeBase,
     }));
 
+    registerMigrateKnowledgeBaseEntriesTask({
+      core,
+      taskManager: plugins.taskManager,
+      logger: this.logger,
+    }).catch((error) => {
+      this.logger.error(`Failed to register migrate knowledge base entries task: ${error}`);
+    });
+
     service.register(registerFunctions);
 
     registerServerRoutes({

diff --git a/...s/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts b/...s/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts
@@ -5,14 +5,16 @@
  * 2.0.
  */
 
-import type {
-  MlDeploymentAllocationState,
-  MlDeploymentState,
-} from '@elastic/elasticsearch/lib/api/types';
 import pLimit from 'p-limit';
 import { notImplemented } from '@hapi/boom';
 import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils';
 import * as t from 'io-ts';
+import {
+  InferenceInferenceEndpointInfo,
+  MlDeploymentAllocationState,
+  MlDeploymentState,
+} from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import moment from 'moment';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
 import { Instruction, KnowledgeBaseEntry, KnowledgeBaseEntryRole } from '../../../common/types';
 
@@ -21,44 +23,86 @@ const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
   options: {
     tags: ['access:ai_assistant'],
   },
-  handler: async (
-    resources
-  ): Promise<{
-    enabled: boolean;
+  handler: async ({
+    service,
+    request,
+  }): Promise<{
+    errorMessage?: string;
     ready: boolean;
-    error?: any;
-    deployment_state?: MlDeploymentState;
-    allocation_state?: MlDeploymentAllocationState;
-    model_name?: string;
+    enabled: boolean;
+    endpoint?: Partial<InferenceInferenceEndpointInfo>;
+    model_stats?: {
+      deployment_state: MlDeploymentState | undefined;
+      allocation_state: MlDeploymentAllocationState | undefined;
+    };
   }> => {
-    const client = await resources.service.getClient({ request: resources.request });
+    const client = await service.getClient({ request });
 
     if (!client) {
       throw notImplemented();
     }
 
-    return await client.getKnowledgeBaseStatus();
+    return client.getKnowledgeBaseStatus();
   },
 });
 
 const setupKnowledgeBase = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
+  params: t.partial({
+    query: t.partial({
+      model_id: t.string,
+    }),
+  }),
   options: {
     tags: ['access:ai_assistant'],
     timeout: {
-      idleSocket: 20 * 60 * 1000, // 20 minutes
+      idleSocket: moment.duration(20, 'minutes').asMilliseconds(),
     },
   },
-  handler: async (resources): Promise<{}> => {
+  handler: async (resources): Promise<InferenceInferenceEndpointInfo> => {
     const client = await resources.service.getClient({ request: resources.request });
 
     if (!client) {
       throw notImplemented();
     }
 
-    await client.setupKnowledgeBase();
+    const { model_id: modelId } = resources.params?.query ?? {};
+
+    return await client.setupKnowledgeBase(modelId);
+  },
+});
+
+const resetKnowledgeBase = createObservabilityAIAssistantServerRoute({
+  endpoint: 'POST /internal/observability_ai_assistant/kb/reset',
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (resources): Promise<{ result: string }> => {
+    const client = await resources.service.getClient({ request: resources.request });
+
+    if (!client) {
+      throw notImplemented();
+    }
+
+    await client.resetKnowledgeBase();
+
+    return { result: 'success' };
+  },
+});
+
+const semanticTextMigrationKnowledgeBase = createObservabilityAIAssistantServerRoute({
+  endpoint: 'POST /internal/observability_ai_assistant/kb/semantic_text_migration',
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (resources): Promise<void> => {
+    const client = await resources.service.getClient({ request: resources.request });
+
+    if (!client) {
+      throw notImplemented();
+    }
 
-    return {};
+    return client.migrateKnowledgeBaseToSemanticText();
   },
 });
 
@@ -225,8 +269,8 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
       throw notImplemented();
     }
 
-    const status = await client.getKnowledgeBaseStatus();
-    if (!status.ready) {
+    const { ready } = await client.getKnowledgeBaseStatus();
+    if (!ready) {
       throw new Error('Knowledge base is not ready');
     }
 
@@ -252,7 +296,9 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
 });
 
 export const knowledgeBaseRoutes = {
+  ...semanticTextMigrationKnowledgeBase,
   ...setupKnowledgeBase,
+  ...resetKnowledgeBase,
   ...getKnowledgeBaseStatus,
   ...getKnowledgeBaseEntries,
   ...saveKnowledgeBaseUserInstruction,

diff --git a/...lugins/observability_solution/observability_ai_assistant/server/routes/register_routes.ts b/...lugins/observability_solution/observability_ai_assistant/server/routes/register_routes.ts
@@ -9,13 +9,14 @@ import type { Logger } from '@kbn/logging';
 import { registerRoutes } from '@kbn/server-route-repository';
 import { getGlobalObservabilityAIAssistantServerRouteRepository } from './get_global_observability_ai_assistant_route_repository';
 import type { ObservabilityAIAssistantRouteHandlerResources } from './types';
+import { ObservabilityAIAssistantPluginStartDependencies } from '../types';
 
 export function registerServerRoutes({
   core,
   logger,
   dependencies,
 }: {
-  core: CoreSetup;
+  core: CoreSetup<ObservabilityAIAssistantPluginStartDependencies>;
   logger: Logger;
   dependencies: Omit<
     ObservabilityAIAssistantRouteHandlerResources,

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/types.ts
@@ -68,6 +68,7 @@ export interface ObservabilityAIAssistantRouteHandlerResources {
 export interface ObservabilityAIAssistantRouteCreateOptions {
   options: {
     timeout?: {
+      payload?: number;
       idleSocket?: number;
     };
     tags: Array<'access:ai_assistant'>;