add API retrieve_elastic_doc tests (#214880)

Related: https://github.com/elastic/kibana/issues/180787

- Adds test for `retrieve_elastic_doc` function

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

@@ -0,0 +1,134 @@
+/*
+ * 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 expect from '@kbn/expect';
+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 { installProductDoc, uninstallProductDoc } from '../../utils/product_doc_base';
+
+export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
+  const log = getService('log');
+  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
+
+  describe('retrieve_elastic_doc', function () {
+    // Fails on MKI: https://github.com/elastic/kibana/issues/205581
+    this.tags(['failsOnMKI']);
+    const supertest = getService('supertest');
+    const USER_MESSAGE = 'What is Kibana Lens?';
+
+    describe('POST /internal/observability_ai_assistant/chat/complete without product doc installed', function () {
+      let llmProxy: LlmProxy;
+      let connectorId: string;
+
+      before(async () => {
+        llmProxy = await createLlmProxy(log);
+        connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+          port: llmProxy.getPort(),
+        });
+        void llmProxy.interceptConversation('Hello from LLM Proxy');
+
+        await chatComplete({
+          userPrompt: USER_MESSAGE,
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        });
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+      });
+
+      after(async () => {
+        llmProxy.close();
+        await observabilityAIAssistantAPIClient.deleteActionConnector({
+          actionId: connectorId,
+        });
+      });
+
+      it('makes 1 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(1);
+      });
+
+      it('not contain retrieve_elastic_doc function when product doc is not installed', () => {
+        expect(
+          llmProxy.interceptedRequests.flatMap(({ requestBody }) =>
+            requestBody.tools?.map((t) => t.function.name)
+          )
+        ).to.not.contain('retrieve_elastic_doc');
+      });
+
+      it('contains the original user message', () => {
+        const everyRequestHasUserMessage = llmProxy.interceptedRequests.every(({ requestBody }) =>
+          requestBody.messages.some(
+            (message) => message.role === 'user' && (message.content as string) === USER_MESSAGE
+          )
+        );
+        expect(everyRequestHasUserMessage).to.be(true);
+      });
+    });
+
+    // Calling `retrieve_elastic_doc` via the chat/complete endpoint
+    describe('POST /internal/observability_ai_assistant/chat/complete', function () {
+      let llmProxy: LlmProxy;
+      let connectorId: string;
+      before(async () => {
+        llmProxy = await createLlmProxy(log);
+        connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({
+          port: llmProxy.getPort(),
+        });
+        await installProductDoc(supertest);
+
+        void llmProxy.interceptWithFunctionRequest({
+          name: 'retrieve_elastic_doc',
+          arguments: () => JSON.stringify({}),
+          when: () => true,
+        });
+
+        void llmProxy.interceptConversation('Hello from LLM Proxy');
+
+        await chatComplete({
+          userPrompt: USER_MESSAGE,
+          connectorId,
+          observabilityAIAssistantAPIClient,
+        });
+
+        await llmProxy.waitForAllInterceptorsToHaveBeenCalled();
+      });
+
+      after(async () => {
+        await uninstallProductDoc(supertest);
+        llmProxy.close();
+        await observabilityAIAssistantAPIClient.deleteActionConnector({
+          actionId: connectorId,
+        });
+      });
+
+      it('makes 6 requests to the LLM', () => {
+        expect(llmProxy.interceptedRequests.length).to.be(6);
+      });
+
+      it('every request contain retrieve_elastic_doc function', () => {
+        const everyRequestHasRetrieveElasticDoc = llmProxy.interceptedRequests.every(
+          ({ requestBody }) =>
+            requestBody.tools?.some((t) => t.function.name === 'retrieve_elastic_doc')
+        );
+        expect(everyRequestHasRetrieveElasticDoc).to.be(true);
+      });
+
+      it('contains the original user message', () => {
+        const everyRequestHasUserMessage = llmProxy.interceptedRequests.every(({ requestBody }) =>
+          requestBody.messages.some(
+            (message) => message.role === 'user' && (message.content as string) === USER_MESSAGE
+          )
+        );
+        expect(everyRequestHasUserMessage).to.be(true);
+      });
+    });
+  });
+}

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

@@ -20,6 +20,7 @@ export default function aiAssistantApiIntegrationTests({
     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/retrieve_elastic_doc.spec.ts'));
     loadTestFile(require.resolve('./complete/functions/summarize.spec.ts'));
     loadTestFile(require.resolve('./complete/functions/recall.spec.ts'));
     loadTestFile(require.resolve('./public_complete/public_complete.spec.ts'));

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/product_doc_base.ts

@@ -0,0 +1,31 @@
+/*
+ * 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 {
+  ELASTIC_HTTP_VERSION_HEADER,
+  X_ELASTIC_INTERNAL_ORIGIN_REQUEST,
+} from '@kbn/core-http-common';
+
+import type SuperTest from 'supertest';
+
+export async function installProductDoc(supertest: SuperTest.Agent) {
+  return supertest
+    .post('/internal/product_doc_base/install')
+    .set(ELASTIC_HTTP_VERSION_HEADER, '1')
+    .set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana')
+    .set('kbn-xsrf', 'foo')
+    .expect(200);
+}
+
+export async function uninstallProductDoc(supertest: SuperTest.Agent) {
+  return supertest
+    .post('/internal/product_doc_base/uninstall')
+    .set(ELASTIC_HTTP_VERSION_HEADER, '1')
+    .set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana')
+    .set('kbn-xsrf', 'foo')
+    .expect(200);
+}