[Obs AI Assistant] Register and retrieve get_data_on_screen function …

…description as an adhoc instruction (elastic#184214)
viduni94 · Oct 30, 2024 · 9b6cbae · 9b6cbae
1 parent 0ce486a
commit 9b6cbae
Show file tree

Hide file tree

Showing 10 changed files with 93 additions and 17 deletions.
diff --git a/...lugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts b/...lugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts
@@ -68,7 +68,7 @@ const getFunctionsRoute = createObservabilityAIAssistantServerRoute({
       systemMessage: getSystemMessageFromInstructions({
         applicationInstructions: functionClient.getInstructions(),
         userInstructions,
-        adHocInstructions: [],
+        adHocInstructions: functionClient.getAdhocInstructions(),
         availableFunctionNames,
       }),
     };

diff --git a/...ity_solution/observability_ai_assistant/server/service/chat_function_client/index.test.ts b/...ity_solution/observability_ai_assistant/server/service/chat_function_client/index.test.ts
@@ -7,6 +7,7 @@
 import dedent from 'dedent';
 import { ChatFunctionClient, GET_DATA_ON_SCREEN_FUNCTION_NAME } from '.';
 import { FunctionVisibility } from '../../../common/functions/types';
+import { AdHocInstruction } from '../../../common/types';
 
 describe('chatFunctionClient', () => {
   describe('when executing a function with invalid arguments', () => {
@@ -86,6 +87,7 @@ describe('chatFunctionClient', () => {
       ]);
 
       const functions = client.getFunctions();
+      const adHocInstructions = client.getAdhocInstructions();
 
       expect(functions[0]).toEqual({
         definition: {
@@ -97,7 +99,7 @@ describe('chatFunctionClient', () => {
         respond: expect.any(Function),
       });
 
-      expect(functions[0].definition.description).toContain(
+      expect(adHocInstructions[0].text).toContain(
         dedent(`my_dummy_data: My dummy data
         my_other_dummy_data: My other dummy data
         `)
@@ -128,4 +130,52 @@ describe('chatFunctionClient', () => {
       });
     });
   });
+
+  describe('when adhoc instructions are provided', () => {
+    let client: ChatFunctionClient;
+
+    beforeEach(() => {
+      client = new ChatFunctionClient([]);
+    });
+
+    describe('register an adhoc Instruction', () => {
+      it('should register a new adhoc instruction', () => {
+        const adhocInstruction: AdHocInstruction = {
+          text: 'Test adhoc instruction',
+          instruction_type: 'application_instruction',
+        };
+
+        client.registerAdhocInstruction(adhocInstruction);
+
+        expect(client.getAdhocInstructions()).toContainEqual(adhocInstruction);
+      });
+    });
+
+    describe('retrieve adHoc instructions', () => {
+      it('should return all registered adhoc instructions', () => {
+        const firstAdhocInstruction: AdHocInstruction = {
+          text: 'First adhoc instruction',
+          instruction_type: 'application_instruction',
+        };
+
+        const secondAdhocInstruction: AdHocInstruction = {
+          text: 'Second adhoc instruction',
+          instruction_type: 'application_instruction',
+        };
+
+        client.registerAdhocInstruction(firstAdhocInstruction);
+        client.registerAdhocInstruction(secondAdhocInstruction);
+
+        const adhocInstructions = client.getAdhocInstructions();
+
+        expect(adhocInstructions).toEqual([firstAdhocInstruction, secondAdhocInstruction]);
+      });
+
+      it('should return an empty array if no adhoc instructions are registered', () => {
+        const adhocInstructions = client.getAdhocInstructions();
+
+        expect(adhocInstructions).toEqual([]);
+      });
+    });
+  });
 });
diff --git a/...vability_solution/observability_ai_assistant/server/service/chat_function_client/index.ts b/...vability_solution/observability_ai_assistant/server/service/chat_function_client/index.ts
@@ -10,13 +10,18 @@ import Ajv, { type ErrorObject, type ValidateFunction } from 'ajv';
 import dedent from 'dedent';
 import { compact, keyBy } from 'lodash';
 import { FunctionVisibility, type FunctionResponse } from '../../../common/functions/types';
-import type { Message, ObservabilityAIAssistantScreenContextRequest } from '../../../common/types';
+import type {
+  AdHocInstruction,
+  Message,
+  ObservabilityAIAssistantScreenContextRequest,
+} from '../../../common/types';
 import { filterFunctionDefinitions } from '../../../common/utils/filter_function_definitions';
 import type {
   FunctionCallChatFunction,
   FunctionHandler,
   FunctionHandlerRegistry,
   InstructionOrCallback,
+  RegisterAdHocInstruction,
   RegisterFunction,
   RegisterInstruction,
 } from '../types';
@@ -35,6 +40,8 @@ export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
 
 export class ChatFunctionClient {
   private readonly instructions: InstructionOrCallback[] = [];
+  private readonly adhocInstructions: AdHocInstruction[] = [];
+
   private readonly functionRegistry: FunctionHandlerRegistry = new Map();
   private readonly validators: Map<string, ValidateFunction> = new Map();
 
@@ -49,9 +56,7 @@ export class ChatFunctionClient {
       this.registerFunction(
         {
           name: GET_DATA_ON_SCREEN_FUNCTION_NAME,
-          description: dedent(`Get data that is on the screen:
-            ${allData.map((data) => `${data.name}: ${data.description}`).join('\n')}
-          `),
+          description: dedent('Get data that is on the screen'),
           visibility: FunctionVisibility.AssistantOnly,
           parameters: {
             type: 'object',
@@ -75,6 +80,13 @@ export class ChatFunctionClient {
           };
         }
       );
+
+      this.registerAdhocInstruction({
+        text: `The ${GET_DATA_ON_SCREEN_FUNCTION_NAME} function will ${dedent(`Get data that is on the screen:
+          ${allData.map((data) => `${data.name}: ${data.description}`).join('\n')}
+        `)}`,
+        instruction_type: 'application_instruction',
+      });
     }
 
     this.actions.forEach((action) => {
@@ -95,6 +107,10 @@ export class ChatFunctionClient {
     this.instructions.push(instruction);
   };
 
+  registerAdhocInstruction: RegisterAdHocInstruction = (instruction: AdHocInstruction) => {
+    this.adhocInstructions.push(instruction);
+  };
+
   validate(name: string, parameters: unknown) {
     const validator = this.validators.get(name)!;
     if (!validator) {
@@ -111,6 +127,10 @@ export class ChatFunctionClient {
     return this.instructions;
   }
 
+  getAdhocInstructions(): AdHocInstruction[] {
+    return this.adhocInstructions;
+  }
+
   hasAction(name: string) {
     return !!this.actions.find((action) => action.name === name)!;
   }

diff --git a/...lity_solution/observability_ai_assistant/server/service/client/adapters/openai_adapter.ts b/...lity_solution/observability_ai_assistant/server/service/client/adapters/openai_adapter.ts
@@ -16,7 +16,6 @@ import { eventsourceStreamIntoObservable } from '../../util/eventsource_stream_i
 import { LlmApiAdapterFactory } from './types';
 import { parseInlineFunctionCalls } from './simulate_function_calling/parse_inline_function_calls';
 import { getMessagesWithSimulatedFunctionCalling } from './simulate_function_calling/get_messages_with_simulated_function_calling';
-import { GET_DATA_ON_SCREEN_FUNCTION_NAME } from '../../chat_function_client';
 
 function getOpenAIPromptTokenCount({
   messages,
@@ -124,10 +123,7 @@ export const createOpenAiAdapter: LlmApiAdapterFactory = ({
           ...(!!functionsForOpenAI?.length
             ? {
                 tools: functionsForOpenAI.map((fn) => ({
-                  function:
-                    fn.name === GET_DATA_ON_SCREEN_FUNCTION_NAME
-                      ? pick(fn, 'name', 'parameters')
-                      : pick(fn, 'name', 'description', 'parameters'),
+                  function: pick(fn, 'name', 'description', 'parameters'),
                   type: 'function',
                 })),
               }

diff --git a/...ins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts b/...ins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts
@@ -124,6 +124,7 @@ describe('Observability AI Assistant client', () => {
     getActions: jest.fn(),
     validate: jest.fn(),
     getInstructions: jest.fn(),
+    getAdhocInstructions: jest.fn(),
   } as any;
 
   let llmSimulator: LlmSimulator;

diff --git a/.../plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts b/.../plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts
@@ -47,12 +47,12 @@ import {
 } from '../../../common/conversation_complete';
 import { CompatibleJSONSchema } from '../../../common/functions/types';
 import {
+  AdHocInstruction,
   type Conversation,
   type ConversationCreateRequest,
   type ConversationUpdateRequest,
   type KnowledgeBaseEntry,
   type Message,
-  type AdHocInstruction,
 } from '../../../common/types';
 import { withoutTokenCountEvents } from '../../../common/utils/without_token_count_events';
 import { CONTEXT_FUNCTION_NAME } from '../../functions/context';
@@ -210,6 +210,9 @@ export class ObservabilityAIAssistantClient {
 
         const userInstructions$ = from(this.getKnowledgeBaseUserInstructions()).pipe(shareReplay());
 
+        const registeredAdhocInstructions = functionClient.getAdhocInstructions();
+        const allAdHocInstructions = adHocInstructions.concat(registeredAdhocInstructions);
+
         // from the initial messages, override any system message with
         // the one that is based on the instructions (registered, request, kb)
         const messagesWithUpdatedSystemMessage$ = userInstructions$.pipe(
@@ -219,7 +222,7 @@ export class ObservabilityAIAssistantClient {
               getSystemMessageFromInstructions({
                 applicationInstructions: functionClient.getInstructions(),
                 userInstructions,
-                adHocInstructions,
+                adHocInstructions: allAdHocInstructions,
                 availableFunctionNames: functionClient
                   .getFunctions()
                   .map((fn) => fn.definition.name),

diff --git a/...ution/observability_ai_assistant/server/service/client/operators/continue_conversation.ts b/...ution/observability_ai_assistant/server/service/client/operators/continue_conversation.ts
@@ -178,7 +178,7 @@ export function continueConversation({
   chat,
   signal,
   functionCallsLeft,
-  adHocInstructions,
+  adHocInstructions = [],
   userInstructions,
   logger,
   disableFunctions,
@@ -213,11 +213,14 @@ export function continueConversation({
     disableFunctions,
   });
 
+  const registeredAdhocInstructions = functionClient.getAdhocInstructions();
+  const allAdHocInstructions = adHocInstructions.concat(registeredAdhocInstructions);
+
   const messagesWithUpdatedSystemMessage = replaceSystemMessage(
     getSystemMessageFromInstructions({
       applicationInstructions: functionClient.getInstructions(),
       userInstructions,
-      adHocInstructions,
+      adHocInstructions: allAdHocInstructions,
       availableFunctionNames: definitions.map((def) => def.name),
     }),
     initialMessages

diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/types.ts
@@ -18,6 +18,7 @@ import type {
   Message,
   ObservabilityAIAssistantScreenContextRequest,
   InstructionOrPlainText,
+  AdHocInstruction,
 } from '../../common/types';
 import type { ObservabilityAIAssistantRouteHandlerResources } from '../routes/types';
 import { ChatFunctionClient } from './chat_function_client';
@@ -76,6 +77,8 @@ export type RegisterInstructionCallback = ({
 
 export type RegisterInstruction = (...instruction: InstructionOrCallback[]) => void;
 
+export type RegisterAdHocInstruction = (...instruction: AdHocInstruction[]) => void;
+
 export type RegisterFunction = <
   TParameters extends CompatibleJSONSchema = any,
   TResponse extends FunctionResponse = any,

diff --git a/...on/observability_ai_assistant/server/service/util/get_system_message_from_instructions.ts b/...on/observability_ai_assistant/server/service/util/get_system_message_from_instructions.ts
@@ -45,7 +45,7 @@ export function getSystemMessageFromInstructions({
 
   const adHocInstructionsWithId = adHocInstructions.map((adHocInstruction) => ({
     ...adHocInstruction,
-    doc_id: adHocInstruction.doc_id ?? v4(),
+    doc_id: adHocInstruction?.doc_id ?? v4(),
   }));
 
   // split ad hoc instructions into user instructions and application instructions

diff --git a/...gins/observability_solution/observability_ai_assistant_app/server/rule_connector/index.ts b/...gins/observability_solution/observability_ai_assistant_app/server/rule_connector/index.ts
@@ -230,7 +230,7 @@ If available, include the link of the conversation at the end of your answer.`
               availableFunctionNames: functionClient.getFunctions().map((fn) => fn.definition.name),
               applicationInstructions: functionClient.getInstructions(),
               userInstructions: [],
-              adHocInstructions: [],
+              adHocInstructions: functionClient.getAdhocInstructions(),
             }),
           },
         },