From 93695d8b0adf68659c84e1d8443de33b634a14fc Mon Sep 17 00:00:00 2001 From: Quynh Nguyen Date: Tue, 15 Jul 2025 07:38:43 -0500 Subject: [PATCH 1/6] Fix inference id missing retrievel documentation function --- .../plugins/shared/ai_infra/llm_tasks/README.md | 13 +++++++++++-- .../shared/ai_infra/llm_tasks/server/plugin.ts | 4 ++-- .../shared/ai_infra/llm_tasks/server/types.ts | 2 +- .../server/functions/documentation.ts | 13 ++++++------- 4 files changed, 20 insertions(+), 12 deletions(-) diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md index e019d456cd65a..9be8de3ff25a8 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md @@ -17,14 +17,22 @@ context. That API receive the inbound request as parameter. -Example: +Example, by default it will check with the default ELSER model: ```ts -if (await llmTasksStart.retrieveDocumentationAvailable({ request })) { +if (await llmTasksStart.retrieveDocumentationAvailable()) { // task is available } else { // task is not available } ``` +To check if documentation is available fora different Inference ID: + +```ts +if (await llmTasksStart.retrieveDocumentationAvailable({inferenceId})) { + // task is available +} else { + // task is not available +} ### Executing the task @@ -37,6 +45,7 @@ const result = await llmTasksStart.retrieveDocumentation({ searchTerm: "How to create a space in Kibana?", request, connectorId: 'my-connector-id', + inferenceId: 'my-inference-id', }); const { success, documents } = result; diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts index 50d9e2e340fab..ff1e82e55275f 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts @@ -41,9 +41,9 @@ export class LlmTasksPlugin start(core: CoreStart, startDependencies: PluginStartDependencies): LlmTasksPluginStart { const { inference, productDocBase } = startDependencies; return { - retrieveDocumentationAvailable: async () => { + retrieveDocumentationAvailable: async (options?: { inferenceId?: string }) => { const docBaseStatus = await startDependencies.productDocBase.management.getStatus({ - inferenceId: defaultInferenceEndpoints.ELSER, + inferenceId: options?.inferenceId ?? defaultInferenceEndpoints.ELSER, }); return docBaseStatus.status === 'installed'; }, diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts index d550e4398b509..7acdd6e0624f8 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts @@ -32,7 +32,7 @@ export interface LlmTasksPluginStart { * are respected. Can be used to check if the task can be registered * as LLM tool for example. */ - retrieveDocumentationAvailable: () => Promise; + retrieveDocumentationAvailable: (options?: { inferenceId: string }) => Promise; /** * Perform the `retrieveDocumentation` task. * diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/documentation.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/documentation.ts index 5f628d8c93859..182b7a6480a0c 100644 --- a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/documentation.ts +++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/documentation.ts @@ -17,7 +17,12 @@ export async function registerDocumentationFunction({ resources, pluginsStart: { llmTasks }, }: FunctionRegistrationParameters) { - const isProductDocAvailable = (await llmTasks.retrieveDocumentationAvailable()) ?? false; + const esClient = (await resources.context.core).elasticsearch.client; + const inferenceId = + (await getInferenceIdFromWriteIndex(esClient, resources.logger)) ?? + defaultInferenceEndpoints.ELSER; + const isProductDocAvailable = + (await llmTasks.retrieveDocumentationAvailable({ inferenceId })) ?? false; if (isProductDocAvailable) { functions.registerInstruction(({ availableFunctionNames }) => { @@ -65,12 +70,6 @@ export async function registerDocumentationFunction({ } as const, }, async ({ arguments: { query, product }, connectorId, simulateFunctionCalling }) => { - const esClient = (await resources.context.core).elasticsearch.client; - - const inferenceId = - (await getInferenceIdFromWriteIndex(esClient, resources.logger)) ?? - defaultInferenceEndpoints.ELSER; - const response = await llmTasks!.retrieveDocumentation({ searchTerm: query, products: product ? [product] : undefined, From acb8478dd95335f85467ad59a5aed90d20618bae Mon Sep 17 00:00:00 2001 From: Quynh Nguyen Date: Tue, 15 Jul 2025 09:04:21 -0500 Subject: [PATCH 2/6] Make inference id required --- .../plugins/shared/ai_infra/llm_tasks/README.md | 11 +---------- .../shared/ai_infra/llm_tasks/server/plugin.ts | 5 ++--- .../plugins/shared/ai_infra/llm_tasks/server/types.ts | 2 +- .../server/routes/chat/chat_complete_route.ts | 5 ++++- .../server/routes/evaluate/post_evaluate.ts | 5 ++++- .../server/routes/post_actions_connector_execute.ts | 5 ++++- 6 files changed, 16 insertions(+), 17 deletions(-) diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md index 9be8de3ff25a8..5c5bacf1029c6 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/README.md @@ -19,16 +19,7 @@ That API receive the inbound request as parameter. Example, by default it will check with the default ELSER model: ```ts -if (await llmTasksStart.retrieveDocumentationAvailable()) { - // task is available -} else { - // task is not available -} -``` -To check if documentation is available fora different Inference ID: - -```ts -if (await llmTasksStart.retrieveDocumentationAvailable({inferenceId})) { +if (await llmTasksStart.retrieveDocumentationAvailable({ inferenceId })) { // task is available } else { // task is not available diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts index ff1e82e55275f..14090d90e4387 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/plugin.ts @@ -7,7 +7,6 @@ import type { Logger } from '@kbn/logging'; import type { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '@kbn/core/server'; -import { defaultInferenceEndpoints } from '@kbn/inference-common'; import type { LlmTasksConfig } from './config'; import type { LlmTasksPluginSetup, @@ -41,9 +40,9 @@ export class LlmTasksPlugin start(core: CoreStart, startDependencies: PluginStartDependencies): LlmTasksPluginStart { const { inference, productDocBase } = startDependencies; return { - retrieveDocumentationAvailable: async (options?: { inferenceId?: string }) => { + retrieveDocumentationAvailable: async (options: { inferenceId: string }) => { const docBaseStatus = await startDependencies.productDocBase.management.getStatus({ - inferenceId: options?.inferenceId ?? defaultInferenceEndpoints.ELSER, + inferenceId: options.inferenceId, }); return docBaseStatus.status === 'installed'; }, diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts index 7acdd6e0624f8..5905070d4027d 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/types.ts @@ -32,7 +32,7 @@ export interface LlmTasksPluginStart { * are respected. Can be used to check if the task can be registered * as LLM tool for example. */ - retrieveDocumentationAvailable: (options?: { inferenceId: string }) => Promise; + retrieveDocumentationAvailable: (options: { inferenceId: string }) => Promise; /** * Perform the `retrieveDocumentation` task. * diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts index ea289abb14d1e..09dd6bccebb3d 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts @@ -22,6 +22,7 @@ import { } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { INVOKE_ASSISTANT_ERROR_EVENT } from '../../lib/telemetry/event_based_telemetry'; import { ElasticAssistantPluginRouter } from '../../types'; import { buildResponse } from '../../lib/build_response'; @@ -86,7 +87,9 @@ export const chatCompleteRoute = ( telemetry = ctx.elasticAssistant.telemetry; const inference = ctx.elasticAssistant.inference; const productDocsAvailable = - (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false; + (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable({ + inferenceId: defaultInferenceEndpoints.ELSER, + })) ?? false; // Perform license and authenticated user checks const checkResponse = await performChecks({ diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index e5634d9e4a870..72ecea9c5f3fb 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -30,6 +30,7 @@ import { getDefaultArguments } from '@kbn/langchain/server'; import { StructuredTool } from '@langchain/core/tools'; import { AgentFinish } from 'langchain/agents'; import { omit } from 'lodash/fp'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { getDefendInsightsPrompt } from '../../lib/defend_insights/graphs/default_defend_insights_graph/prompts'; import { evaluateDefendInsights } from '../../lib/defend_insights/evaluation'; import { localToolPrompts, promptGroupId as toolsGroupId } from '../../lib/prompt/tool_prompts'; @@ -176,7 +177,9 @@ export const postEvaluateRoute = ( const inference = ctx.elasticAssistant.inference; const productDocsAvailable = - (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false; + (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable({ + inferenceId: defaultInferenceEndpoints.ELSER, + })) ?? false; const { featureFlags } = await context.core; const inferenceChatModelEnabled = await featureFlags.getBooleanValue( diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index 1c7bbca8e1ca7..78913f395ac87 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -22,6 +22,7 @@ import { INFERENCE_CHAT_MODEL_ENABLED_FEATURE_FLAG, } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { getPrompt } from '../lib/prompt'; import { INVOKE_ASSISTANT_ERROR_EVENT } from '../lib/telemetry/event_based_telemetry'; import { buildResponse } from '../lib/build_response'; @@ -124,7 +125,9 @@ export const postActionsConnectorExecuteRoute = ( const inference = ctx.elasticAssistant.inference; const savedObjectsClient = ctx.elasticAssistant.savedObjectsClient; const productDocsAvailable = - (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable()) ?? false; + (await ctx.elasticAssistant.llmTasks.retrieveDocumentationAvailable({ + inferenceId: defaultInferenceEndpoints.ELSER, + })) ?? false; const actionsClient = await actions.getActionsClientWithRequest(request); const connectors = await actionsClient.getBulk({ ids: [connectorId] }); const connector = connectors.length > 0 ? connectors[0] : undefined; From 414ad954646527961da570e2619d190f0bcbbb48 Mon Sep 17 00:00:00 2001 From: Quynh Nguyen Date: Wed, 16 Jul 2025 06:24:48 -0500 Subject: [PATCH 3/6] Change to tiny elser --- .../src/is_default_inference_endpoint.ts | 3 ++- .../complete/functions/retrieve_elastic_doc.spec.ts | 9 ++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/is_default_inference_endpoint.ts b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/is_default_inference_endpoint.ts index a2c0bfc01d7a3..c7f6de079e2af 100644 --- a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/is_default_inference_endpoint.ts +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/is_default_inference_endpoint.ts @@ -17,6 +17,7 @@ export const isImpliedDefaultElserInferenceId = (inferenceId: string | null | un inferenceId === null || inferenceId === undefined || inferenceId === defaultInferenceEndpoints.ELSER || - inferenceId === defaultInferenceEndpoints.ELSER_IN_EIS_INFERENCE_ID + inferenceId === defaultInferenceEndpoints.ELSER_IN_EIS_INFERENCE_ID || + (typeof inferenceId === 'string' && inferenceId.toLowerCase().includes('elser')) ); }; diff --git a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts index 918388c7a4f55..48876c6968118 100644 --- a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts +++ b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts @@ -14,8 +14,12 @@ import { LlmProxy, createLlmProxy } from '../../utils/create_llm_proxy'; import { chatComplete } from '../../utils/conversation'; import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context'; import { installProductDoc, uninstallProductDoc } from '../../utils/product_doc_base'; +import { + deployTinyElserAndSetupKb, + teardownTinyElserModelAndInferenceEndpoint, +} from '../../utils/model_and_inference'; -const DEFAULT_INFERENCE_ID = '.elser-2-elasticsearch'; +const DEFAULT_INFERENCE_ID = 'tiny_elser_inference_id'; export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { const log = getService('log'); @@ -92,6 +96,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({ port: llmProxy.getPort(), }); + await deployTinyElserAndSetupKb(getService); + await installProductDoc(supertest, DEFAULT_INFERENCE_ID); void llmProxy.interceptWithFunctionRequest({ @@ -122,6 +128,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon await observabilityAIAssistantAPIClient.deleteActionConnector({ actionId: connectorId, }); + await teardownTinyElserModelAndInferenceEndpoint(getService); }); it('makes 2 requests to the LLM', () => { From 51bfba62e2bec1a1ff55b7c32adee2b68d4ca825 Mon Sep 17 00:00:00 2001 From: Quynh Nguyen Date: Wed, 16 Jul 2025 08:44:06 -0500 Subject: [PATCH 4/6] Fix tiny elser inference id --- .../shared/ai-infra/product-doc-common/src/artifact.ts | 3 ++- .../server/services/package_installer/package_installer.ts | 5 +++-- .../complete/functions/retrieve_elastic_doc.spec.ts | 6 +++--- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.ts b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.ts index d67138a8d39f0..42f2fb89136ca 100644 --- a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.ts +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.ts @@ -5,6 +5,7 @@ * 2.0. */ +import { isImpliedDefaultElserInferenceId } from './is_default_inference_endpoint'; import { type ProductName, DocumentationProduct } from './product'; const allowedProductNames: ProductName[] = Object.values(DocumentationProduct); @@ -24,7 +25,7 @@ export const getArtifactName = ({ }): string => { const ext = excludeExtension ? '' : '.zip'; return `kb-product-doc-${productName}-${productVersion}${ - inferenceId && inferenceId !== DEFAULT_ELSER ? `--${inferenceId}` : '' + inferenceId && !isImpliedDefaultElserInferenceId(inferenceId) ? `--${inferenceId}` : '' }${ext}`.toLowerCase(); }; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.ts index e9f23f30626c1..0332ef5606530 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.ts @@ -17,6 +17,7 @@ import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { cloneDeep } from 'lodash'; import type { InferenceInferenceEndpointInfo } from '@elastic/elasticsearch/lib/api/types'; import { i18n } from '@kbn/i18n'; +import { isImpliedDefaultElserInferenceId } from '@kbn/product-doc-common/src/is_default_inference_endpoint'; import type { ProductDocInstallClient } from '../doc_install_status'; import { downloadToDisk, @@ -179,7 +180,7 @@ export class PackageInstaller { inferenceId, }); - if (customInference && customInference?.inference_id !== this.elserInferenceId) { + if (customInference && !isImpliedDefaultElserInferenceId(customInference?.inference_id)) { if (customInference?.task_type !== 'text_embedding') { throw new Error( `Inference [${inferenceId}]'s task type ${customInference?.task_type} is not supported. Please use a model with task type 'text_embedding'.` @@ -191,7 +192,7 @@ export class PackageInstaller { }); } - if (!customInference || customInference?.inference_id === this.elserInferenceId) { + if (!customInference || isImpliedDefaultElserInferenceId(customInference?.inference_id)) { await ensureDefaultElserDeployed({ client: this.esClient, }); diff --git a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts index 48876c6968118..99f0f19081b13 100644 --- a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts +++ b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts @@ -10,6 +10,7 @@ import { ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream'; import { ChatCompletionMessageParam } from 'openai/resources'; import { last } from 'lodash'; import { MessageAddEvent, MessageRole } from '@kbn/observability-ai-assistant-plugin/common'; +import { TINY_ELSER_INFERENCE_ID } from '@kbn/test-suites-xpack/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/model_and_inference'; import { LlmProxy, createLlmProxy } from '../../utils/create_llm_proxy'; import { chatComplete } from '../../utils/conversation'; import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context'; @@ -19,13 +20,12 @@ import { teardownTinyElserModelAndInferenceEndpoint, } from '../../utils/model_and_inference'; -const DEFAULT_INFERENCE_ID = 'tiny_elser_inference_id'; - +const DEFAULT_INFERENCE_ID = TINY_ELSER_INFERENCE_ID; export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { const log = getService('log'); const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); - describe('tool: retrieve_elastic_doc', function () { + describe.only('tool: retrieve_elastic_doc', function () { // Fails on MKI: https://github.com/elastic/kibana/issues/205581 this.tags(['skipCloud']); const supertest = getService('supertest'); From 1fde492138f4fb1e66c9f608057cfe0002d556fb Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 16 Jul 2025 13:57:12 +0000 Subject: [PATCH 5/6] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- x-pack/solutions/observability/test/tsconfig.json | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/solutions/observability/test/tsconfig.json b/x-pack/solutions/observability/test/tsconfig.json index b6c03de9a1dbd..43f5fbab911fa 100644 --- a/x-pack/solutions/observability/test/tsconfig.json +++ b/x-pack/solutions/observability/test/tsconfig.json @@ -91,5 +91,6 @@ "@kbn/typed-react-router-config", "@kbn/observability-ai-assistant-app-plugin", "@kbn/dev-utils", + "@kbn/test-suites-xpack", ] } From a1b313f6de89e90d6ce0ddc7775eb30f399d4b49 Mon Sep 17 00:00:00 2001 From: Viduni Wickramarachchi Date: Wed, 16 Jul 2025 12:34:40 -0400 Subject: [PATCH 6/6] Update the retrieve_elastic_doc tool test to intercept query rewrite when the KB is installed --- .../functions/retrieve_elastic_doc.spec.ts | 37 +++++++++++-------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts index 99f0f19081b13..518e7128bd06c 100644 --- a/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts +++ b/x-pack/solutions/observability/test/api_integration_deployment_agnostic/apis/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts @@ -20,12 +20,11 @@ import { teardownTinyElserModelAndInferenceEndpoint, } from '../../utils/model_and_inference'; -const DEFAULT_INFERENCE_ID = TINY_ELSER_INFERENCE_ID; export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { const log = getService('log'); const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); - describe.only('tool: retrieve_elastic_doc', function () { + describe('tool: retrieve_elastic_doc', function () { // Fails on MKI: https://github.com/elastic/kibana/issues/205581 this.tags(['skipCloud']); const supertest = getService('supertest'); @@ -89,8 +88,9 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon let llmProxy: LlmProxy; let connectorId: string; let messageAddedEvents: MessageAddEvent[]; - let firstRequestBody: ChatCompletionStreamParams; - let secondRequestBody: ChatCompletionStreamParams; + let toolCallRequestBody: ChatCompletionStreamParams; + let userPromptRequestBody: ChatCompletionStreamParams; + before(async () => { llmProxy = await createLlmProxy(log); connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({ @@ -98,7 +98,9 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon }); await deployTinyElserAndSetupKb(getService); - await installProductDoc(supertest, DEFAULT_INFERENCE_ID); + await installProductDoc(supertest, TINY_ELSER_INFERENCE_ID); + + void llmProxy.interceptQueryRewrite('This is a rewritten user prompt.'); void llmProxy.interceptWithFunctionRequest({ name: 'retrieve_elastic_doc', @@ -118,12 +120,12 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon })); await llmProxy.waitForAllInterceptorsToHaveBeenCalled(); - firstRequestBody = llmProxy.interceptedRequests[0].requestBody; - secondRequestBody = llmProxy.interceptedRequests[1].requestBody; + toolCallRequestBody = llmProxy.interceptedRequests[1].requestBody; + userPromptRequestBody = llmProxy.interceptedRequests[2].requestBody; }); after(async () => { - await uninstallProductDoc(supertest, DEFAULT_INFERENCE_ID); + await uninstallProductDoc(supertest, TINY_ELSER_INFERENCE_ID); llmProxy.close(); await observabilityAIAssistantAPIClient.deleteActionConnector({ actionId: connectorId, @@ -131,8 +133,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon await teardownTinyElserModelAndInferenceEndpoint(getService); }); - it('makes 2 requests to the LLM', () => { - expect(llmProxy.interceptedRequests.length).to.be(2); + it('makes 3 requests to the LLM', () => { + expect(llmProxy.interceptedRequests.length).to.be(3); }); it('emits 5 messageAdded events', () => { @@ -141,8 +143,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon describe('The first request', () => { it('enables the LLM to call `retrieve_elastic_doc`', () => { - expect(firstRequestBody.tool_choice).to.be('auto'); - expect(firstRequestBody.tools?.map((t) => t.function.name)).to.contain( + expect(toolCallRequestBody.tool_choice).to.be('auto'); + expect(toolCallRequestBody.tools?.map((t) => t.function.name)).to.contain( 'retrieve_elastic_doc' ); }); @@ -151,14 +153,16 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon describe('The second request - Sending the user prompt', () => { let lastMessage: ChatCompletionMessageParam; let parsedContent: { documents: Array<{ title: string; content: string; url: string }> }; + before(() => { - lastMessage = last(secondRequestBody.messages) as ChatCompletionMessageParam; + lastMessage = last(userPromptRequestBody.messages) as ChatCompletionMessageParam; parsedContent = JSON.parse(lastMessage.content as string); }); + it('includes the retrieve_elastic_doc function call', () => { - expect(secondRequestBody.messages[4].role).to.be(MessageRole.Assistant); + expect(userPromptRequestBody.messages[4].role).to.be(MessageRole.Assistant); // @ts-expect-error - expect(secondRequestBody.messages[4].tool_calls[0].function.name).to.be( + expect(userPromptRequestBody.messages[4].tool_calls[0].function.name).to.be( 'retrieve_elastic_doc' ); }); @@ -168,9 +172,10 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon // @ts-expect-error expect(lastMessage?.tool_call_id).to.equal( // @ts-expect-error - secondRequestBody.messages[4].tool_calls[0].id + userPromptRequestBody.messages[4].tool_calls[0].id ); }); + it('sends the retrieved documents from Elastic docs to the LLM', () => { expect(lastMessage.content).to.be.a('string'); });