[Obs AI Assistant] Add KB user instructions

x-pack/plugins/observability_solution/observability_ai_assistant/common/index.ts

@@ -51,3 +51,5 @@ export { DEFAULT_LANGUAGE_OPTION, LANGUAGE_OPTIONS } from './ui_settings/languag
 export { isSupportedConnectorType } from './connectors';
 
 export { ShortIdTable } from './utils/short_id_table';
+
+export { KnowledgeBaseType } from './types';

x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts

@@ -84,6 +84,7 @@ export interface KnowledgeBaseEntry {
   doc_id: string;
   confidence: 'low' | 'medium' | 'high';
   is_correction: boolean;
+  type?: 'user_instruction' | 'contextual';
   public: boolean;
   labels?: Record<string, string>;
   role: KnowledgeBaseEntryRole;
@@ -92,13 +93,18 @@ export interface KnowledgeBaseEntry {
   };
 }
 
-export interface UserInstruction {
+export interface Instruction {
   doc_id: string;
   text: string;
-  system?: boolean;
+  user_instruction?: boolean;
 }
 
-export type UserInstructionOrPlainText = string | UserInstruction;
+export type InstructionOrPlainText = string | Instruction;
+
+export enum KnowledgeBaseType {
+  UserInstruction = 'user_instruction',
+  Contextual = 'contextual',
+}
 
 export interface ObservabilityAIAssistantScreenContextRequest {
   screenDescription?: string;

x-pack/plugins/observability_solution/observability_ai_assistant/public/index.ts

@@ -78,6 +78,8 @@ export type {
   ShortIdTable,
 } from '../common';
 
+export { KnowledgeBaseType } from '../common';
+
 export type { TelemetryEventTypeWithPayload } from './analytics';
 export { ObservabilityAIAssistantTelemetryEventType } from './analytics/telemetry_event_type';
 

x-pack/plugins/observability_solution/observability_ai_assistant/public/types.ts

@@ -18,7 +18,7 @@ import type {
   Message,
   ObservabilityAIAssistantScreenContext,
   PendingMessage,
-  UserInstructionOrPlainText,
+  InstructionOrPlainText,
 } from '../common/types';
 import type { TelemetryEventTypeWithPayload } from './analytics';
 import type { ObservabilityAIAssistantAPIClient } from './api';
@@ -68,7 +68,7 @@ export interface ObservabilityAIAssistantChatService {
         };
     signal: AbortSignal;
     responseLanguage?: string;
-    instructions?: UserInstructionOrPlainText[];
+    instructions?: InstructionOrPlainText[];
   }) => Observable<StreamingChatResponseEventWithoutError>;
   getFunctions: (options?: { contexts?: string[]; filter?: string }) => FunctionDefinition[];
   hasFunction: (name: string) => boolean;

x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts

@@ -5,6 +5,7 @@
  * 2.0.
  */
 
+import { KnowledgeBaseType } from '../../common/types';
 import type { FunctionRegistrationParameters } from '.';
 import { KnowledgeBaseEntryRole } from '../../common';
 
