[8.x] [Obs AI Assistant] Add API test for get_alerts_dataset_info t…

…ool (#212858) (#213202)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[Obs AI Assistant] Add API test for `get_alerts_dataset_info` tool
(#212858)](#212858)

<!--- Backport version: 9.6.6 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sorenlouv/backport)

<!--BACKPORT [{"author":{"name":"Søren
Louv-Jansen","email":"soren.louv@elastic.co"},"sourceCommit":{"committedDate":"2025-03-05T08:09:22Z","message":"[Obs
AI Assistant] Add API test for `get_alerts_dataset_info` tool
(#212858)\n\nFollow-up to:
https://github.com/elastic/kibana/pull/212077\n\nThis PR includes an API
test that covers `get_alerts_dataset_info` and\nwould have caught the
bug fixed in\nhttps://github.com//pull/212077.\n\nIt also
contains the following bug fixes:\n\n- Fix system message in
`select_relevant_fields`\n- Change prompt in `select_relevant_fields` so
that the LLM consistently\nuses the right format when
responding.","sha":"0fb83efd82ae3ebd8a9fe27813e436b80cd240d3","branchLabelMapping":{"^v9.1.0$":"main","^v8.19.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","backport:prev-minor","backport:prev-major","Team:Obs
AI Assistant","ci:project-deploy-observability","v9.1.0"],"title":"[Obs
AI Assistant] Add API test for `get_alerts_dataset_info`
tool","number":212858,"url":"https://github.com/elastic/kibana/pull/212858","mergeCommit":{"message":"[Obs
AI Assistant] Add API test for `get_alerts_dataset_info` tool
(#212858)\n\nFollow-up to:
https://github.com/elastic/kibana/pull/212077\n\nThis PR includes an API
test that covers `get_alerts_dataset_info` and\nwould have caught the
bug fixed in\nhttps://github.com//pull/212077.\n\nIt also
contains the following bug fixes:\n\n- Fix system message in
`select_relevant_fields`\n- Change prompt in `select_relevant_fields` so
that the LLM consistently\nuses the right format when
responding.","sha":"0fb83efd82ae3ebd8a9fe27813e436b80cd240d3"}},"sourceBranch":"main","suggestedTargetBranches":[],"targetPullRequestStates":[{"branch":"main","label":"v9.1.0","branchLabelMappingKey":"^v9.1.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/212858","number":212858,"mergeCommit":{"message":"[Obs
AI Assistant] Add API test for `get_alerts_dataset_info` tool
(#212858)\n\nFollow-up to:
https://github.com/elastic/kibana/pull/212077\n\nThis PR includes an API
test that covers `get_alerts_dataset_info` and\nwould have caught the
bug fixed in\nhttps://github.com//pull/212077.\n\nIt also
contains the following bug fixes:\n\n- Fix system message in
`select_relevant_fields`\n- Change prompt in `select_relevant_fields` so
that the LLM consistently\nuses the right format when
responding.","sha":"0fb83efd82ae3ebd8a9fe27813e436b80cd240d3"}}]}]
BACKPORT-->

Co-authored-by: Søren Louv-Jansen <soren.louv@elastic.co>

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/get_dataset_info/get_relevant_field_names.ts

@@ -13,6 +13,12 @@ import { MessageRole, ShortIdTable, type Message } from '../../../common';
 import { concatenateChatCompletionChunks } from '../../../common/utils/concatenate_chat_completion_chunks';
 import { FunctionCallChatFunction } from '../../service/types';
 
+const SELECT_RELEVANT_FIELDS_NAME = 'select_relevant_fields';
+export const GET_RELEVANT_FIELD_NAMES_SYSTEM_MESSAGE = `You are a helpful assistant for Elastic Observability. 
+Your task is to determine which fields are relevant to the conversation by selecting only the field IDs from the provided list. 
+The list in the user message consists of JSON objects that map a human-readable "field" name to its unique "id". 
+You must not output any field names — only the corresponding "id" values. Ensure that your output follows the exact JSON format specified.`;
+
 export async function getRelevantFieldNames({
   index,
   start,
@@ -100,19 +106,15 @@ export async function getRelevantFieldNames({
         await chat('get_relevant_dataset_names', {
           signal,
           stream: true,
-          systemMessage: `You are a helpful assistant for Elastic Observability.
-            Your task is to create a list of field names that are relevant
-            to the conversation, using ONLY the list of fields and
-            types provided in the last user message. DO NOT UNDER ANY
-            CIRCUMSTANCES include fields not mentioned in this list.`,
+          systemMessage: GET_RELEVANT_FIELD_NAMES_SYSTEM_MESSAGE,
           messages: [
             // remove the last function request
             ...messages.slice(0, -1),
             {
               '@timestamp': new Date().toISOString(),
               message: {
                 role: MessageRole.User,
-                content: `This is the list:
+                content: `Below is a list of fields. Each entry is a JSON object that contains a "field" (the field name) and an "id" (the unique identifier). Use only the "id" values from this list when selecting relevant fields:
 
             ${fieldsInChunk
               .map((field) => JSON.stringify({ field, id: shortIdTable.take(field) }))
@@ -122,8 +124,12 @@ export async function getRelevantFieldNames({
           ],
           functions: [
             {
-              name: 'select_relevant_fields',
-              description: 'The IDs of the fields you consider relevant to the conversation',
+              name: SELECT_RELEVANT_FIELDS_NAME,
+              description: `Return only the field IDs (from the provided list) that you consider relevant to the conversation. Do not use any of the field names. Your response must be in the exact JSON format:
+              {
+                "fieldIds": ["id1", "id2", "id3"]
+              }
+              Only include IDs from the list provided in the user message.`,
               parameters: {
                 type: 'object',
                 properties: {
@@ -138,7 +144,7 @@ export async function getRelevantFieldNames({
               } as const,
             },
           ],
-          functionCall: 'select_relevant_fields',
+          functionCall: SELECT_RELEVANT_FIELDS_NAME,
         })
       ).pipe(concatenateChatCompletionChunks());
 

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/index.ts

@@ -39,52 +39,46 @@ export const registerFunctions: RegistrationCallback = async ({
   };
 
   const isServerless = !!resources.plugins.serverless;
-  if (scopes.includes('observability')) {
-    functions.registerInstruction(`You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities.
 
-    It's very important to not assume what the user is meaning. Ask them for clarification if needed.
-
-    If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
-
-    In KQL ("kqlFilter")) escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
-    /\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
-
-    You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
-
-    Note that ES|QL (the Elasticsearch Query Language which is a new piped language) is the preferred query language.
-
-    If you want to call a function or tool, only call it a single time per message. Wait until the function has been executed and its results
-    returned to you, before executing the same tool or another tool again if needed.
-
-    DO NOT UNDER ANY CIRCUMSTANCES USE ES|QL syntax (\`service.name == "foo"\`) with "kqlFilter" (\`service.name:"foo"\`).
-
-    The user is able to change the language which they want you to reply in on the settings page of the AI Assistant for Observability and Search, which can be found in the ${
-      isServerless ? `Project settings.` : `Stack Management app under the option AI Assistants`
-    }.
-    If the user asks how to change the language, reply in the same language the user asked in.`);
-  }
-
-  if (scopes.length === 0 || (scopes.length === 1 && scopes[0] === 'all')) {
-    functions.registerInstruction(
-      `You are a helpful assistant for Elasticsearch. Your goal is to help Elasticsearch users accomplish tasks using Kibana and Elasticsearch. You can help them construct queries, index data, search data, use Elasticsearch APIs, generate sample data, visualise and analyze data.
-
-  It's very important to not assume what the user means. Ask them for clarification if needed.
-
-  If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
-
-  In KQL ("kqlFilter")) escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
-  /\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
-
-  You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
-
-  If you want to call a function or tool, only call it a single time per message. Wait until the function has been executed and its results
-  returned to you, before executing the same tool or another tool again if needed.
-
-  The user is able to change the language which they want you to reply in on the settings page of the AI Assistant for Observability and Search, which can be found in the ${
-    isServerless ? `Project settings.` : `Stack Management app under the option AI Assistants`
-  }.
-  If the user asks how to change the language, reply in the same language the user asked in.`
-    );
+  const isObservabilityDeployment = scopes.includes('observability');
+  const isGenericDeployment = scopes.length === 0 || (scopes.length === 1 && scopes[0] === 'all');
+
+  if (isObservabilityDeployment || isGenericDeployment) {
+    functions.registerInstruction(`
+${
+  isObservabilityDeployment
+    ? `You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities.`
+    : `You are a helpful assistant for Elasticsearch. Your goal is to help Elasticsearch users accomplish tasks using Kibana and Elasticsearch. You can help them construct queries, index data, search data, use Elasticsearch APIs, generate sample data, visualise and analyze data.`
+}
+      It's very important to not assume what the user means. Ask them for clarification if needed.
+  
+      If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
+  
+      In KQL ("kqlFilter")) escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
+      /\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
+  
+      You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
+  
+      ${
+        isObservabilityDeployment
+          ? 'Note that ES|QL (the Elasticsearch Query Language which is a new piped language) is the preferred query language.'
+          : ''
+      }
+  
+      If you want to call a function or tool, only call it a single time per message. Wait until the function has been executed and its results
+      returned to you, before executing the same tool or another tool again if needed.
+  
+
+      ${
+        isObservabilityDeployment
+          ? 'DO NOT UNDER ANY CIRCUMSTANCES USE ES|QL syntax (`service.name == "foo"`) with "kqlFilter" (`service.name:"foo"`).'
+          : ''
+      }
+  
+      The user is able to change the language which they want you to reply in on the settings page of the AI Assistant for Observability and Search, which can be found in the ${
+        isServerless ? `Project settings.` : `Stack Management app under the option AI Assistants`
+      }.
+      If the user asks how to change the language, reply in the same language the user asked in.`);
   }
 
   const { ready: isKnowledgeBaseReady } = await client.getKnowledgeBaseStatus();

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

@@ -274,8 +274,8 @@ export class ObservabilityAIAssistantClient {
                 chat: (name, chatParams) => {
                   // inject a chat function with predefined parameters
                   return this.chat(name, {
-                    ...chatParams,
                     systemMessage,
+                    ...chatParams,
                     signal,
                     simulateFunctionCalling,
                     connectorId,

x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/alerts.ts

@@ -112,9 +112,10 @@ export function registerAlertsFunction({
           signal,
           chat: (
             operationName,
-            { messages: nextMessages, functionCall, functions: nextFunctions }
+            { messages: nextMessages, functionCall, functions: nextFunctions, systemMessage }
           ) => {
             return chat(operationName, {
+              systemMessage,
               messages: nextMessages,
               functionCall,
               functions: nextFunctions,

x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/index.ts

@@ -18,7 +18,6 @@ import { createFunctionResponseMessage } from '@kbn/observability-ai-assistant-p
 import { convertMessagesForInference } from '@kbn/observability-ai-assistant-plugin/common/convert_messages_for_inference';
 import { map } from 'rxjs';
 import { v4 } from 'uuid';
-import { RegisterInstructionCallback } from '@kbn/observability-ai-assistant-plugin/server/service/types';
 import type { FunctionRegistrationParameters } from '..';
 import { runAndValidateEsqlQuery } from './validate_esql_query';
 
@@ -30,9 +29,12 @@ export function registerQueryFunction({
   resources,
   pluginsStart,
 }: FunctionRegistrationParameters) {
-  const instruction: RegisterInstructionCallback = ({ availableFunctionNames }) =>
-    availableFunctionNames.includes(QUERY_FUNCTION_NAME)
-      ? `You MUST use the "${QUERY_FUNCTION_NAME}" function when the user wants to:
+  functions.registerInstruction(({ availableFunctionNames }) => {
+    if (!availableFunctionNames.includes(QUERY_FUNCTION_NAME)) {
+      return;
+    }
+
+    return `You MUST use the "${QUERY_FUNCTION_NAME}" function when the user wants to:
   - visualize data
   - run any arbitrary query
   - breakdown or filter ES|QL queries that are displayed on the current page
@@ -48,9 +50,8 @@ export function registerQueryFunction({
   even if it has been called before.
 
   When the "visualize_query" function has been called, a visualization has been displayed to the user. DO NOT UNDER ANY CIRCUMSTANCES follow up a "visualize_query" function call with your own visualization attempt.
-  If the "${EXECUTE_QUERY_NAME}" function has been called, summarize these results for the user. The user does not see a visualization in this case.`
-      : undefined;
-  functions.registerInstruction(instruction);
+  If the "${EXECUTE_QUERY_NAME}" function has been called, summarize these results for the user. The user does not see a visualization in this case.`;
+  });
 
   functions.registerFunction(
     {

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/chat/chat.spec.ts

@@ -85,7 +85,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
           },
         },
       });
-      await proxy.waitForAllInterceptorsSettled();
+      await proxy.waitForAllInterceptorsToHaveBeenCalled();
       expect(status).to.be(200);
     });
 
@@ -104,7 +104,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
           },
         },
       });
-      await proxy.waitForAllInterceptorsSettled();
+      await proxy.waitForAllInterceptorsToHaveBeenCalled();
       const simulator = await simulatorPromise;
       const requestData = simulator.requestBody; // This is the request sent to the LLM
       expect(requestData.messages[0].content).to.eql(SYSTEM_MESSAGE);

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/complete.spec.ts

@@ -76,7 +76,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
           scopes: ['all'],
         });
 
-      await proxy.waitForAllInterceptorsSettled();
+      await proxy.waitForAllInterceptorsToHaveBeenCalled();
 
       return String(response.body)
         .split('\n')
@@ -133,7 +133,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
         await new Promise<void>((resolve) => passThrough.on('end', () => resolve()));
 
-        await proxy.waitForAllInterceptorsSettled();
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
 
         parsedEvents = decodeEvents(receivedChunks.join(''));
       });
@@ -243,7 +243,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
             },
           },
         });
-        await proxy.waitForAllInterceptorsSettled();
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
         const simulator = await simulatorPromise;
         const requestData = simulator.requestBody;
         expect(requestData.messages[0].role).to.eql('system');
@@ -420,7 +420,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
         expect(createResponse.status).to.be(200);
 
-        await proxy.waitForAllInterceptorsSettled();
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
 
         conversationCreatedEvent = getConversationCreatedEvent(createResponse.body);
 
@@ -463,7 +463,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
         expect(updatedResponse.status).to.be(200);
 
-        await proxy.waitForAllInterceptorsSettled();
+        await proxy.waitForAllInterceptorsToHaveBeenCalled();
       });
 
       after(async () => {

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/alerts.spec.ts

@@ -46,7 +46,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         },
       });
 
-      await proxy.waitForAllInterceptorsSettled();
+      await proxy.waitForAllInterceptorsToHaveBeenCalled();
 
       alertsEvents = getMessageAddedEvents(alertsResponseBody);
     });

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/elasticsearch.spec.ts

@@ -65,7 +65,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         },
       });
 
-      await proxy.waitForAllInterceptorsSettled();
+      await proxy.waitForAllInterceptorsToHaveBeenCalled();
 
       events = getMessageAddedEvents(responseBody);
     });