From bcff7b63d2b2914e9ce1a8ee812f9c93e97fe3c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arturo=20Lidue=C3=B1a?= Date: Mon, 24 Mar 2025 22:31:42 +0100 Subject: [PATCH 1/4] add test: visualize_query --- .../functions/visualize_query.spec.ts | 97 +++++++++++++++++++ .../apis/observability/ai_assistant/index.ts | 1 + 2 files changed, 98 insertions(+) create mode 100644 x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts new file mode 100644 index 0000000000000..88b21e637e645 --- /dev/null +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts @@ -0,0 +1,97 @@ +/* + * 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 { MessageAddEvent, MessageRole } from '@kbn/observability-ai-assistant-plugin/common'; +import { + LlmProxy, + createLlmProxy, +} from '../../../../../../../observability_ai_assistant_api_integration/common/create_llm_proxy'; +import { + getMessageAddedEvents, + invokeChatCompleteWithFunctionRequest, +} from '../../utils/conversation'; +import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context'; + +export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { + const es = getService('es'); + const log = getService('log'); + const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); + + describe('visualize_query', function () { + this.tags(['failsOnMKI']); + let llmProxy: LlmProxy; + let connectorId: string; + let events: MessageAddEvent[]; + + before(async () => { + llmProxy = await createLlmProxy(log); + connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({ + port: llmProxy.getPort(), + }); + await es.index({ + index: 'test_index', + refresh: true, + id: 'index_id', + document: { bar: 'foo' }, + }); + void llmProxy.interceptConversation('Hello from LLM Proxy'); + const responseBody = await invokeChatCompleteWithFunctionRequest({ + connectorId, + observabilityAIAssistantAPIClient, + functionCall: { + name: 'visualize_query', + trigger: MessageRole.Assistant, + arguments: JSON.stringify({ query: 'FROM test_index', intention: 'visualizeAuto' }), + }, + }); + + await llmProxy.waitForAllInterceptorsToHaveBeenCalled(); + + events = getMessageAddedEvents(responseBody); + }); + + after(async () => { + await es.indices.delete({ + index: 'test_index', + }); + llmProxy.close(); + await observabilityAIAssistantAPIClient.deleteActionConnector({ + actionId: connectorId, + }); + }); + + it('should execute the visualize_query function and return expected messages', async () => { + expect(events.length).to.be(2); + }); + + it('should correctly structure the function response', async () => { + const functionResponse = events[0]; + expect(functionResponse.message.message.name).to.be('visualize_query'); + + const parsedResponse = JSON.parse(functionResponse.message.message.content!); + expect(parsedResponse.message).to.contain('Only following query is visualized'); + expect(parsedResponse.errorMessages).to.be.an('array'); + expect(parsedResponse.errorMessages.length).to.be(0); + }); + + it('should contain expected data structure in response', async () => { + const functionResponse = events[0]; + const parsedData = JSON.parse(functionResponse.message.message.data!); + + expect(parsedData).to.have.property('columns'); + expect(parsedData).to.have.property('rows'); + expect(parsedData).to.have.property('correctedQuery', 'FROM test_index'); + }); + + it('should return a valid LLM proxy response', async () => { + const llmResponse = events[1]; + expect(llmResponse.message.message.content).to.be('Hello from LLM Proxy'); + expect(llmResponse.message.message.role).to.be('assistant'); + }); + }); +} 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 46f7e22bf7698..06bb2af4d583a 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 @@ -22,6 +22,7 @@ export default function aiAssistantApiIntegrationTests({ loadTestFile(require.resolve('./complete/functions/retrieve_elastic_doc.spec.ts')); loadTestFile(require.resolve('./complete/functions/summarize.spec.ts')); loadTestFile(require.resolve('./complete/functions/title_conversation.spec.ts')); + loadTestFile(require.resolve('./complete/functions/visualize_query.spec.ts')); // knowledge base loadTestFile(require.resolve('./knowledge_base/knowledge_base_8.10_upgrade_test.spec.ts')); From 73a27cd48784a6e68d3893ebac646bfdd97091eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arturo=20Lidue=C3=B1a?= Date: Thu, 1 May 2025 22:55:12 +0200 Subject: [PATCH 2/4] update function visualize_query test --- .../functions/visualize_query.spec.ts | 29 +++++++------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts index 88b21e637e645..fd5180c755101 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts @@ -7,6 +7,7 @@ import expect from '@kbn/expect'; import { MessageAddEvent, MessageRole } from '@kbn/observability-ai-assistant-plugin/common'; +import { VisualizeESQLUserIntention } from '@kbn/observability-ai-assistant-plugin/common/functions/visualize_esql'; import { LlmProxy, createLlmProxy, @@ -27,6 +28,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon let llmProxy: LlmProxy; let connectorId: string; let events: MessageAddEvent[]; + const query = `FROM test_index`; before(async () => { llmProxy = await createLlmProxy(log); @@ -46,7 +48,10 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon functionCall: { name: 'visualize_query', trigger: MessageRole.Assistant, - arguments: JSON.stringify({ query: 'FROM test_index', intention: 'visualizeAuto' }), + arguments: JSON.stringify({ + query, + intention: VisualizeESQLUserIntention.visualizeAuto, + }), }, }); @@ -66,32 +71,20 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon }); it('should execute the visualize_query function and return expected messages', async () => { - expect(events.length).to.be(2); - }); - - it('should correctly structure the function response', async () => { const functionResponse = events[0]; expect(functionResponse.message.message.name).to.be('visualize_query'); const parsedResponse = JSON.parse(functionResponse.message.message.content!); - expect(parsedResponse.message).to.contain('Only following query is visualized'); - expect(parsedResponse.errorMessages).to.be.an('array'); - expect(parsedResponse.errorMessages.length).to.be(0); + expect(parsedResponse.message).to.contain(query); }); - it('should contain expected data structure in response', async () => { + it('should contain expected document data in response', async () => { const functionResponse = events[0]; const parsedData = JSON.parse(functionResponse.message.message.data!); - expect(parsedData).to.have.property('columns'); - expect(parsedData).to.have.property('rows'); - expect(parsedData).to.have.property('correctedQuery', 'FROM test_index'); - }); - - it('should return a valid LLM proxy response', async () => { - const llmResponse = events[1]; - expect(llmResponse.message.message.content).to.be('Hello from LLM Proxy'); - expect(llmResponse.message.message.role).to.be('assistant'); + expect(parsedData.columns[0].id).to.be('bar'); + expect(parsedData.rows[0][0]).to.be('foo'); + expect(parsedData).to.have.property('correctedQuery', query); }); }); } From e174d71ce86aa645b572afb45169b69c4cfb02ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arturo=20Lidue=C3=B1a?= Date: Tue, 6 May 2025 08:44:36 +0200 Subject: [PATCH 3/4] use interceptWithResponse llmProxy methos in visualize_query test --- .../ai_assistant/complete/functions/visualize_query.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts index fd5180c755101..6cc166b012cb1 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts @@ -41,7 +41,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon id: 'index_id', document: { bar: 'foo' }, }); - void llmProxy.interceptConversation('Hello from LLM Proxy'); + void llmProxy.interceptWithResponse('Hello from LLM Proxy'); const responseBody = await invokeChatCompleteWithFunctionRequest({ connectorId, observabilityAIAssistantAPIClient, From cc627f2d7050ba2dad5a7e78f0955dfab096539e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arturo=20Lidue=C3=B1a?= Date: Tue, 6 May 2025 09:13:04 +0200 Subject: [PATCH 4/4] replace 'failsOnMKI' tag with 'skipCloud' in visualize_query test --- .../ai_assistant/complete/functions/visualize_query.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts index 6cc166b012cb1..0929d5fa133fe 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/visualize_query.spec.ts @@ -24,7 +24,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); describe('visualize_query', function () { - this.tags(['failsOnMKI']); + this.tags(['skipCloud']); let llmProxy: LlmProxy; let connectorId: string; let events: MessageAddEvent[];