[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,26 @@ export interface KnowledgeBaseEntry {
   };
 }
 
-export interface UserInstruction {
+export interface Instruction {
   doc_id: string;
   text: string;
-  system?: boolean;
 }
 
-export type UserInstructionOrPlainText = string | UserInstruction;
+export interface AdHocInstruction {
+  doc_id?: string;
+  text: string;
+  instruction_type: 'user_instruction' | 'application_instruction';
+}
+
+export type InstructionOrPlainText = string | Instruction;
+
+export enum KnowledgeBaseType {
+  // user instructions are included in the system prompt regardless of the user's input
+  UserInstruction = 'user_instruction',
+
+  // contextual entries are only included in the system prompt if the user's input matches the context
+  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,
+  AdHocInstruction,
 } 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?: AdHocInstruction[];
   }) => 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';
 
@@ -66,13 +67,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

@@ -41,17 +41,15 @@ const chatCompleteBaseRt = t.type({
         }),
       ]),
       instructions: t.array(
-        t.union([
-          t.string,
-          t.intersection([
-            t.type({
-              doc_id: t.string,
-              text: t.string,
-            }),
-            t.partial({
-              system: t.boolean,
-            }),
-          ]),
+        t.intersection([
+          t.partial({ doc_id: t.string }),
+          t.type({
+            text: t.string,
+            instruction_type: t.union([
+              t.literal('user_instruction'),
+              t.literal('application_instruction'),
+            ]),
+          }),
         ])
       ),
     }),

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

@@ -41,7 +41,7 @@ const getFunctionsRoute = createObservabilityAIAssistantServerRoute({
         screenContexts: [],
       }),
       // error is caught in client
-      client.fetchUserInstructions(),
+      client.getKnowledgeBaseUserInstructions(),
     ]);
 
     const functionDefinitions = functionClient.getFunctions().map((fn) => fn.definition);
