Repository: elastic/kibana
Pull request #213517: [Obs AI Assistant] Add API test for `execute_query` tool
Author: login: sorenlouv
Labels: name: release_note:skip
description: Skip the PR/issue when compiling release notes
name: v9.0.0
name: Team:Obs AI Assistant
description: Observability AI Assistant
name: ci:project-deploy-observability
description: Create an Observability project
name: backport:version
description: Backport to applied version labels
name: v9.1.0
name: v8.19.0
Created: 2025-03-07T08:42:45Z | Merged: 2025-03-11T21:45:28Z

Related: https://github.com/elastic/kibana/issues/180787

- Adds test for `execute_query` function
- Add `drop_null_columns` to esql query in order to avoid column explosion 

Depends on: https://github.com/elastic/kibana/pull/213231

<!--ONMERGE {"backportTargets":["8.x","9.0"]} ONMERGE-->


Diff:
diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/index.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/index.ts
index 3e740fbac28f6..3853c2d2b9657 100644
--- a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/index.ts
+++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/index.ts
@@ -28,6 +28,7 @@ export function registerQueryFunction({
   functions,
   resources,
   pluginsStart,
+  signal,
 }: FunctionRegistrationParameters) {
   functions.registerInstruction(({ availableFunctionNames }) => {
     if (!availableFunctionNames.includes(QUERY_FUNCTION_NAME)) {
@@ -83,6 +84,7 @@ export function registerQueryFunction({
       const { error, errorMessages, rows, columns } = await runAndValidateEsqlQuery({
         query: correctedQuery,
         client,
+        signal,
       });
 
       if (!!error) {
@@ -114,7 +116,7 @@ export function registerQueryFunction({
       function takes no input.`,
       visibility: FunctionVisibility.AssistantOnly,
     },
-    async ({ messages, connectorId, simulateFunctionCalling }, signal) => {
+    async ({ messages, connectorId, simulateFunctionCalling }) => {
       const esqlFunctions = functions
         .getFunctions()
         .filter(
diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts
index 1c36d08594521..b3e3b1c25ad02 100644
--- a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts
+++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts
@@ -16,9 +16,11 @@ import { splitIntoCommands } from '@kbn/inference-plugin/common';
 export async function runAndValidateEsqlQuery({
   query,
   client,
+  signal,
 }: {
   query: string;
   client: ElasticsearchClient;
+  signal: AbortSignal;
 }): Promise<{
   columns?: DatatableColumn[];
   rows?: ESQLRow[];
@@ -47,30 +49,22 @@ export async function runAndValidateEsqlQuery({
     return 'text' in error ? error.text : error.message;
   });
 
-  return client.transport
-    .request({
-      method: 'POST',
-      path: '_query',
-      body: {
-        query,
-      },
-    })
-    .then((res) => {
-      const esqlResponse = res as ESQLSearchResponse;
+  try {
+    const res = await client.esql.query({ query, drop_null_columns: true }, { signal });
+    const esqlResponse = res as unknown as ESQLSearchResponse;
 
-      const columns =
-        esqlResponse.columns?.map(({ name, type }) => ({
-          id: name,
-          name,
-          meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType },
-        })) ?? [];
+    const columns =
+      esqlResponse.columns?.map(({ name, type }) => ({
+        id: name,
+        name,
+        meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType },
+      })) ?? [];
 
-      return { columns, rows: esqlResponse.values };
-    })
-    .catch((error) => {
-      return {
-        error,
-        ...(errorMessages.length ? { errorMessages } : {}),
-      };
-    });
+    return { columns, rows: esqlResponse.values };
+  } catch (error) {
+    return {
+      error,
+      ...(errorMessages.length ? { errorMessages } : {}),
+    };
+  }
 }
diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts
index 4eeba0450e6e4..e7a4e8c0dc029 100644
--- a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts
+++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts
@@ -30,6 +30,7 @@ const getMessageForLLM = (
 export function registerVisualizeESQLFunction({
   functions,
   resources,
+  signal,
 }: FunctionRegistrationParameters) {
   functions.registerFunction(
     visualizeESQLFunction,
@@ -43,6 +44,7 @@ export function registerVisualizeESQLFunction({
       const { columns, errorMessages, rows, error } = await runAndValidateEsqlQuery({
         query: correctedQuery,
         client: (await resources.context.core).elasticsearch.client.asCurrentUser,
+        signal,
       });
 
       const message = getMessageForLLM(intention, query, Boolean(errorMessages?.length));
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
new file mode 100644
index 0000000000000..a389bc83a711e
--- /dev/null
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
@@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('First LLM request - Initial tool selection', () => {
+        it('exposes the right tools', () => {
+          expect(firstRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'query',
+            'get_alerts_dataset_info',
+            'alerts',
+            'changes',
+            'elasticsearch',
+            'kibana',
+            'get_dataset_info',
+            'execute_connector',
+          ]);
+        });
+      });
+
+      describe('The second request - Structured output validation', () => {
+        it('contains the correct number of messages', () => {
+          expect(secondRequestBody.messages.length).to.be(6);
+        });
+
+        it('contains the `structuredOutput` tool choice', () => {
+          // @ts-expect-error
+          const hasToolChoice = secondRequestBody.tool_choice.function?.name === 'structuredOutput';
+          expect(hasToolChoice).to.be(true);
+        });
+
+        it('contains user message with information about how to request ESQL documentation', () => {
+          expect(last(secondRequestBody.messages)?.content).to.contain(
+            'Based on the previous conversation, request documentation'
+          );
+        });
+      });
+
+      describe('The third request - Requesting ESQL documentation', () => {
+        it('contains the `request_documentation` tool call request', () => {
+          const hasToolCall = thirdRequestBody.messages.some(
+            // @ts-expect-error
+            (message) => message.tool_calls?.[0]?.function?.name === 'request_documentation'
+          );
+          expect(hasToolCall).to.be(true);
+        });
+
+        it('contains ESQL documentation', () => {
+          const parsed = JSON.parse(last(thirdRequestBody.messages)?.content as string);
+          expect(parsed.documentation.OPERATORS).to.contain('Binary Operators');
+        });
+
+        it('allows the LLM to call the tools execute_query, visualize_query and request_documentation', () => {
+          expect(thirdRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'execute_query',
+            'visualize_query',
+            'request_documentation',
+          ]);
+        });
+      });
+
+      describe('The fourth request - Executing the ES|QL query', () => {
+        it('contains the `execute_query` tool call request', () => {
+          const hasToolCall = fourthRequestBody.messages.some(
+            // @ts-expect-error
+            (message) => message.tool_calls?.[0]?.function?.name === 'execute_query'
+          );
+          expect(hasToolCall).to.be(true);
+        });
+
+        it('emits a messageAdded event with the `execute_query` tool response', () => {
+          const event = messageAddedEvents.find(
+            ({ message }) => message.message.name === 'execute_query'
+          );
+          expect(event?.message.message.content).to.contain('simple log message');
+        });
+
+        describe('the `execute_query` tool call response', () => {
+          let toolCallResponse: { columns: EsqlResponse['columns']; rows: EsqlResponse['values'] };
+          before(async () => {
+            toolCallResponse = JSON.parse(last(fourthRequestBody.messages)?.content as string);
+          });
+
+          it('has the correct columns', () => {
+            expect(toolCallResponse.columns.map(({ name }) => name)).to.eql([
+              'message',
+              '@timestamp',
+            ]);
+          });
+
+          it('has the correct number of rows', () => {
+            expect(toolCallResponse.rows.length).to.be(10);
+          });
+
+          it('has the right log message', () => {
+            expect(toolCallResponse.rows[0][0]).to.be('simple log message');
+          });
+        });
+      });
+    });
+  });
+}
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts
index c7db9d728337c..5e1efaa295310 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/index.ts
@@ -18,6 +18,7 @@ export default function aiAssistantApiIntegrationTests({
     loadTestFile(require.resolve('./complete/functions/alerts.spec.ts'));
     loadTestFile(require.resolve('./complete/functions/get_alerts_dataset_info.spec.ts'));
     loadTestFile(require.resolve('./complete/functions/get_dataset_info.spec.ts'));
+    loadTestFile(require.resolve('./complete/functions/execute_query.spec.ts'));
     loadTestFile(require.resolve('./complete/functions/elasticsearch.spec.ts'));
     loadTestFile(require.resolve('./complete/functions/summarize.spec.ts'));
     loadTestFile(require.resolve('./public_complete/public_complete.spec.ts'));
diff --git a/x-pack/test/observability_ai_assistant_functional/tests/conversations/index.spec.ts b/x-pack/test/observability_ai_assistant_functional/tests/conversations/index.spec.ts
index 49eb42d06d8ad..3f1d0c89b2282 100644
--- a/x-pack/test/observability_ai_assistant_functional/tests/conversations/index.spec.ts
+++ b/x-pack/test/observability_ai_assistant_functional/tests/conversations/index.spec.ts
@@ -11,6 +11,7 @@ import { ChatFeedback } from '@kbn/observability-ai-assistant-plugin/public/anal
 import { pick } from 'lodash';
 import { parse as parseCookie } from 'tough-cookie';
 import { kbnTestConfig } from '@kbn/test';
+import { systemMessageSorted } from '../../../api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/helpers';
 import {
   createLlmProxy,
   LlmProxy,
@@ -282,8 +283,8 @@ export default function ApiTest({ getService, getPageObjects }: FtrProviderConte
                   'You are a helpful assistant for Elastic Observability. Your goal is '
                 );
 
-                expect(sortSystemMessage(systemMessage!)).to.eql(
-                  sortSystemMessage(primarySystemMessage)
+                expect(systemMessageSorted(systemMessage!)).to.eql(
+                  systemMessageSorted(primarySystemMessage)
                 );
 
                 expect(firstUserMessage.content).to.eql('hello');
@@ -465,11 +466,3 @@ export default function ApiTest({ getService, getPageObjects }: FtrProviderConte
     });
   });
 }
-
-// order of instructions can vary, so we sort to compare them
-function sortSystemMessage(message: string) {
-  return message
-    .split('\n\n')
-    .map((line) => line.trim())
-    .sort();
-}


Reviews:
user: login: sorenlouv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T22:53:31Z
user: login: sorenlouv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T08:04:18Z
user: login: sorenlouv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T08:05:59Z
user: login: viduni94
body: LGTM, just a few nits
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T12:41:53Z
user: login: viduni94
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T14:23:41Z

Review comments:
diff_hunk: @@ -47,30 +49,22 @@ export async function runAndValidateEsqlQuery({
     return 'text' in error ? error.text : error.message;
   });
 
-  return client.transport
-    .request({
-      method: 'POST',
-      path: '_query',
-      body: {
-        query,
-      },
-    })
-    .then((res) => {
-      const esqlResponse = res as ESQLSearchResponse;
+  try {
+    const res = await client.esql.query({ query, drop_null_columns: true }, { signal });
+    const esqlResponse = res as unknown as ESQLSearchResponse;
 
-      const columns =
-        esqlResponse.columns?.map(({ name, type }) => ({
-          id: name,
-          name,
-          meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType },
-        })) ?? [];
+    const columns =
+      esqlResponse.columns?.map(({ name, type }) => ({
+        id: name,
+        name,
+        meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType },
+      })) ?? [];
 
-      return { columns, rows: esqlResponse.values };
-    })
-    .catch((error) => {
-      return {
-        error,
-        ...(errorMessages.length ? { errorMessages } : {}),
-      };
-    });
+    return { columns, rows: esqlResponse.values };
path: x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts
user: login: sorenlouv
body: @dgieselaar nit Do you recall why we rename `values` to `rows`? (asking because we cannot use the built-in types and it seems like and odd change to make).
created_at: 2025-03-07T22:53:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -47,30 +49,22 @@ export async function runAndValidateEsqlQuery({
     return 'text' in error ? error.text : error.message;
   });
 
-  return client.transport
-    .request({
-      method: 'POST',
-      path: '_query',
-      body: {
-        query,
-      },
-    })
-    .then((res) => {
-      const esqlResponse = res as ESQLSearchResponse;
+  try {
+    const res = await client.esql.query({ query, drop_null_columns: true }, { signal });
path: x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts
user: login: sorenlouv
body: @dgieselaar I assume the ES client didn't have the `esql.query` method previously which is why we used the `transport` method. LMK if there's any reason we should keep doing this
created_at: 2025-03-10T08:04:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -16,9 +16,11 @@ import { splitIntoCommands } from '@kbn/inference-plugin/common';
 export async function runAndValidateEsqlQuery({
   query,
   client,
+  signal,
 }: {
   query: string;
   client: ElasticsearchClient;
+  signal: AbortSignal;
path: x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/validate_esql_query.ts
user: login: sorenlouv
body: Added missing `signal`. This should mean that we cancel the ES request when appropriate.
created_at: 2025-03-10T08:05:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('The first request', () => {
path: x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
user: login: viduni94
body: Maybe we could add a description about what we expect in the first request, second request, etc.,  to the block (just a suggestion) :

```suggestion
      describe('First LLM request - Initial tool selection', () => {
```
created_at: 2025-03-11T12:34:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('The first request', () => {
+        it('exposes the right tools', () => {
+          expect(firstRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'query',
+            'get_alerts_dataset_info',
+            'alerts',
+            'changes',
+            'elasticsearch',
+            'kibana',
+            'get_dataset_info',
+            'execute_connector',
+          ]);
+        });
+      });
+
+      describe('The second request', () => {
+        it('contains the correct number of messages', () => {
+          expect(secondRequestBody.messages.length).to.be(6);
+        });
+
+        it('contains the `structuredOutput` tool choice', () => {
+          // @ts-expect-error
+          const hasToolChoice = secondRequestBody.tool_choice.function?.name === 'structuredOutput';
+          expect(hasToolChoice).to.be(true);
+        });
+
+        it('contains user message with information about how to request ESQL documentation', () => {
+          expect(last(secondRequestBody.messages)?.content).to.contain(
+            'Based on the previous conversation, request documentation'
+          );
+        });
+      });
+
+      describe('The third request', () => {
+        it('contains the `request_documentation` tool call request', () => {
path: x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
user: login: viduni94
body: There is a scenario where this tool call won't be requested if the user has not installed the documentation. Do we want to validate that flow as well?
created_at: 2025-03-11T12:37:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('The first request', () => {
+        it('exposes the right tools', () => {
+          expect(firstRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'query',
+            'get_alerts_dataset_info',
+            'alerts',
+            'changes',
+            'elasticsearch',
+            'kibana',
+            'get_dataset_info',
+            'execute_connector',
+          ]);
+        });
+      });
+
+      describe('The second request', () => {
+        it('contains the correct number of messages', () => {
+          expect(secondRequestBody.messages.length).to.be(6);
+        });
+
+        it('contains the `structuredOutput` tool choice', () => {
+          // @ts-expect-error
+          const hasToolChoice = secondRequestBody.tool_choice.function?.name === 'structuredOutput';
+          expect(hasToolChoice).to.be(true);
+        });
+
+        it('contains user message with information about how to request ESQL documentation', () => {
+          expect(last(secondRequestBody.messages)?.content).to.contain(
+            'Based on the previous conversation, request documentation'
+          );
+        });
+      });
+
+      describe('The third request', () => {
+        it('contains the `request_documentation` tool call request', () => {
+          const hasToolCall = thirdRequestBody.messages.some(
+            // @ts-expect-error
+            (message) => message.tool_calls?.[0]?.function?.name === 'request_documentation'
+          );
+          expect(hasToolCall).to.be(true);
+        });
+
+        it('contains ESQL documentation', () => {
+          const parsed = JSON.parse(last(thirdRequestBody.messages)?.content as string);
+          expect(parsed.documentation.OPERATORS).to.contain('Binary Operators');
+        });
+
+        it('allows the LLM to call the tools execute_query, visualize_query and request_documentation', () => {
+          expect(thirdRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'execute_query',
+            'visualize_query',
+            'request_documentation',
+          ]);
+        });
+      });
+
+      describe('The fourth request', () => {
path: x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
user: login: viduni94
body: (just a suggestion, can be ignored if it doesn't add much value)
```suggestion
      describe('The fourth request - Executing the ES|QL query', () => {
```
created_at: 2025-03-11T12:38:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('The first request', () => {
+        it('exposes the right tools', () => {
+          expect(firstRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'query',
+            'get_alerts_dataset_info',
+            'alerts',
+            'changes',
+            'elasticsearch',
+            'kibana',
+            'get_dataset_info',
+            'execute_connector',
+          ]);
+        });
+      });
+
+      describe('The second request', () => {
+        it('contains the correct number of messages', () => {
+          expect(secondRequestBody.messages.length).to.be(6);
+        });
+
+        it('contains the `structuredOutput` tool choice', () => {
+          // @ts-expect-error
+          const hasToolChoice = secondRequestBody.tool_choice.function?.name === 'structuredOutput';
+          expect(hasToolChoice).to.be(true);
+        });
+
+        it('contains user message with information about how to request ESQL documentation', () => {
+          expect(last(secondRequestBody.messages)?.content).to.contain(
+            'Based on the previous conversation, request documentation'
+          );
+        });
+      });
+
+      describe('The third request', () => {
path: x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
user: login: viduni94
body: (just a suggestion, can be ignored if it doesn't add much value)
```suggestion
      describe('The third request - Requesting ESQL documentation', () => {
```
created_at: 2025-03-11T12:38:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('The first request', () => {
+        it('exposes the right tools', () => {
+          expect(firstRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'query',
+            'get_alerts_dataset_info',
+            'alerts',
+            'changes',
+            'elasticsearch',
+            'kibana',
+            'get_dataset_info',
+            'execute_connector',
+          ]);
+        });
+      });
+
+      describe('The second request', () => {
path: x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
user: login: viduni94
body: (just a suggestion, can be ignored if it doesn't add much value)
```suggestion
      describe('The second request - Structured output validation', () => {
```
created_at: 2025-03-11T12:39:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,210 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { MessageAddEvent } from '@kbn/observability-ai-assistant-plugin/common';
+import expect from '@kbn/expect';
+import { LogsSynthtraceEsClient } from '@kbn/apm-synthtrace';
+import { last } from 'lodash';
+import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+import { EsqlResponse } from '@elastic/elasticsearch/lib/helpers';
+import {
+  LlmProxy,
+  createLlmProxy,
+} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy';
+import { chatComplete } from './helpers';
+import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context';
+import { createSimpleSyntheticLogs } from '../../synthtrace_scenarios/simple_logs';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+  const synthtrace = getService('synthtrace');
+
+  describe('execute_query', function () {
+    this.tags(['failsOnMKI']);
+    let llmProxy: LlmProxy;
+    let connectorId: string;
+
+    before(async () => {
+      llmProxy = await createLlmProxy(log);
+      connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+        port: llmProxy.getPort(),
+      });
+    });
+
+    after(async () => {
+      llmProxy.close();
+      await observabilityAIAssistantAPIClient.deleteActionConnector({
+        actionId: connectorId,
+      });
+    });
+
+    // Calling `execute_query` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let messageAddedEvents: MessageAddEvent[];
+      let logsSynthtraceEsClient: LogsSynthtraceEsClient;
+      let firstRequestBody: ChatCompletionStreamParams;
+      let secondRequestBody: ChatCompletionStreamParams;
+      let thirdRequestBody: ChatCompletionStreamParams;
+      let fourthRequestBody: ChatCompletionStreamParams;
+
+      before(async () => {
+        logsSynthtraceEsClient = synthtrace.createLogsSynthtraceEsClient();
+        await createSimpleSyntheticLogs({
+          logsSynthtraceEsClient,
+          dataset: 'apache.access',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'query',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'structuredOutput',
+          arguments: () => JSON.stringify({}),
+          // @ts-expect-error
+          when: (requestBody) => requestBody.tool_choice?.function?.name === 'structuredOutput',
+        });
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'execute_query',
+          arguments: () =>
+            JSON.stringify({
+              query: `FROM logs-apache.access-default
+            | KEEP message
+            | SORT @timestamp DESC
+            | LIMIT 10`,
+            }),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation({ content: 'Hello from user' });
+
+        ({ messageAddedEvents } = await chatComplete({
+          userPrompt: 'Please retrieve the most recent Apache log messages',
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        }));
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+        firstRequestBody = llmProxy.interceptedRequests[0].requestBody;
+        secondRequestBody = llmProxy.interceptedRequests[1].requestBody;
+        thirdRequestBody = llmProxy.interceptedRequests[2].requestBody;
+        fourthRequestBody = llmProxy.interceptedRequests[3].requestBody;
+      });
+
+      after(async () => {
+        await logsSynthtraceEsClient.clean();
+      });
+
+      it('makes 4 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(4);
+      });
+
+      it('emits 7 messageAdded events', () => {
+        expect(messageAddedEvents.length).to.be(7);
+      });
+
+      describe('The first request', () => {
+        it('exposes the right tools', () => {
+          expect(firstRequestBody.tools?.map((t) => t.function.name)).to.eql([
+            'query',
+            'get_alerts_dataset_info',
+            'alerts',
+            'changes',
+            'elasticsearch',
+            'kibana',
+            'get_dataset_info',
+            'execute_connector',
+          ]);
+        });
+      });
+
+      describe('The second request', () => {
+        it('contains the correct number of messages', () => {
+          expect(secondRequestBody.messages.length).to.be(6);
+        });
+
+        it('contains the `structuredOutput` tool choice', () => {
+          // @ts-expect-error
+          const hasToolChoice = secondRequestBody.tool_choice.function?.name === 'structuredOutput';
+          expect(hasToolChoice).to.be(true);
+        });
+
+        it('contains user message with information about how to request ESQL documentation', () => {
+          expect(last(secondRequestBody.messages)?.content).to.contain(
+            'Based on the previous conversation, request documentation'
+          );
+        });
+      });
+
+      describe('The third request', () => {
+        it('contains the `request_documentation` tool call request', () => {
path: x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/execute_query.spec.ts
user: login: viduni94
body: Ignore this comment. I got `request_documentation` mixed up with `retrieve_elastic_doc`
created_at: 2025-03-11T14:23:41Z
author_association: CONTRIBUTOR

PR comments:
user: login: elasticmachine
created_at: 2025-03-07T22:53:36Z
body: Pinging @elastic/obs-ai-assistant (Team:Obs AI Assistant)
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-03-07T22:53:43Z
body: 
### :robot: GitHub comments

<details><summary>Expand to view the GitHub comments</summary>
<p>

Just comment with:
- `/oblt-deploy` : Deploy a Kibana instance using the Observability test environments.
- `run` `docs-build` : Re-trigger the docs validation. (use unformatted text in the comment!)
</p>
</details>

author_association: CONTRIBUTOR
user: login: elasticmachine
created_at: 2025-03-11T15:19:24Z
body: ## :green_heart: Build Succeeded
* [Buildkite Build](https://buildkite.com/elastic/kibana-pull-request/builds/282760)
* Commit: 4f9c179dcdede098b863e7fbc81f2e2fc1b82a27
* Kibana Serverless Image: `docker.elastic.co/kibana-ci/kibana-serverless:pr-213517-4f9c179dcded`

### Metrics [[docs](https://www.elastic.co/guide/en/kibana/master/ci-metrics.html)]
> ✅ unchanged

### History
* :green_heart: [Build #282030](https://buildkite.com/elastic/kibana-pull-request/builds/282030) succeeded b55cb2e228c13f4344b6333c37750e090912ee23
* :green_heart: [Build #281961](https://buildkite.com/elastic/kibana-pull-request/builds/281961) succeeded 33928d161a4cae0349173cadcf098bed73fd014f

<!--buildkite-pr-comment-kibana-pull-request
{"builds":[{"buildStatus":{"state":"passed","success":true,"hasRetries":true,"hasNonPreemptionRetries":false},"url":"https://buildkite.com/elastic/kibana-pull-request/builds/282760","number":282760,"commit":"4f9c179dcdede098b863e7fbc81f2e2fc1b82a27"},{"buildStatus":{"state":"passed","success":true,"hasRetries":false,"hasNonPreemptionRetries":false},"url":"https://buildkite.com/elastic/kibana-pull-request/builds/282030","number":282030,"commit":"b55cb2e228c13f4344b6333c37750e090912ee23"},{"buildStatus":{"state":"passed","success":true,"hasRetries":false,"hasNonPreemptionRetries":false},"url":"https://buildkite.com/elastic/kibana-pull-request/builds/281961","number":281961,"commit":"33928d161a4cae0349173cadcf098bed73fd014f"}],"number":282760}
buildkite-pr-comment-->
author_association: CONTRIBUTOR
user: login: kibanamachine
created_at: 2025-03-11T21:45:48Z
body: Starting backport for target branches: 8.x, 9.0

https://github.com/elastic/kibana/actions/runs/13798836730
author_association: CONTRIBUTOR
user: login: kibanamachine
created_at: 2025-03-11T21:51:37Z
body: ## 💚 All backports created successfully

| Status | Branch | Result |
|:------:|:------:|:------|
|✅|8.x|[<img src="https://img.shields.io/github/pulls/detail/state/elastic/kibana/214021">](https://github.com/elastic/kibana/pull/214021)|
|✅|9.0|[<img src="https://img.shields.io/github/pulls/detail/state/elastic/kibana/214022">](https://github.com/elastic/kibana/pull/214022)|

Note: Successful backport PRs will be merged automatically after passing CI.

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

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