@@ -41,6 +42,11 @@ export function registerSummarizationFunction({
             type: 'boolean',
             description: 'Whether this is a correction for a previous learning.',
           },
+          // type: {
+          //   type: 'string',
+          //   description: `The type can be: "user_instruction" or "contextual". A "user_instruction" entry will be included in the system message if it fits within the token budget and it will be marked as a user instruction. A "contextual" entry will be included as part of the context function response (and not the system message) if it semantically matches the user's prompt.`,
+          //   enum: [KnowledgeBaseType.UserInstruction, KnowledgeBaseType.Contextual],
+          // },
           confidence: {
             type: 'string',
             description: 'How confident you are about this being a correct and useful learning',
@@ -56,6 +62,7 @@ export function registerSummarizationFunction({
           'id' as const,
           'text' as const,
           'is_correction' as const,
+          // 'type' as const,
           'confidence' as const,
           'public' as const,
         ],
@@ -66,13 +73,14 @@ export function registerSummarizationFunction({
       signal
     ) => {
       return client
-        .createKnowledgeBaseEntry({
+        .addKnowledgeBaseEntry({
           entry: {
             doc_id: id,
             role: KnowledgeBaseEntryRole.AssistantSummarization,
             id,
             text,
             is_correction: isCorrection,
+            type: KnowledgeBaseType.Contextual,
             confidence,
             public: isPublic,
             labels: {},

x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/chat/route.ts

@@ -49,7 +49,7 @@ const chatCompleteBaseRt = t.type({
               text: t.string,
             }),
             t.partial({
-              system: t.boolean,
+              user_instruction: t.boolean,
             }),
           ]),
         ])

x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts

@@ -33,15 +33,15 @@ const getFunctionsRoute = createObservabilityAIAssistantServerRoute({
 
     const client = await service.getClient({ request });
 
-    const [functionClient, userInstructions] = await Promise.all([
+    const [functionClient, kbUserInstructions] = await Promise.all([
       service.getFunctionClient({
         signal: controller.signal,
         resources,
         client,
         screenContexts: [],
       }),
       // error is caught in client
-      client.fetchUserInstructions(),
+      client.getKnowledgeBaseUserInstructions(),
     ]);
 
     const functionDefinitions = functionClient.getFunctions().map((fn) => fn.definition);
@@ -52,7 +52,7 @@ const getFunctionsRoute = createObservabilityAIAssistantServerRoute({
       functionDefinitions: functionClient.getFunctions().map((fn) => fn.definition),
       systemMessage: getSystemMessageFromInstructions({
         registeredInstructions: functionClient.getInstructions(),
-        userInstructions,
+        kbUserInstructions,
         requestInstructions: [],
         availableFunctionNames,
       }),
@@ -111,6 +111,7 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
       text: nonEmptyStringRt,
       confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]),
       is_correction: toBooleanRt,
+      type: t.union([t.literal('user_instruction'), t.literal('contextual')]),
       public: toBooleanRt,
       labels: t.record(t.string, t.string),
     }),
@@ -129,17 +130,19 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
       confidence,
       id,
       is_correction: isCorrection,
+      type,
       text,
       public: isPublic,
       labels,
     } = resources.params.body;
 
-    return client.createKnowledgeBaseEntry({
+    return client.addKnowledgeBaseEntry({
       entry: {
         confidence,
         id,
         doc_id: id,
         is_correction: isCorrection,
+        type,
         text,
         public: isPublic,
         labels,

x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts

@@ -100,6 +100,7 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
       t.partial({
         confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]),
         is_correction: toBooleanRt,
+        type: t.union([t.literal('user_instruction'), t.literal('contextual')]),
         public: toBooleanRt,
         labels: t.record(t.string, t.string),
         role: t.union([
@@ -126,17 +127,19 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
       public: isPublic,
       confidence,
       is_correction: isCorrection,
+      type,
       labels,
       role,
     } = resources.params.body;
 
-    return client.createKnowledgeBaseEntry({
+    return client.addKnowledgeBaseEntry({
       entry: {
         id,
         text,
         doc_id: id,
         confidence: confidence ?? 'high',
         is_correction: isCorrection ?? false,
+        type: type ?? 'contextual',
         public: isPublic ?? true,
         labels: labels ?? {},
         role: (role as KnowledgeBaseEntryRole) ?? KnowledgeBaseEntryRole.UserEntry,
@@ -192,6 +195,7 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
       doc_id: entry.id,
       confidence: 'high' as KnowledgeBaseEntry['confidence'],
       is_correction: false,
+      type: 'contextual' as const,
       public: true,
       labels: {},
       role: KnowledgeBaseEntryRole.UserEntry,

x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts

@@ -112,7 +112,7 @@ describe('Observability AI Assistant client', () => {
 
   const knowledgeBaseServiceMock: DeeplyMockedKeys<KnowledgeBaseService> = {
     recall: jest.fn(),
-    getUserInstructions: jest.fn(),
+    getUserInstructions: jest.fn(), // TODO: change
   } as any;
 
   let loggerMock: DeeplyMockedKeys<Logger> = {} as any;

x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts

@@ -46,7 +46,7 @@ import {
 } from '../../../common/conversation_complete';
 import { CompatibleJSONSchema } from '../../../common/functions/types';
 import {
-  UserInstructionOrPlainText,
+  InstructionOrPlainText,
   type Conversation,
   type ConversationCreateRequest,
   type ConversationUpdateRequest,
@@ -159,7 +159,21 @@ export class ObservabilityAIAssistantClient {
     });
   };
 
-  complete = (params: {
+  complete = ({
+    functionClient,
+    connectorId,
+    simulateFunctionCalling,
+    instructions: requestInstructions = [],
+    messages: initialMessages,
+    signal,
+    responseLanguage = 'English',
+    persist,
+    kibanaPublicUrl,
+    isPublic,
+    title: predefinedTitle,
+    conversationId: predefinedConversationId,
+    disableFunctions = false,
+  }: {
     messages: Message[];
     connectorId: string;
     signal: AbortSignal;
@@ -170,7 +184,7 @@ export class ObservabilityAIAssistantClient {
     title?: string;
     isPublic?: boolean;
     kibanaPublicUrl?: string;
-    instructions?: UserInstructionOrPlainText[];
+    instructions?: InstructionOrPlainText[];
     simulateFunctionCalling?: boolean;
     disableFunctions?:
       | boolean
@@ -181,22 +195,6 @@ export class ObservabilityAIAssistantClient {
     return new LangTracer(context.active()).startActiveSpan(
       'complete',
       ({ tracer: completeTracer }) => {
-        const {
-          functionClient,
-          connectorId,
-          simulateFunctionCalling,
-          instructions: requestInstructions = [],
-          messages: initialMessages,
-          signal,
-          responseLanguage = 'English',
-          persist,
-          kibanaPublicUrl,
-          isPublic,
-          title: predefinedTitle,
-          conversationId: predefinedConversationId,
-          disableFunctions = false,
-        } = params;
-
         if (responseLanguage) {
           requestInstructions.push(
             `You MUST respond in the users preferred language which is: ${responseLanguage}.`
@@ -215,17 +213,19 @@ export class ObservabilityAIAssistantClient {
           );
         }
 
-        const userInstructions$ = from(this.fetchUserInstructions()).pipe(shareReplay());
+        const kbUserInstructions$ = from(this.getKnowledgeBaseUserInstructions()).pipe(
+          shareReplay()
+        );
 
         // from the initial messages, override any system message with
         // the one that is based on the instructions (registered, request, kb)
-        const messagesWithUpdatedSystemMessage$ = userInstructions$.pipe(
-          map((userInstructions) => {
+        const messagesWithUpdatedSystemMessage$ = kbUserInstructions$.pipe(
+          map((kbUserInstructions) => {
             // this is what we eventually store in the conversation
             const messagesWithUpdatedSystemMessage = replaceSystemMessage(
               getSystemMessageFromInstructions({
                 registeredInstructions: functionClient.getInstructions(),
-                userInstructions,
+                kbUserInstructions,
                 requestInstructions,
                 availableFunctionNames: functionClient
                   .getFunctions()
@@ -272,9 +272,9 @@ export class ObservabilityAIAssistantClient {
         // messages and the knowledge base instructions
         const nextEvents$ = combineLatest([
           messagesWithUpdatedSystemMessage$,
-          userInstructions$,
+          kbUserInstructions$,
         ]).pipe(
-          switchMap(([messagesWithUpdatedSystemMessage, userInstructions]) => {
+          switchMap(([messagesWithUpdatedSystemMessage, kbUserInstructions]) => {
             // if needed, inject a context function request here
             const contextRequest = functionClient.hasFunction(CONTEXT_FUNCTION_NAME)
               ? getContextFunctionRequestIfNeeded(messagesWithUpdatedSystemMessage)
@@ -302,7 +302,7 @@ export class ObservabilityAIAssistantClient {
                 // start out with the max number of function calls
                 functionCallsLeft: MAX_FUNCTION_CALLS,
                 functionClient,
-                userInstructions,
+                kbUserInstructions,
                 requestInstructions,
                 signal,
                 logger: this.dependencies.logger,
@@ -731,7 +731,7 @@ export class ObservabilityAIAssistantClient {
     return this.dependencies.knowledgeBaseService.setup();
   };
 
-  createKnowledgeBaseEntry = async ({
+  addKnowledgeBaseEntry = async ({
     entry,
   }: {
     entry: Omit<KnowledgeBaseEntry, '@timestamp'>;
@@ -772,7 +772,7 @@ export class ObservabilityAIAssistantClient {
     return this.dependencies.knowledgeBaseService.deleteEntry({ id });
   };
 
-  fetchUserInstructions = async () => {
+  getKnowledgeBaseUserInstructions = async () => {
     return this.dependencies.knowledgeBaseService.getUserInstructions(
       this.dependencies.namespace,
       this.dependencies.user

x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/continue_conversation.ts

@@ -28,7 +28,7 @@ import {
   MessageOrChatEvent,
 } from '../../../../common/conversation_complete';
 import { FunctionVisibility } from '../../../../common/functions/types';
-import { UserInstruction } from '../../../../common/types';
+import { Instruction } from '../../../../common/types';
 import { createFunctionResponseMessage } from '../../../../common/utils/create_function_response_message';
 import { emitWithConcatenatedMessage } from '../../../../common/utils/emit_with_concatenated_message';
 import { withoutTokenCountEvents } from '../../../../common/utils/without_token_count_events';
@@ -172,7 +172,7 @@ export function continueConversation({
   signal,
   functionCallsLeft,
   requestInstructions,
-  userInstructions,
+  kbUserInstructions,
   logger,
   disableFunctions,
   tracer,
@@ -182,8 +182,8 @@ export function continueConversation({
   chat: ChatFunctionWithoutConnector;
   signal: AbortSignal;
   functionCallsLeft: number;
-  requestInstructions: Array<string | UserInstruction>;
-  userInstructions: UserInstruction[];
+  requestInstructions: Array<string | Instruction>;
+  kbUserInstructions: Instruction[];
   logger: Logger;
   disableFunctions:
     | boolean
@@ -205,7 +205,7 @@ export function continueConversation({
   const messagesWithUpdatedSystemMessage = replaceSystemMessage(
     getSystemMessageFromInstructions({
       registeredInstructions: functionClient.getInstructions(),
-      userInstructions,
+      kbUserInstructions,
       requestInstructions,
       availableFunctionNames: definitions.map((def) => def.name),
     }),
@@ -324,7 +324,7 @@ export function continueConversation({
               functionCallsLeft: nextFunctionCallsLeft,
               functionClient,
               signal,
-              userInstructions,
+              kbUserInstructions,
               requestInstructions,
               logger,
               disableFunctions,