@@ -51,9 +51,9 @@ const getFunctionsRoute = createObservabilityAIAssistantServerRoute({
     return {
       functionDefinitions: functionClient.getFunctions().map((fn) => fn.definition),
       systemMessage: getSystemMessageFromInstructions({
-        registeredInstructions: functionClient.getInstructions(),
+        applicationInstructions: functionClient.getInstructions(),
         userInstructions,
-        requestInstructions: [],
+        adHocInstructions: [],
         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

@@ -13,7 +13,12 @@ import { notImplemented } from '@hapi/boom';
 import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils';
 import * as t from 'io-ts';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
-import { KnowledgeBaseEntry, KnowledgeBaseEntryRole } from '../../../common/types';
+import {
+  Instruction,
+  KnowledgeBaseEntry,
+  KnowledgeBaseEntryRole,
+  KnowledgeBaseType,
+} from '../../../common/types';
 
 const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
   endpoint: 'GET /internal/observability_ai_assistant/kb/status',
@@ -60,6 +65,64 @@ const setupKnowledgeBase = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const getKnowledgeBaseUserInstructions = createObservabilityAIAssistantServerRoute({
+  endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (
+    resources
+  ): Promise<{
+    userInstructions: Array<Instruction & { public?: boolean }>;
+  }> => {
+    const client = await resources.service.getClient({ request: resources.request });
+
+    if (!client) {
+      throw notImplemented();
+    }
+
+    return {
+      userInstructions: await client.getKnowledgeBaseUserInstructions(),
+    };
+  },
+});
+
+const saveKnowledgeBaseUserInstruction = createObservabilityAIAssistantServerRoute({
+  endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
+  params: t.type({
+    body: t.type({
+      id: t.string,
+      text: nonEmptyStringRt,
+      public: toBooleanRt,
+    }),
+  }),
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (resources): Promise<void> => {
+    const client = await resources.service.getClient({ request: resources.request });
+
+    if (!client) {
+      throw notImplemented();
+    }
+
+    const { id, text, public: isPublic } = resources.params.body;
+    return client.addKnowledgeBaseEntry({
+      entry: {
+        id,
+        doc_id: id,
+        text,
+        public: isPublic,
+        confidence: 'high',
+        is_correction: false,
+        type: KnowledgeBaseType.UserInstruction,
+        labels: {},
+        role: KnowledgeBaseEntryRole.UserEntry,
+      },
+    });
+  },
+});
+
 const getKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
   endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
   options: {
@@ -130,13 +193,14 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
       role,
     } = resources.params.body;
 
-    return client.createKnowledgeBaseEntry({
+    return client.addKnowledgeBaseEntry({
       entry: {
         id,
         text,
         doc_id: id,
         confidence: confidence ?? 'high',
         is_correction: isCorrection ?? false,
+        type: 'contextual',
         public: isPublic ?? true,
         labels: labels ?? {},
         role: (role as KnowledgeBaseEntryRole) ?? KnowledgeBaseEntryRole.UserEntry,
@@ -192,6 +256,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,
@@ -206,6 +271,8 @@ export const knowledgeBaseRoutes = {
   ...setupKnowledgeBase,
   ...getKnowledgeBaseStatus,
   ...getKnowledgeBaseEntries,
+  ...saveKnowledgeBaseUserInstruction,
+  ...getKnowledgeBaseUserInstructions,
   ...importKnowledgeBaseEntries,
   ...saveKnowledgeBaseEntry,
   ...deleteKnowledgeBaseEntry,

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

@@ -16,7 +16,7 @@ import type {
   FunctionCallChatFunction,
   FunctionHandler,
   FunctionHandlerRegistry,
-  RegisteredInstruction,
+  InstructionOrCallback,
   RegisterFunction,
   RegisterInstruction,
 } from '../types';
@@ -34,7 +34,7 @@ const ajv = new Ajv({
 export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
 
 export class ChatFunctionClient {
-  private readonly instructions: RegisteredInstruction[] = [];
+  private readonly instructions: InstructionOrCallback[] = [];
   private readonly functionRegistry: FunctionHandlerRegistry = new Map();
   private readonly validators: Map<string, ValidateFunction> = new Map();
 
@@ -107,7 +107,7 @@ export class ChatFunctionClient {
     }
   }
 
-  getInstructions(): RegisteredInstruction[] {
+  getInstructions(): InstructionOrCallback[] {
     return this.instructions;
   }
 

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

@@ -26,15 +26,14 @@ import { createFunctionResponseMessage } from '../../../common/utils/create_func
 import { CONTEXT_FUNCTION_NAME } from '../../functions/context';
 import { ChatFunctionClient } from '../chat_function_client';
 import type { KnowledgeBaseService } from '../knowledge_base_service';
-import { USER_INSTRUCTIONS_HEADER } from '../util/get_system_message_from_instructions';
 import { observableIntoStream } from '../util/observable_into_stream';
 import { CreateChatCompletionResponseChunk } from './adapters/process_openai_stream';
 
 type ChunkDelta = CreateChatCompletionResponseChunk['choices'][number]['delta'];
 
 type LlmSimulator = ReturnType<typeof createLlmSimulator>;
 
-const EXPECTED_STORED_SYSTEM_MESSAGE = `system\n\n${USER_INSTRUCTIONS_HEADER}\n\nYou MUST respond in the users preferred language which is: English.`;
+const EXPECTED_STORED_SYSTEM_MESSAGE = `system\n\nYou MUST respond in the users preferred language which is: English.`;
 
 const nextTick = () => {
   return new Promise(process.nextTick);
@@ -367,8 +366,8 @@ describe('Observability AI Assistant client', () => {
               last_updated: expect.any(String),
               token_count: {
                 completion: 1,
-                prompt: 84,
-                total: 85,
+                prompt: 46,
+                total: 47,
               },
             },
             type: StreamingChatResponseEventType.ConversationCreate,
@@ -424,8 +423,8 @@ describe('Observability AI Assistant client', () => {
               last_updated: expect.any(String),
               token_count: {
                 completion: 6,
-                prompt: 268,
-                total: 274,
+                prompt: 230,
+                total: 236,
               },
             },
             type: StreamingChatResponseEventType.ConversationCreate,
@@ -442,8 +441,8 @@ describe('Observability AI Assistant client', () => {
                 title: 'An auto-generated title',
                 token_count: {
                   completion: 6,
-                  prompt: 268,
-                  total: 274,
+                  prompt: 230,
+                  total: 236,
                 },
               },
               labels: {},
@@ -573,8 +572,8 @@ describe('Observability AI Assistant client', () => {
           last_updated: expect.any(String),
           token_count: {
             completion: 2,
-            prompt: 162,
-            total: 164,
+            prompt: 124,
+            total: 126,
           },
         },
         type: StreamingChatResponseEventType.ConversationUpdate,
@@ -592,8 +591,8 @@ describe('Observability AI Assistant client', () => {
             title: 'My stored conversation',
             token_count: {
               completion: 2,
-              prompt: 162,
-              total: 164,
+              prompt: 124,
+              total: 126,
             },
           },
           labels: {},
@@ -1609,7 +1608,10 @@ describe('Observability AI Assistant client', () => {
       .subscribe(() => {}); // To trigger call to chat
     await nextTick();
 
-    expect(chatSpy.mock.calls[0][1].messages[0].message.content).toEqual(
+    const systemMessage = chatSpy.mock.calls[0][1].messages[0];
+
+    expect(systemMessage.message.role).toEqual(MessageRole.System);
+    expect(systemMessage.message.content).toEqual(
       EXPECTED_STORED_SYSTEM_MESSAGE.replace('English', 'Orcish')
     );
   });