diff --git a/packages/kbn-check-mappings-update-cli/current_fields.json b/packages/kbn-check-mappings-update-cli/current_fields.json index 814f3404fd66f..3c1c290916b5b 100644 --- a/packages/kbn-check-mappings-update-cli/current_fields.json +++ b/packages/kbn-check-mappings-update-cli/current_fields.json @@ -915,6 +915,7 @@ ], "product-doc-install-status": [ "index_name", + "inference_id", "installation_status", "last_installation_date", "product_name", diff --git a/packages/kbn-check-mappings-update-cli/current_mappings.json b/packages/kbn-check-mappings-update-cli/current_mappings.json index 7378ec08bd0dd..4a285a6cf9b96 100644 --- a/packages/kbn-check-mappings-update-cli/current_mappings.json +++ b/packages/kbn-check-mappings-update-cli/current_mappings.json @@ -3042,6 +3042,9 @@ "index_name": { "type": "keyword" }, + "inference_id": { + "type": "keyword" + }, "installation_status": { "type": "keyword" }, diff --git a/src/core/server/integration_tests/ci_checks/saved_objects/check_registered_types.test.ts b/src/core/server/integration_tests/ci_checks/saved_objects/check_registered_types.test.ts index 497d9b37676b9..d31eb1f418ed2 100644 --- a/src/core/server/integration_tests/ci_checks/saved_objects/check_registered_types.test.ts +++ b/src/core/server/integration_tests/ci_checks/saved_objects/check_registered_types.test.ts @@ -154,7 +154,7 @@ describe('checking migration metadata changes on all registered SO types', () => "osquery-saved-query": "a8ef11610473e3d1b51a8fdacb2799d8a610818e", "policy-settings-protection-updates-note": "c05c4c33a5e5bd1fa153991f300d040ac5d6f38d", "privilege-monitoring-status": "4daec76df427409bcd64250f5c23f5ab86c8bac3", - "product-doc-install-status": "ee7817c45bf1c41830290c8ef535e726c86f7c19", + "product-doc-install-status": "f94e3e5ad2cc933df918f2cd159044c626e01011", "query": "1966ccce8e9853018111fb8a1dee500228731d9e", "risk-engine-configuration": "533a0a3f2dbef1c95129146ec4d5714de305be1a", "rules-settings": "53f94e5ce61f5e75d55ab8adbc1fb3d0937d2e0b", diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/README.md b/x-pack/packages/ai-infra/product-doc-artifact-builder/README.md index 49949def3e5e7..0a4d8de5a204e 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/README.md +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/README.md @@ -8,17 +8,23 @@ Script to build the knowledge base artifacts. node scripts/build_product_doc_artifacts.js --stack-version {version} --product-name {product} ``` +Example: + +``` +node scripts/build_product_doc_artifacts.js --product-name=security --stack-version=8.18 --inference-id=.multilingual-e5-small-elasticsearch +``` + ### parameters -#### `stack-version`: +#### `stack-version`: the stack version to generate the artifacts for. -#### `product-name`: +#### `product-name`: (multi-value) the list of products to generate artifacts for. -possible values: +possible values: - "kibana" - "elasticsearch" - "observability" @@ -34,6 +40,11 @@ Defaults to `{REPO_ROOT}/build-kb-artifacts`. The folder to use for temporary files. +#### inference-id: + +The inference endpoint to use to generate the embeddings. If the inference ID provided and is not the ELSER default, the artifacts will be generated with `{artifactName}--{inference-id}.zip`. Note the double dash before inference-id. + + Defaults to `{REPO_ROOT}/build/temp-kb-artifacts` #### Cluster infos @@ -46,4 +57,7 @@ Defaults to `{REPO_ROOT}/build/temp-kb-artifacts` - params for the embedding cluster: `embeddingClusterUrl` / env.KIBANA_EMBEDDING_CLUSTER_URL `embeddingClusterUsername` / env.KIBANA_EMBEDDING_CLUSTER_USERNAME -`embeddingClusterPassword` / env.KIBANA_EMBEDDING_CLUSTER_PASSWORD \ No newline at end of file +`embeddingClusterPassword` / env.KIBANA_EMBEDDING_CLUSTER_PASSWORD + +- params for the inference endpoint: +`inferenceId` diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/artifact/mappings.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/artifact/mappings.ts index 979845ec31844..763288f49f52a 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/artifact/mappings.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/artifact/mappings.ts @@ -6,30 +6,31 @@ */ import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; +import { + DEFAULT_ELSER, + getSemanticTextMapping, + type SemanticTextMapping, +} from '../tasks/create_index'; -export const getArtifactMappings = (inferenceEndpoint: string): MappingTypeMapping => { +export const getArtifactMappings = ( + customSemanticTextMapping?: SemanticTextMapping +): MappingTypeMapping => { + const semanticTextMapping = customSemanticTextMapping + ? customSemanticTextMapping + : getSemanticTextMapping(DEFAULT_ELSER); return { dynamic: 'strict', properties: { content_title: { type: 'text' }, - content_body: { - type: 'semantic_text', - inference_id: inferenceEndpoint, - }, + content_body: semanticTextMapping, product_name: { type: 'keyword' }, root_type: { type: 'keyword' }, slug: { type: 'keyword' }, url: { type: 'keyword' }, version: { type: 'version' }, ai_subtitle: { type: 'text' }, - ai_summary: { - type: 'semantic_text', - inference_id: inferenceEndpoint, - }, - ai_questions_answered: { - type: 'semantic_text', - inference_id: inferenceEndpoint, - }, + ai_summary: semanticTextMapping, + ai_questions_answered: semanticTextMapping, ai_tags: { type: 'keyword' }, }, }; diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/build_artifacts.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/build_artifacts.ts index 805f35c4460ae..371c72bd02aef 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/build_artifacts.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/build_artifacts.ts @@ -7,8 +7,14 @@ import Path from 'path'; import { Client, HttpConnection } from '@elastic/elasticsearch'; +import { + Client as ElasticsearchClient8, + HttpConnection as Elasticsearch8HttpConnection, +} from 'elasticsearch-8.x'; + import { ToolingLog } from '@kbn/tooling-log'; import type { ProductName } from '@kbn/product-doc-common'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { // checkConnectivity, createTargetIndex, @@ -21,9 +27,10 @@ import { processDocuments, } from './tasks'; import type { TaskConfig } from './types'; +import { getSemanticTextMapping } from './tasks/create_index'; const getSourceClient = (config: TaskConfig) => { - return new Client({ + return new ElasticsearchClient8({ compression: true, nodes: [config.sourceClusterUrl], sniffOnStart: false, @@ -31,7 +38,7 @@ const getSourceClient = (config: TaskConfig) => { username: config.sourceClusterUsername, password: config.sourceClusterPassword, }, - Connection: HttpConnection, + Connection: Elasticsearch8HttpConnection, requestTimeout: 30_000, }); }; @@ -79,6 +86,7 @@ export const buildArtifacts = async (config: TaskConfig) => { sourceClient, embeddingClient, log, + inferenceId: config.inferenceId ?? defaultInferenceEndpoints.ELSER, }); } @@ -93,18 +101,41 @@ const buildArtifact = async ({ embeddingClient, sourceClient, log, + inferenceId, }: { productName: ProductName; stackVersion: string; buildFolder: string; targetFolder: string; - sourceClient: Client; + sourceClient: ElasticsearchClient8; embeddingClient: Client; log: ToolingLog; + inferenceId: string; }) => { - log.info(`Starting building artifact for product [${productName}] and version [${stackVersion}]`); + log.info( + `Starting building artifact for product [${productName}] and version [${stackVersion}] with inference id [${inferenceId}]` + ); - const targetIndex = getTargetIndexName({ productName, stackVersion }); + const semanticTextMapping = getSemanticTextMapping(inferenceId); + + log.info( + `Detected semantic text mapping for Inference ID ${inferenceId}:\n ${JSON.stringify( + semanticTextMapping, + null, + 2 + )}` + ); + + const targetIndex = getTargetIndexName({ + productName, + stackVersion, + inferenceId: semanticTextMapping?.inference_id, + }); + await deleteIndex({ + indexName: targetIndex, + client: embeddingClient, + log, + }); let documents = await extractDocumentation({ client: sourceClient, @@ -119,6 +150,7 @@ const buildArtifact = async ({ await createTargetIndex({ client: embeddingClient, indexName: targetIndex, + semanticTextMapping, }); await indexDocuments({ @@ -142,12 +174,7 @@ const buildArtifact = async ({ productName, stackVersion, log, - }); - - await deleteIndex({ - indexName: targetIndex, - client: embeddingClient, - log, + semanticTextMapping, }); log.info(`Finished building artifact for product [${productName}] and version [${stackVersion}]`); @@ -156,9 +183,13 @@ const buildArtifact = async ({ const getTargetIndexName = ({ productName, stackVersion, + inferenceId, }: { productName: string; stackVersion: string; + inferenceId?: string; }) => { - return `kb-artifact-builder-${productName}-${stackVersion}`.toLowerCase(); + return `kb-artifact-builder-${productName}-${stackVersion}${ + inferenceId ? `-${inferenceId}` : '' + }`.toLowerCase(); }; diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/command.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/command.ts index e8d0d9486e331..7e4ebda200f25 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/command.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/command.ts @@ -71,6 +71,10 @@ function options(y: yargs.Argv) { demandOption: true, default: process.env.KIBANA_EMBEDDING_CLUSTER_PASSWORD, }) + .option('inferenceId', { + describe: 'The inference id to use for the artifacts', + string: true, + }) .locale('en'); } @@ -89,6 +93,7 @@ export function runScript() { embeddingClusterUrl: argv.embeddingClusterUrl!, embeddingClusterUsername: argv.embeddingClusterUsername!, embeddingClusterPassword: argv.embeddingClusterPassword!, + inferenceId: argv.inferenceId, }; return buildArtifacts(taskConfig); diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_artifact.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_artifact.ts index bd6d005936574..b0468163ecce6 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_artifact.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_artifact.ts @@ -15,7 +15,7 @@ import { } from '@kbn/product-doc-common'; import { getArtifactMappings } from '../artifact/mappings'; import { getArtifactManifest } from '../artifact/manifest'; -import { DEFAULT_ELSER } from './create_index'; +import { DEFAULT_ELSER, SemanticTextMapping } from './create_index'; export const createArtifact = async ({ productName, @@ -23,12 +23,14 @@ export const createArtifact = async ({ buildFolder, targetFolder, log, + semanticTextMapping, }: { buildFolder: string; targetFolder: string; productName: ProductName; stackVersion: string; log: ToolingLog; + semanticTextMapping?: SemanticTextMapping; }) => { log.info( `Starting to create artifact from build folder [${buildFolder}] into target [${targetFolder}]` @@ -36,7 +38,9 @@ export const createArtifact = async ({ const zip = new AdmZip(); - const mappings = getArtifactMappings(DEFAULT_ELSER); + const inferenceId = semanticTextMapping?.inference_id || DEFAULT_ELSER; + + const mappings = getArtifactMappings(semanticTextMapping); const mappingFileContent = JSON.stringify(mappings, undefined, 2); zip.addFile('mappings.json', Buffer.from(mappingFileContent, 'utf-8')); @@ -53,6 +57,7 @@ export const createArtifact = async ({ const artifactName = getArtifactName({ productName, productVersion: stackVersion, + inferenceId, }); zip.writeZip(Path.join(targetFolder, artifactName)); diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_index.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_index.ts index b867edc31b85a..e1b501a341af9 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_index.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/create_index.ts @@ -6,43 +6,66 @@ */ import type { Client } from '@elastic/elasticsearch'; -import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; +import { getArtifactMappings } from '../artifact/mappings'; export const DEFAULT_ELSER = '.elser-2-elasticsearch'; +export const DEFAULT_E5_SMALL = '.multilingual-e5-small-elasticsearch'; -const mappings: MappingTypeMapping = { - dynamic: 'strict', - properties: { - content_title: { type: 'text' }, - content_body: { - type: 'semantic_text', - inference_id: DEFAULT_ELSER, - }, - product_name: { type: 'keyword' }, - root_type: { type: 'keyword' }, - slug: { type: 'keyword' }, - url: { type: 'keyword' }, - version: { type: 'version' }, - ai_subtitle: { type: 'text' }, - ai_summary: { - type: 'semantic_text', - inference_id: DEFAULT_ELSER, - }, - ai_questions_answered: { - type: 'semantic_text', - inference_id: DEFAULT_ELSER, +interface BaseSemanticTextMapping { + type: 'semantic_text'; + inference_id: string; +} +export interface SemanticTextMapping extends BaseSemanticTextMapping { + model_settings?: { + service?: string; + task_type?: string; + dimensions?: number; + similarity?: string; + element_type?: string; + }; +} + +type SupportedInferenceId = typeof DEFAULT_E5_SMALL | typeof DEFAULT_ELSER; +const isSupportedInferenceId = (inferenceId: string): inferenceId is SupportedInferenceId => { + return inferenceId === DEFAULT_E5_SMALL || inferenceId === DEFAULT_ELSER; +}; + +const INFERENCE_ID_TO_SEMANTIC_TEXT_MAPPING: Record = { + [DEFAULT_E5_SMALL]: { + type: 'semantic_text', + inference_id: DEFAULT_E5_SMALL, + model_settings: { + service: 'elasticsearch', + task_type: 'text_embedding', + dimensions: 384, + similarity: 'cosine', + element_type: 'float', }, - ai_tags: { type: 'keyword' }, }, + [DEFAULT_ELSER]: { + type: 'semantic_text', + inference_id: DEFAULT_ELSER, + }, +}; +export const getSemanticTextMapping = (inferenceId: string): SemanticTextMapping => { + if (isSupportedInferenceId(inferenceId)) { + return INFERENCE_ID_TO_SEMANTIC_TEXT_MAPPING[inferenceId]; + } + throw new Error(`Semantic text mapping for Inference ID ${inferenceId} not found`); }; export const createTargetIndex = async ({ indexName, client, + semanticTextMapping, }: { indexName: string; client: Client; + semanticTextMapping?: SemanticTextMapping; }) => { + const mappings = semanticTextMapping + ? getArtifactMappings(semanticTextMapping) + : getArtifactMappings(getSemanticTextMapping(DEFAULT_ELSER)); await client.indices.create({ index: indexName, mappings, diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/extract_documentation.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/extract_documentation.ts index 6aa8bb49b0cfd..4f5e519837811 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/extract_documentation.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/tasks/extract_documentation.ts @@ -5,7 +5,7 @@ * 2.0. */ -import type { Client } from '@elastic/elasticsearch'; +import { Client as ElasticsearchClient8 } from 'elasticsearch-8.x'; import type { SearchHit } from '@elastic/elasticsearch/lib/api/types'; import type { ToolingLog } from '@kbn/tooling-log'; import type { ProductName } from '@kbn/product-doc-common'; @@ -64,7 +64,7 @@ export const extractDocumentation = async ({ productName, log, }: { - client: Client; + client: ElasticsearchClient8; index: string; stackVersion: string; productName: ProductName; diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/types.ts b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/types.ts index 1eb4a4348d218..82bbebabf6fe1 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/src/types.ts +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/src/types.ts @@ -18,4 +18,5 @@ export interface TaskConfig { embeddingClusterUrl: string; embeddingClusterUsername: string; embeddingClusterPassword: string; + inferenceId?: string; } diff --git a/x-pack/packages/ai-infra/product-doc-artifact-builder/tsconfig.json b/x-pack/packages/ai-infra/product-doc-artifact-builder/tsconfig.json index 68ff27852c4d1..7e6082692c495 100644 --- a/x-pack/packages/ai-infra/product-doc-artifact-builder/tsconfig.json +++ b/x-pack/packages/ai-infra/product-doc-artifact-builder/tsconfig.json @@ -17,5 +17,6 @@ "@kbn/tooling-log", "@kbn/repo-info", "@kbn/product-doc-common", + "@kbn/inference-common", ] } diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/inference_endpoints.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/inference_endpoints.ts index 1fc2116d3b585..e3e908c7d20ff 100644 --- a/x-pack/platform/packages/shared/ai-infra/inference-common/src/inference_endpoints.ts +++ b/x-pack/platform/packages/shared/ai-infra/inference-common/src/inference_endpoints.ts @@ -10,6 +10,7 @@ */ export const defaultInferenceEndpoints = { ELSER: '.elser-2-elasticsearch', + ELSER_IN_EIS_INFERENCE_ID: '.elser-v2-elastic', MULTILINGUAL_E5_SMALL: '.multilingual-e5-small-elasticsearch', } as const; diff --git a/x-pack/platform/packages/shared/ai-infra/product-doc-common/kibana.jsonc b/x-pack/platform/packages/shared/ai-infra/product-doc-common/kibana.jsonc index ed5332676de2e..b39155b2a51d8 100644 --- a/x-pack/platform/packages/shared/ai-infra/product-doc-common/kibana.jsonc +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/kibana.jsonc @@ -6,4 +6,5 @@ ], "group": "platform", "visibility": "shared" -} \ No newline at end of file + +} diff --git a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.test.ts b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.test.ts index 2b6362dbf4aad..b4a763ea9e9db 100644 --- a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.test.ts +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/artifact.test.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { getArtifactName, parseArtifactName } from './artifact'; +import { getArtifactName, parseArtifactName, DEFAULT_ELSER } from './artifact'; describe('getArtifactName', () => { it('builds the name based on the provided product name and version', () => { @@ -37,6 +37,32 @@ describe('getArtifactName', () => { }) ).toEqual('kb-product-doc-elasticsearch-8.17'); }); + it('generates a name with inference id when inference_id is not the ELSER default', () => { + expect( + getArtifactName({ + productName: 'kibana', + productVersion: '8.16', + inferenceId: '.multilingual-e5-small-elasticsearch', + }) + ).toEqual('kb-product-doc-kibana-8.16--.multilingual-e5-small-elasticsearch.zip'); + expect( + getArtifactName({ + productName: 'kibana', + productVersion: '8.16', + inferenceId: '.multilingual-e5-small-elasticsearch', + excludeExtension: true, + }) + ).toEqual('kb-product-doc-kibana-8.16--.multilingual-e5-small-elasticsearch'); + }); + it('generates a name with inference id when inference_id is the ELSER default', () => { + expect( + getArtifactName({ + productName: 'kibana', + productVersion: '8.16', + inferenceId: DEFAULT_ELSER, + }) + ).toEqual('kb-product-doc-kibana-8.16.zip'); + }); }); describe('parseArtifactName', () => { @@ -61,4 +87,22 @@ describe('parseArtifactName', () => { it('returns undefined if the provided string is not strictly lowercase', () => { expect(parseArtifactName('kb-product-doc-Security-8.17')).toEqual(undefined); }); + it('parses an artifact name with inference id and extension', () => { + expect( + parseArtifactName('kb-product-doc-kibana-8.16--.multilingual-e5-small-elasticsearch.zip') + ).toEqual({ + productName: 'kibana', + productVersion: '8.16', + inferenceId: '.multilingual-e5-small-elasticsearch', + }); + }); + it('parses an artifact name with inference id when it is not the default', () => { + expect( + parseArtifactName('kb-product-doc-kibana-8.16--.multilingual-e5-small-elasticsearch') + ).toEqual({ + productName: 'kibana', + productVersion: '8.16', + inferenceId: '.multilingual-e5-small-elasticsearch', + }); + }); }); 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 1a6745abd733d..68cc21aafd684 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 @@ -9,30 +9,44 @@ import { type ProductName, DocumentationProduct } from './product'; // kb-product-doc-elasticsearch-8.15.zip const artifactNameRegexp = /^kb-product-doc-([a-z]+)-([0-9]+\.[0-9]+)(\.zip)?$/; +const inferenceIdRegexp = /--([a-z0-9.-]+)(\.zip)?$/; const allowedProductNames: ProductName[] = Object.values(DocumentationProduct); +export const DEFAULT_ELSER = '.elser-2-elasticsearch'; + export const getArtifactName = ({ productName, productVersion, excludeExtension = false, + inferenceId, }: { productName: ProductName; productVersion: string; excludeExtension?: boolean; + inferenceId?: string; }): string => { const ext = excludeExtension ? '' : '.zip'; - return `kb-product-doc-${productName}-${productVersion}${ext}`.toLowerCase(); + return `kb-product-doc-${productName}-${productVersion}${ + inferenceId && inferenceId !== DEFAULT_ELSER ? `--${inferenceId}` : '' + }${ext}`.toLowerCase(); }; export const parseArtifactName = (artifactName: string) => { - const match = artifactNameRegexp.exec(artifactName); + let name = artifactName.replace(/\.zip$/, ''); + // First, extract out the inference Id which is prefixed by -- + const inferenceIdMatch = name.match(inferenceIdRegexp); + name = inferenceIdMatch ? name.replace(inferenceIdMatch[0], '') : artifactName; + + const match = name.match(artifactNameRegexp); if (match) { const productName = match[1].toLowerCase() as ProductName; const productVersion = match[2].toLowerCase(); + const inferenceId = inferenceIdMatch ? inferenceIdMatch[1] : undefined; if (allowedProductNames.includes(productName)) { return { productName, productVersion, + ...(inferenceId ? { inferenceId } : {}), }; } } diff --git a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/indices.ts b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/indices.ts index 90e416ff48c46..a7f264938dab8 100644 --- a/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/indices.ts +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/indices.ts @@ -5,11 +5,14 @@ * 2.0. */ +import { isImpliedDefaultElserInferenceId } from './is_default_inference_endpoint'; import type { ProductName } from './product'; export const productDocIndexPrefix = '.kibana_ai_product_doc'; export const productDocIndexPattern = `${productDocIndexPrefix}_*`; -export const getProductDocIndexName = (productName: ProductName): string => { - return `${productDocIndexPrefix}_${productName.toLowerCase()}`; +export const getProductDocIndexName = (productName: ProductName, inferenceId?: string): string => { + return `${productDocIndexPrefix}_${productName.toLowerCase()}${ + !isImpliedDefaultElserInferenceId(inferenceId) ? `-${inferenceId}` : '' + }`; }; 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 new file mode 100644 index 0000000000000..a2c0bfc01d7a3 --- /dev/null +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/src/is_default_inference_endpoint.ts @@ -0,0 +1,22 @@ +/* + * 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 { defaultInferenceEndpoints } from '@kbn/inference-common'; + +/** + * Returns true if inferenceId is not provided, or when provided, it is a default ELSER inference ID + * @param inferenceId + * @returns + */ +export const isImpliedDefaultElserInferenceId = (inferenceId: string | null | undefined) => { + return ( + inferenceId === null || + inferenceId === undefined || + inferenceId === defaultInferenceEndpoints.ELSER || + inferenceId === defaultInferenceEndpoints.ELSER_IN_EIS_INFERENCE_ID + ); +}; diff --git a/x-pack/platform/packages/shared/ai-infra/product-doc-common/tsconfig.json b/x-pack/platform/packages/shared/ai-infra/product-doc-common/tsconfig.json index 63f0b5ff33faa..d40c9dc9a23d1 100644 --- a/x-pack/platform/packages/shared/ai-infra/product-doc-common/tsconfig.json +++ b/x-pack/platform/packages/shared/ai-infra/product-doc-common/tsconfig.json @@ -13,5 +13,7 @@ "exclude": [ "target/**/*" ], - "kbn_references": [] + "kbn_references": [ + "@kbn/inference-common" + ] } diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_current_inference_id.ts b/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_current_inference_id.ts new file mode 100644 index 0000000000000..bd369ca3196ac --- /dev/null +++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/hooks/use_current_inference_id.ts @@ -0,0 +1,19 @@ +/* + * 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 { useMemo } from 'react'; +import { useKnowledgeBase } from './use_knowledge_base'; + +export const useCurrentlyDeployedInferenceId = () => { + const knowledgeBase = useKnowledgeBase(); + return useMemo( + () => + knowledgeBase.status.value?.currentInferenceId ?? + knowledgeBase.status.value?.endpoint?.inference_id, + [knowledgeBase.status.value] + ); +}; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_get_product_doc_status.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_get_product_doc_status.ts index 5e516a4207cfa..1e6b5663da9b3 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_get_product_doc_status.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_get_product_doc_status.ts @@ -6,6 +6,7 @@ */ import { useQuery } from '@tanstack/react-query'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { REACT_QUERY_KEYS } from './const'; import { useAssistantContext } from '../../../..'; @@ -15,7 +16,9 @@ export function useGetProductDocStatus() { const { isLoading, isError, isSuccess, isRefetching, data, refetch } = useQuery({ queryKey: [REACT_QUERY_KEYS.GET_PRODUCT_DOC_STATUS], queryFn: async () => { - return productDocBase.installation.getStatus(); + return productDocBase.installation.getStatus({ + inferenceId: defaultInferenceEndpoints.ELSER, + }); }, keepPreviousData: false, refetchOnWindowFocus: false, diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.test.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.test.ts index 3b3c12d6b9dc8..d2e16a6799c91 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.test.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.test.ts @@ -9,6 +9,7 @@ import { waitFor, renderHook } from '@testing-library/react'; import { useInstallProductDoc } from './use_install_product_doc'; import { useAssistantContext } from '../../../..'; import { TestProviders } from '../../../mock/test_providers/test_providers'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; jest.mock('../../../..', () => ({ useAssistantContext: jest.fn(), @@ -39,7 +40,7 @@ describe('useInstallProductDoc', () => { wrapper: TestProviders, }); - result.current.mutate(); + result.current.mutate(defaultInferenceEndpoints.ELSER); await waitFor(() => result.current.isSuccess); expect(mockAddSuccess).toHaveBeenCalledWith( @@ -54,7 +55,7 @@ describe('useInstallProductDoc', () => { wrapper: TestProviders, }); - result.current.mutate(); + result.current.mutate(defaultInferenceEndpoints.ELSER); await waitFor(() => result.current.isError); expect(mockAddError).toHaveBeenCalledWith( diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.ts index b17dab7826c48..f84b13bcf9ddf 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/api/product_docs/use_install_product_doc.ts @@ -11,17 +11,16 @@ import type { IHttpFetchError, ResponseErrorBody } from '@kbn/core/public'; import type { PerformInstallResponse } from '@kbn/product-doc-base-plugin/common/http_api/installation'; import { REACT_QUERY_KEYS } from './const'; import { useAssistantContext } from '../../../..'; - type ServerError = IHttpFetchError; export function useInstallProductDoc() { const { productDocBase, toasts } = useAssistantContext(); const queryClient = useQueryClient(); - return useMutation( + return useMutation( [REACT_QUERY_KEYS.INSTALL_PRODUCT_DOC], - () => { - return productDocBase.installation.install(); + (inferenceId: string) => { + return productDocBase.installation.install({ inferenceId }); }, { onSuccess: () => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.test.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.test.tsx index d8d865b9353f4..b425a07c19281 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.test.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.test.tsx @@ -9,6 +9,7 @@ import { render, screen, fireEvent, waitFor } from '@testing-library/react'; import { ProductDocumentationManagement } from '.'; import * as i18n from './translations'; import { useInstallProductDoc } from '../../api/product_docs/use_install_product_doc'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; jest.mock('../../api/product_docs/use_install_product_doc'); jest.mock('../../api/product_docs/use_get_product_doc_status'); @@ -26,12 +27,22 @@ describe('ProductDocumentationManagement', () => { }); it('renders install button when not installed', () => { - render(); + render( + + ); expect(screen.getByText(i18n.INSTALL)).toBeInTheDocument(); }); it('does not render anything when already installed', () => { - const { container } = render(); + const { container } = render( + + ); expect(container).toBeEmptyDOMElement(); }); @@ -41,7 +52,12 @@ describe('ProductDocumentationManagement', () => { isLoading: false, isSuccess: false, }); - const { container } = render(); + const { container } = render( + + ); expect(container).toBeEmptyDOMElement(); }); @@ -51,7 +67,12 @@ describe('ProductDocumentationManagement', () => { isLoading: true, isSuccess: false, }); - render(); + render( + + ); expect(screen.getByTestId('installing')).toBeInTheDocument(); expect(screen.getByText(i18n.INSTALLING)).toBeInTheDocument(); }); @@ -63,7 +84,12 @@ describe('ProductDocumentationManagement', () => { isSuccess: true, }); mockInstallProductDoc.mockResolvedValueOnce({}); - render(); + render( + + ); expect(screen.queryByText(i18n.INSTALL)).not.toBeInTheDocument(); }); @@ -74,7 +100,12 @@ describe('ProductDocumentationManagement', () => { isSuccess: false, }); mockInstallProductDoc.mockRejectedValueOnce(new Error('Installation failed')); - render(); + render( + + ); fireEvent.click(screen.getByText(i18n.INSTALL)); await waitFor(() => expect(screen.getByText(i18n.INSTALL)).toBeInTheDocument()); }); diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.tsx index 98a487ec9d187..6eeebbda3aa96 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/settings/product_documentation/index.tsx @@ -21,7 +21,8 @@ import * as i18n from './translations'; export const ProductDocumentationManagement = React.memo<{ status?: InstallationStatus; -}>(({ status }) => { + inferenceId: string; +}>(({ status, inferenceId }) => { const { mutateAsync: installProductDoc, isSuccess: isInstalled, @@ -29,8 +30,8 @@ export const ProductDocumentationManagement = React.memo<{ } = useInstallProductDoc(); const onClickInstall = useCallback(() => { - installProductDoc(); - }, [installProductDoc]); + installProductDoc(inferenceId); + }, [installProductDoc, inferenceId]); const content = useMemo(() => { if (isInstalling) { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index b600d3eb5efeb..ee9ad482cbb3a 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -32,6 +32,7 @@ import { css } from '@emotion/react'; import { DataViewsContract } from '@kbn/data-views-plugin/public'; import useAsync from 'react-use/lib/useAsync'; import { useSearchParams } from 'react-router-dom-v5-compat'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { useKnowledgeBaseUpdater } from '../../assistant/settings/use_settings_updater/use_knowledge_base_updater'; import { ProductDocumentationManagement } from '../../assistant/settings/product_documentation'; import { KnowledgeBaseTour } from '../../tour/knowledge_base'; @@ -338,7 +339,10 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ d return ( <> - + { - return productDocBase!.installation.getStatus(); + return productDocBase!.installation.getStatus({ inferenceId }); }, keepPreviousData: false, refetchOnWindowFocus: false, diff --git a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_install_product_doc.ts b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_install_product_doc.ts index cb32efa7e3908..7fabf7c3306fc 100644 --- a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_install_product_doc.ts +++ b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_install_product_doc.ts @@ -20,11 +20,10 @@ export function useInstallProductDoc() { notifications: { toasts }, } = useKibana().services; const queryClient = useQueryClient(); - - return useMutation( + return useMutation( [REACT_QUERY_KEYS.INSTALL_PRODUCT_DOC], - () => { - return productDocBase!.installation.install(); + (inferenceId: string) => { + return productDocBase!.installation.install({ inferenceId }); }, { onSuccess: () => { diff --git a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_uninstall_product_doc.ts b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_uninstall_product_doc.ts index 4aa3b5423faa1..8646b92c6b8d9 100644 --- a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_uninstall_product_doc.ts +++ b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/hooks/use_uninstall_product_doc.ts @@ -21,10 +21,10 @@ export function useUninstallProductDoc() { } = useKibana().services; const queryClient = useQueryClient(); - return useMutation( + return useMutation( [REACT_QUERY_KEYS.UNINSTALL_PRODUCT_DOC], - () => { - return productDocBase!.installation.uninstall(); + (inferenceId: string) => { + return productDocBase!.installation.uninstall({ inferenceId }); }, { onSuccess: () => { diff --git a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/change_kb_model.tsx b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/change_kb_model.tsx index 2697639f95c7c..11b9d7b828719 100644 --- a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/change_kb_model.tsx +++ b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/change_kb_model.tsx @@ -26,13 +26,20 @@ import { } from '@kbn/ai-assistant/src/utils/get_model_options_for_inference_endpoints'; import { useInferenceEndpoints, UseKnowledgeBaseResult } from '@kbn/ai-assistant/src/hooks'; import { KnowledgeBaseState, useKibana } from '@kbn/observability-ai-assistant-plugin/public'; +import { useInstallProductDoc } from '../../../hooks/use_install_product_doc'; export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBaseResult }) { const { overlays } = useKibana().services; + const currentlyDeployedInferenceId = knowledgeBase.status.value?.currentInferenceId; + + const [selectedInferenceId, setSelectedInferenceId] = useState( + currentlyDeployedInferenceId || '' + ); + const [hasLoadedCurrentModel, setHasLoadedCurrentModel] = useState(false); - const [selectedInferenceId, setSelectedInferenceId] = useState(''); const [isUpdatingModel, setIsUpdatingModel] = useState(false); + const { mutateAsync: installProductDoc } = useInstallProductDoc(); const { inferenceEndpoints, isLoading: isLoadingEndpoints, error } = useInferenceEndpoints(); @@ -44,8 +51,7 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa knowledgeBase.status?.value?.kbState === KnowledgeBaseState.MODEL_PENDING_ALLOCATION || knowledgeBase.status?.value?.kbState === KnowledgeBaseState.MODEL_PENDING_DEPLOYMENT; - const isSelectedModelCurrentModel = - selectedInferenceId === knowledgeBase.status?.value?.endpoint?.inference_id; + const isSelectedModelCurrentModel = selectedInferenceId === currentlyDeployedInferenceId; const isKnowledgeBaseInLoadingState = knowledgeBase.isInstalling || @@ -55,11 +61,16 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa useEffect(() => { if (!hasLoadedCurrentModel && modelOptions?.length && knowledgeBase.status?.value) { - const currentModel = knowledgeBase.status.value.currentInferenceId; - setSelectedInferenceId(currentModel || modelOptions[0].key); + setSelectedInferenceId(currentlyDeployedInferenceId || modelOptions[0].key); setHasLoadedCurrentModel(true); } - }, [hasLoadedCurrentModel, modelOptions, knowledgeBase.status?.value, selectedInferenceId]); + }, [ + hasLoadedCurrentModel, + modelOptions, + knowledgeBase.status?.value, + setSelectedInferenceId, + currentlyDeployedInferenceId, + ]); useEffect(() => { if (isUpdatingModel && !knowledgeBase.isInstalling && !knowledgeBase.isPolling) { @@ -145,6 +156,7 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa if (isConfirmed) { setIsUpdatingModel(true); knowledgeBase.install(selectedInferenceId); + installProductDoc(selectedInferenceId); } }); } @@ -156,6 +168,7 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa isSelectedModelCurrentModel, overlays, confirmationMessages, + installProductDoc, ]); const superSelectOptions = modelOptions.map((option: ModelOptionsData) => ({ @@ -227,6 +240,7 @@ export function ChangeKbModel({ knowledgeBase }: { knowledgeBase: UseKnowledgeBa isLoadingEndpoints, superSelectOptions, selectedInferenceId, + setSelectedInferenceId, isKnowledgeBaseInLoadingState, doesModelNeedRedeployment, knowledgeBase.status?.value?.kbState, diff --git a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/product_doc_entry.tsx b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/product_doc_entry.tsx index 668e363d071ee..afafee70b9a4f 100644 --- a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/product_doc_entry.tsx +++ b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/product_doc_entry.tsx @@ -18,6 +18,7 @@ import { } from '@elastic/eui'; import { i18n } from '@kbn/i18n'; import { FormattedMessage } from '@kbn/i18n-react'; +import { useKnowledgeBase } from '@kbn/ai-assistant'; import { useKibana } from '../../../hooks/use_kibana'; import { useGetProductDocStatus } from '../../../hooks/use_get_product_doc_status'; import { useInstallProductDoc } from '../../../hooks/use_install_product_doc'; @@ -26,22 +27,29 @@ import { useUninstallProductDoc } from '../../../hooks/use_uninstall_product_doc export function ProductDocEntry() { const { overlays } = useKibana().services; - const [isInstalled, setInstalled] = useState(true); + const knowledgeBase = useKnowledgeBase(); + const selectedInferenceId: string | undefined = knowledgeBase.status.value?.currentInferenceId; + + const [isInstalled, setInstalled] = useState(false); const [isInstalling, setInstalling] = useState(false); const { mutateAsync: installProductDoc } = useInstallProductDoc(); const { mutateAsync: uninstallProductDoc } = useUninstallProductDoc(); - const { status, isLoading: isStatusLoading } = useGetProductDocStatus(); + const { status, isLoading: isStatusLoading } = useGetProductDocStatus(selectedInferenceId); useEffect(() => { + if (isStatusLoading) return; if (status) { - setInstalled(status.overall === 'installed'); + setInstalled(status.overall === 'installed' && status.inferenceId === selectedInferenceId); } - }, [status]); + }, [selectedInferenceId, status, isStatusLoading]); const onClickInstall = useCallback(() => { + if (!selectedInferenceId) { + throw new Error('Inference ID is required to install product documentation'); + } setInstalling(true); - installProductDoc().then( + installProductDoc(selectedInferenceId).then( () => { setInstalling(false); setInstalled(true); @@ -51,7 +59,7 @@ export function ProductDocEntry() { setInstalled(false); } ); - }, [installProductDoc]); + }, [installProductDoc, selectedInferenceId]); const onClickUninstall = useCallback(() => { overlays @@ -72,18 +80,18 @@ export function ProductDocEntry() { } ) .then((confirmed) => { - if (confirmed) { - uninstallProductDoc().then(() => { + if (confirmed && selectedInferenceId) { + uninstallProductDoc(selectedInferenceId).then(() => { setInstalling(false); setInstalled(false); }); } }); - }, [overlays, uninstallProductDoc]); + }, [overlays, uninstallProductDoc, selectedInferenceId]); const content = useMemo(() => { if (isStatusLoading) { - return <>; + return ; } if (isInstalling) { return ( diff --git a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/settings_tab.test.tsx b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/settings_tab.test.tsx index fef2c84101c81..f23a56c099527 100644 --- a/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/settings_tab.test.tsx +++ b/x-pack/platform/plugins/private/observability_ai_assistant_management/public/routes/components/settings_tab/settings_tab.test.tsx @@ -48,6 +48,11 @@ describe('SettingsTab', () => { basePath: { prepend: prependMock }, }, productDocBase: undefined, + notifications: { + toasts: { + add: jest.fn(), + }, + }, }, }); useKnowledgeBaseMock.mockReturnValue({ 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 d10c495ece159..50d9e2e340fab 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,6 +7,7 @@ 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,7 +42,9 @@ export class LlmTasksPlugin const { inference, productDocBase } = startDependencies; return { retrieveDocumentationAvailable: async () => { - const docBaseStatus = await startDependencies.productDocBase.management.getStatus(); + const docBaseStatus = await startDependencies.productDocBase.management.getStatus({ + inferenceId: defaultInferenceEndpoints.ELSER, + }); return docBaseStatus.status === 'installed'; }, retrieveDocumentation: (options) => { diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.test.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.test.ts index 0214a6bee1a3e..f15da795fba47 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.test.ts @@ -16,6 +16,7 @@ const truncateMock = truncate as jest.MockedFn; const countTokensMock = countTokens as jest.MockedFn; import { summarizeDocument } from './summarize_document'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; jest.mock('./summarize_document'); const summarizeDocumentMock = summarizeDocument as jest.MockedFn; @@ -65,6 +66,7 @@ describe('retrieveDocumentation', () => { max: 5, connectorId: '.my-connector', functionCalling: 'simulated', + inferenceId: defaultInferenceEndpoints.ELSER, }); expect(result).toEqual({ @@ -78,6 +80,7 @@ describe('retrieveDocumentation', () => { products: ['kibana'], max: 5, highlights: 4, + inferenceId: defaultInferenceEndpoints.ELSER, }); }); @@ -92,6 +95,7 @@ describe('retrieveDocumentation', () => { max: 5, connectorId: '.my-connector', functionCalling: 'simulated', + inferenceId: defaultInferenceEndpoints.ELSER, }); expect(searchDocAPI).toHaveBeenCalledTimes(1); @@ -100,6 +104,7 @@ describe('retrieveDocumentation', () => { products: ['kibana'], max: 5, highlights: 0, + inferenceId: defaultInferenceEndpoints.ELSER, }); }); @@ -127,6 +132,7 @@ describe('retrieveDocumentation', () => { connectorId: '.my-connector', maxDocumentTokens: 100, tokenReductionStrategy: 'highlight', + inferenceId: defaultInferenceEndpoints.ELSER, }); expect(result.documents.length).toEqual(3); @@ -162,6 +168,7 @@ describe('retrieveDocumentation', () => { connectorId: '.my-connector', maxDocumentTokens: 100, tokenReductionStrategy: 'truncate', + inferenceId: defaultInferenceEndpoints.ELSER, }); expect(result.documents.length).toEqual(3); @@ -201,6 +208,7 @@ describe('retrieveDocumentation', () => { connectorId: '.my-connector', maxDocumentTokens: 100, tokenReductionStrategy: 'summarize', + inferenceId: defaultInferenceEndpoints.ELSER, }); expect(result.documents.length).toEqual(3); @@ -230,6 +238,7 @@ describe('retrieveDocumentation', () => { connectorId: '.my-connector', maxDocumentTokens: 100, tokenReductionStrategy: 'summarize', + inferenceId: defaultInferenceEndpoints.ELSER, }); expect(result).toEqual({ diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.ts index 38472c9b51647..91fe9c586f057 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/retrieve_documentation.ts @@ -6,7 +6,7 @@ */ import type { Logger } from '@kbn/logging'; -import type { OutputAPI } from '@kbn/inference-common'; +import { type OutputAPI } from '@kbn/inference-common'; import type { ProductDocSearchAPI, DocSearchResult } from '@kbn/product-doc-base-plugin/server'; import { truncate, count as countTokens } from '../../utils/tokens'; import type { RetrieveDocumentationAPI } from './types'; @@ -30,6 +30,7 @@ export const retrieveDocumentation = connectorId, products, functionCalling, + inferenceId, max = MAX_DOCUMENTS_DEFAULT, maxDocumentTokens = MAX_TOKENS_DEFAULT, tokenReductionStrategy = 'highlight', @@ -65,6 +66,7 @@ export const retrieveDocumentation = products, max, highlights, + inferenceId, }); log.debug(`searching with term=[${searchTerm}] returned ${results.length} documents`); diff --git a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/types.ts b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/types.ts index 78bf58bbce87e..dc65ba7d41318 100644 --- a/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/types.ts +++ b/x-pack/platform/plugins/shared/ai_infra/llm_tasks/server/tasks/retrieve_documentation/types.ts @@ -58,6 +58,10 @@ export interface RetrieveDocumentationParams { * Optional functionCalling parameter to pass down to the inference APIs. */ functionCalling?: FunctionCallingMode; + /** + * Inferece ID to route the request to the right index to perform the search. + */ + inferenceId: string; } /** diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/http_api/installation.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/http_api/installation.ts index 0237bd2c3b488..627e5b08e72c6 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/http_api/installation.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/http_api/installation.ts @@ -13,14 +13,20 @@ export const INSTALL_ALL_API_PATH = '/internal/product_doc_base/install'; export const UNINSTALL_ALL_API_PATH = '/internal/product_doc_base/uninstall'; export interface InstallationStatusResponse { + inferenceId: string; overall: InstallationStatus; perProducts: Record; } export interface PerformInstallResponse { installed: boolean; + failureReason?: string; } export interface UninstallResponse { success: boolean; } + +export interface ProductDocInstallParams { + inferenceId: string | undefined; +} diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/install_status.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/install_status.ts index 81102d43c1ff3..20625b268ebdd 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/install_status.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/common/install_status.ts @@ -20,9 +20,11 @@ export interface ProductDocInstallStatus { lastInstallationDate: Date | undefined; lastInstallationFailureReason: string | undefined; indexName?: string; + inferenceId?: string; } export interface ProductInstallState { status: InstallationStatus; version?: string; + failureReason?: string; } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/plugin.tsx b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/plugin.tsx index 6f2c989b6e45d..52a93611f60bd 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/plugin.tsx +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/plugin.tsx @@ -16,6 +16,10 @@ import type { } from './types'; import { InstallationService } from './services/installation'; +interface ProductDocInstallServiceParams { + inferenceId: string; +} + export class ProductDocBasePlugin implements Plugin< @@ -42,9 +46,11 @@ export class ProductDocBasePlugin return { installation: { - getStatus: () => installationService.getInstallationStatus(), - install: () => installationService.install(), - uninstall: () => installationService.uninstall(), + getStatus: (params: ProductDocInstallServiceParams) => + installationService.getInstallationStatus(params), + install: (params: ProductDocInstallServiceParams) => installationService.install(params), + uninstall: (params: ProductDocInstallServiceParams) => + installationService.uninstall(params), }, }; } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.test.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.test.ts index 294aeb99e0fd8..52181af38d5fe 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.test.ts @@ -12,6 +12,9 @@ import { INSTALL_ALL_API_PATH, UNINSTALL_ALL_API_PATH, } from '../../../common/http_api/installation'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; + +const inferenceId = defaultInferenceEndpoints.ELSER; describe('InstallationService', () => { let http: ReturnType; @@ -24,17 +27,32 @@ describe('InstallationService', () => { describe('#getInstallationStatus', () => { it('calls the endpoint with the right parameters', async () => { - await service.getInstallationStatus(); + await service.getInstallationStatus({ inferenceId }); expect(http.get).toHaveBeenCalledTimes(1); - expect(http.get).toHaveBeenCalledWith(INSTALLATION_STATUS_API_PATH); + expect(http.get).toHaveBeenCalledWith(INSTALLATION_STATUS_API_PATH, { + query: { + inferenceId, + }, + }); }); it('returns the value from the server', async () => { const expected = { stubbed: true }; http.get.mockResolvedValue(expected); - const response = await service.getInstallationStatus(); + const response = await service.getInstallationStatus({ inferenceId }); expect(response).toEqual(expected); }); + it('calls the endpoint with the right parameters for different inference IDs', async () => { + await service.getInstallationStatus({ + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }); + expect(http.get).toHaveBeenCalledTimes(1); + expect(http.get).toHaveBeenCalledWith(INSTALLATION_STATUS_API_PATH, { + query: { + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }, + }); + }); }); describe('#install', () => { beforeEach(() => { @@ -42,37 +60,68 @@ describe('InstallationService', () => { }); it('calls the endpoint with the right parameters', async () => { - await service.install(); + await service.install({ inferenceId }); + expect(http.post).toHaveBeenCalledTimes(1); + expect(http.post).toHaveBeenCalledWith(INSTALL_ALL_API_PATH, { + body: JSON.stringify({ + inferenceId, + }), + }); + }); + it('calls the endpoint with the right parameters for different inference IDs', async () => { + await service.install({ + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }); expect(http.post).toHaveBeenCalledTimes(1); - expect(http.post).toHaveBeenCalledWith(INSTALL_ALL_API_PATH); + expect(http.post).toHaveBeenCalledWith(INSTALL_ALL_API_PATH, { + body: JSON.stringify({ + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }), + }); }); it('returns the value from the server', async () => { const expected = { installed: true }; http.post.mockResolvedValue(expected); - const response = await service.install(); + const response = await service.install({ inferenceId }); expect(response).toEqual(expected); }); it('throws when the server returns installed: false', async () => { const expected = { installed: false }; http.post.mockResolvedValue(expected); - await expect(service.install()).rejects.toThrowErrorMatchingInlineSnapshot( - `"Installation did not complete successfully"` + await expect(service.install({ inferenceId })).rejects.toThrowErrorMatchingInlineSnapshot( + `"Installation did not complete successfully."` ); }); }); describe('#uninstall', () => { it('calls the endpoint with the right parameters', async () => { - await service.uninstall(); + await service.uninstall({ inferenceId }); + expect(http.post).toHaveBeenCalledTimes(1); + expect(http.post).toHaveBeenCalledWith(UNINSTALL_ALL_API_PATH, { + body: JSON.stringify({ + inferenceId, + }), + }); + }); + it('calls the endpoint with the right parameters for different inference IDs', async () => { + await service.uninstall({ + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }); expect(http.post).toHaveBeenCalledTimes(1); - expect(http.post).toHaveBeenCalledWith(UNINSTALL_ALL_API_PATH); + expect(http.post).toHaveBeenCalledWith(UNINSTALL_ALL_API_PATH, { + body: JSON.stringify({ + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }), + }); }); + it('returns the value from the server', async () => { const expected = { stubbed: true }; http.post.mockResolvedValue(expected); - const response = await service.uninstall(); + const response = await service.uninstall({ inferenceId }); expect(response).toEqual(expected); }); }); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.ts index ff347f52cb531..d46399749f588 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/installation_service.ts @@ -6,6 +6,7 @@ */ import type { HttpSetup } from '@kbn/core-http-browser'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { INSTALLATION_STATUS_API_PATH, INSTALL_ALL_API_PATH, @@ -22,19 +23,42 @@ export class InstallationService { this.http = http; } - async getInstallationStatus(): Promise { - return await this.http.get(INSTALLATION_STATUS_API_PATH); + async getInstallationStatus(params: { + inferenceId: string; + }): Promise { + const inferenceId = params?.inferenceId ?? defaultInferenceEndpoints.ELSER; + + const response = await this.http.get(INSTALLATION_STATUS_API_PATH, { + query: { inferenceId }, + }); + + return response; } - async install(): Promise { - const response = await this.http.post(INSTALL_ALL_API_PATH); + async install(params: { inferenceId: string }): Promise { + const inferenceId = params?.inferenceId ?? defaultInferenceEndpoints.ELSER; + + const response = await this.http.post(INSTALL_ALL_API_PATH, { + body: JSON.stringify({ inferenceId }), + }); + if (!response.installed) { - throw new Error('Installation did not complete successfully'); + throw new Error( + `Installation did not complete successfully.${ + response.failureReason ? `\n${response.failureReason}` : '' + }` + ); } return response; } - async uninstall(): Promise { - return await this.http.post(UNINSTALL_ALL_API_PATH); + async uninstall(params: { inferenceId: string }): Promise { + const inferenceId = params?.inferenceId ?? defaultInferenceEndpoints.ELSER; + + const response = await this.http.post(UNINSTALL_ALL_API_PATH, { + body: JSON.stringify({ inferenceId }), + }); + + return response; } } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/types.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/types.ts index 5c01c84b24625..5bf369c3b58e4 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/types.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/public/services/installation/types.ts @@ -9,10 +9,11 @@ import type { InstallationStatusResponse, PerformInstallResponse, UninstallResponse, + ProductDocInstallParams, } from '../../../common/http_api/installation'; export interface InstallationAPI { - getStatus(): Promise; - install(): Promise; - uninstall(): Promise; + getStatus(params: ProductDocInstallParams): Promise; + install(params: ProductDocInstallParams): Promise; + uninstall(params: ProductDocInstallParams): Promise; } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/plugin.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/plugin.ts index 53bc6acae4cfc..d3bff7235e9a0 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/plugin.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/plugin.ts @@ -10,6 +10,7 @@ import type { Logger } from '@kbn/logging'; import { getDataPath } from '@kbn/utils'; import type { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '@kbn/core/server'; import { SavedObjectsClient } from '@kbn/core/server'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { productDocInstallStatusSavedObjectTypeName } from '../common/consts'; import type { ProductDocBaseConfig } from './config'; import { @@ -76,7 +77,10 @@ export class ProductDocBasePlugin const soClient = new SavedObjectsClient( core.savedObjects.createInternalRepository([productDocInstallStatusSavedObjectTypeName]) ); - const productDocClient = new ProductDocInstallClient({ soClient }); + const productDocClient = new ProductDocInstallClient({ + soClient, + log: this.logger, + }); const packageInstaller = new PackageInstaller({ esClient: core.elasticsearch.client.asInternalUser, @@ -110,7 +114,7 @@ export class ProductDocBasePlugin taskManager, }; - documentationManager.update().catch((err) => { + documentationManager.update({ inferenceId: defaultInferenceEndpoints.ELSER }).catch((err) => { this.logger.error(`Error scheduling product documentation update task: ${err.message}`); }); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/routes/installation.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/routes/installation.ts index 98d8b5e0d85d7..af04ae85def9e 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/routes/installation.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/routes/installation.ts @@ -7,6 +7,8 @@ import type { IRouter } from '@kbn/core/server'; import { ApiPrivileges } from '@kbn/core-security-server'; +import { schema } from '@kbn/config-schema'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { INSTALLATION_STATUS_API_PATH, INSTALL_ALL_API_PATH, @@ -16,6 +18,7 @@ import { UninstallResponse, } from '../../common/http_api/installation'; import type { InternalServices } from '../types'; +import { ProductInstallState } from '../../common/install_status'; export const registerInstallationRoutes = ({ router, @@ -27,7 +30,11 @@ export const registerInstallationRoutes = ({ router.get( { path: INSTALLATION_STATUS_API_PATH, - validate: false, + validate: { + query: schema.object({ + inferenceId: schema.string({ defaultValue: defaultInferenceEndpoints.ELSER }), + }), + }, options: { access: 'internal', }, @@ -39,11 +46,17 @@ export const registerInstallationRoutes = ({ }, async (ctx, req, res) => { const { installClient, documentationManager } = getServices(); - const installStatus = await installClient.getInstallationStatus(); - const { status: overallStatus } = await documentationManager.getStatus(); + const inferenceId = req.query?.inferenceId; + const installStatus = await installClient.getInstallationStatus({ + inferenceId, + }); + const { status: overallStatus } = await documentationManager.getStatus({ + inferenceId, + }); return res.ok({ body: { + inferenceId, perProducts: installStatus, overall: overallStatus, }, @@ -54,7 +67,11 @@ export const registerInstallationRoutes = ({ router.post( { path: INSTALL_ALL_API_PATH, - validate: false, + validate: { + body: schema.object({ + inferenceId: schema.string({ defaultValue: defaultInferenceEndpoints.ELSER }), + }), + }, options: { access: 'internal', timeout: { idleSocket: 20 * 60 * 1000 }, // install can take time. @@ -68,18 +85,33 @@ export const registerInstallationRoutes = ({ async (ctx, req, res) => { const { documentationManager } = getServices(); + const inferenceId = req.body?.inferenceId; + await documentationManager.install({ request: req, force: false, wait: true, + inferenceId, }); // check status after installation in case of failure - const { status } = await documentationManager.getStatus(); + const { status, installStatus } = await documentationManager.getStatus({ + inferenceId, + }); + let failureReason = null; + if (status === 'error' && installStatus) { + failureReason = Object.values(installStatus) + .filter( + (product: ProductInstallState) => product.status === 'error' && product.failureReason + ) + .map((product: ProductInstallState) => product.failureReason) + .join('\n'); + } return res.ok({ body: { installed: status === 'installed', + ...(failureReason ? { failureReason } : {}), }, }); } @@ -88,7 +120,11 @@ export const registerInstallationRoutes = ({ router.post( { path: UNINSTALL_ALL_API_PATH, - validate: false, + validate: { + body: schema.object({ + inferenceId: schema.string({ defaultValue: defaultInferenceEndpoints.ELSER }), + }), + }, options: { access: 'internal', }, @@ -104,6 +140,7 @@ export const registerInstallationRoutes = ({ await documentationManager.uninstall({ request: req, wait: true, + inferenceId: req.body?.inferenceId, }); return res.ok({ diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/saved_objects/product_doc_install.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/saved_objects/product_doc_install.ts index 47cf7eb50cdd1..c18add6994ae1 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/saved_objects/product_doc_install.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/saved_objects/product_doc_install.ts @@ -7,6 +7,7 @@ import type { SavedObjectsType } from '@kbn/core/server'; import type { ProductName } from '@kbn/product-doc-common'; +import type { SavedObjectsModelVersion } from '@kbn/core-saved-objects-server'; import { productDocInstallStatusSavedObjectTypeName } from '../../common/consts'; import type { InstallationStatus } from '../../common/install_status'; @@ -22,7 +23,18 @@ export interface ProductDocInstallStatusAttributes { last_installation_date?: number; last_installation_failure_reason?: string; index_name?: string; + inference_id?: string; } +const modelVersion1: SavedObjectsModelVersion = { + changes: [ + { + type: 'mappings_addition', + addedMappings: { + inference_id: { type: 'keyword' }, + }, + }, + ], +}; export const productDocInstallStatusSavedObjectType: SavedObjectsType = { @@ -37,10 +49,11 @@ export const productDocInstallStatusSavedObjectType: SavedObjectsType => { return { id: attrs.product_name, @@ -24,10 +26,12 @@ const createObj = (attrs: TypeAttributes): SavedObjectsFindResult { let soClient: ReturnType; let service: ProductDocInstallClient; + let log: Logger; beforeEach(() => { soClient = savedObjectsClientMock.create(); - service = new ProductDocInstallClient({ soClient }); + log = loggingSystemMock.createLogger(); + service = new ProductDocInstallClient({ soClient, log }); }); describe('getInstallationStatus', () => { @@ -50,7 +54,7 @@ describe('ProductDocInstallClient', () => { page: 1, }); - const installStatus = await service.getInstallationStatus(); + const installStatus = await service.getInstallationStatus({ inferenceId }); expect(Object.keys(installStatus).sort()).toEqual(Object.keys(DocumentationProduct).sort()); expect(installStatus.kibana).toEqual({ diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/product_doc_install_service.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/product_doc_install_service.ts index 24625ebc51586..8572e12352eea 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/product_doc_install_service.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/product_doc_install_service.ts @@ -8,74 +8,130 @@ import type { SavedObjectsClientContract } from '@kbn/core/server'; import { SavedObjectsErrorHelpers } from '@kbn/core-saved-objects-server'; import { ProductName, DocumentationProduct } from '@kbn/product-doc-common'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; +import type { Logger } from '@kbn/logging'; +import { isImpliedDefaultElserInferenceId } from '@kbn/product-doc-common/src/is_default_inference_endpoint'; import type { ProductInstallState } from '../../../common/install_status'; import { productDocInstallStatusSavedObjectTypeName as typeName } from '../../../common/consts'; import type { ProductDocInstallStatusAttributes as TypeAttributes } from '../../saved_objects'; export class ProductDocInstallClient { private soClient: SavedObjectsClientContract; + private log: Logger; - constructor({ soClient }: { soClient: SavedObjectsClientContract }) { + constructor({ soClient, log }: { soClient: SavedObjectsClientContract; log: Logger }) { this.soClient = soClient; + this.log = log; } - async getInstallationStatus(): Promise> { - const response = await this.soClient.find({ + async getPreviouslyInstalledInferenceIds(): Promise { + const query = { type: typeName, perPage: 100, - }); + }; + const response = await this.soClient.find(query); + const inferenceIds = new Set( + response?.saved_objects.map( + (so) => so.attributes?.inference_id ?? defaultInferenceEndpoints.ELSER + ) + ); + return Array.from(inferenceIds); + } + async getInstallationStatus({ + inferenceId, + }: { + inferenceId: string; + }): Promise> { + const query = { + type: typeName, + perPage: 100, + }; const installStatus = Object.values(DocumentationProduct).reduce((memo, product) => { memo[product] = { status: 'uninstalled' }; return memo; }, {} as Record); + try { + const response = await this.soClient.find(query); + const savedObjects = isImpliedDefaultElserInferenceId(inferenceId) + ? response?.saved_objects.filter((so) => + isImpliedDefaultElserInferenceId(so.attributes.inference_id) + ) + : response?.saved_objects.filter((so) => so.attributes.inference_id === inferenceId); - response.saved_objects.forEach(({ attributes }) => { - installStatus[attributes.product_name as ProductName] = { - status: attributes.installation_status, - version: attributes.product_version, - }; - }); + savedObjects?.forEach(({ attributes }) => { + installStatus[attributes.product_name as ProductName] = { + status: attributes.installation_status, + version: attributes.product_version, + ...(attributes.last_installation_failure_reason + ? { failureReason: attributes.last_installation_failure_reason } + : {}), + }; + }); - return installStatus; + return installStatus; + } catch (error) { + this.log.error( + `An error occurred getting installation status saved object for inferenceId [${inferenceId}] + Query: ${JSON.stringify(query, null, 2)}`, + error + ); + return installStatus; + } } - async setInstallationStarted(fields: { productName: ProductName; productVersion: string }) { - const { productName, productVersion } = fields; - const objectId = getObjectIdFromProductName(productName); + async setInstallationStarted(fields: { + productName: ProductName; + productVersion: string; + inferenceId: string | undefined; + }) { + const { productName, productVersion, inferenceId } = fields; + const objectId = getObjectIdFromProductName(productName, inferenceId); const attributes = { product_name: productName, product_version: productVersion, installation_status: 'installing' as const, last_installation_failure_reason: '', + inference_id: inferenceId, }; await this.soClient.update(typeName, objectId, attributes, { upsert: attributes, }); } - async setInstallationSuccessful(productName: ProductName, indexName: string) { - const objectId = getObjectIdFromProductName(productName); + async setInstallationSuccessful( + productName: ProductName, + indexName: string, + inferenceId: string | undefined + ) { + const objectId = getObjectIdFromProductName(productName, inferenceId); await this.soClient.update(typeName, objectId, { installation_status: 'installed', index_name: indexName, + inference_id: inferenceId, }); } - async setInstallationFailed(productName: ProductName, failureReason: string) { - const objectId = getObjectIdFromProductName(productName); + async setInstallationFailed( + productName: ProductName, + failureReason: string, + inferenceId: string | undefined + ) { + const objectId = getObjectIdFromProductName(productName, inferenceId); await this.soClient.update(typeName, objectId, { installation_status: 'error', last_installation_failure_reason: failureReason, + inference_id: inferenceId, }); } - async setUninstalled(productName: ProductName) { - const objectId = getObjectIdFromProductName(productName); + async setUninstalled(productName: ProductName, inferenceId: string | undefined) { + const objectId = getObjectIdFromProductName(productName, inferenceId); try { await this.soClient.update(typeName, objectId, { installation_status: 'uninstalled', last_installation_failure_reason: '', + inference_id: inferenceId, }); } catch (e) { if (!SavedObjectsErrorHelpers.isNotFoundError(e)) { @@ -85,5 +141,7 @@ export class ProductDocInstallClient { } } -const getObjectIdFromProductName = (productName: ProductName) => - `kb-product-doc-${productName}-status`.toLowerCase(); +const getObjectIdFromProductName = (productName: ProductName, inferenceId: string | undefined) => { + const inferenceIdPart = !isImpliedDefaultElserInferenceId(inferenceId) ? `-${inferenceId}` : ''; + return `kb-product-doc-${productName}${inferenceIdPart}-status`.toLowerCase(); +}; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/service.mock.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/service.mock.ts index c2a0adbac9f29..91497404f8bf4 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/service.mock.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_install_status/service.mock.ts @@ -16,6 +16,7 @@ const createInstallClientMock = (): InstallClientMock => { setInstallationSuccessful: jest.fn(), setInstallationFailed: jest.fn(), setUninstalled: jest.fn(), + getPreviouslyInstalledInferenceIds: jest.fn().mockResolvedValue([]), } as unknown as InstallClientMock; }; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.test.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.test.ts index 0be913ee6dd71..eed325f327259 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.test.ts @@ -20,6 +20,7 @@ import { getTaskStatus, waitUntilTaskCompleted, } from '../../tasks'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; const scheduleInstallAllTaskMock = scheduleInstallAllTask as jest.MockedFn< typeof scheduleInstallAllTask @@ -35,6 +36,7 @@ const waitUntilTaskCompletedMock = waitUntilTaskCompleted as jest.MockedFn< >; const getTaskStatusMock = getTaskStatus as jest.MockedFn; +const DEFAULT_INFERENCE_ID = defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL; describe('DocumentationManager', () => { let logger: MockedLogger; let taskManager: ReturnType; @@ -85,19 +87,20 @@ describe('DocumentationManager', () => { }); it('calls `scheduleInstallAllTask`', async () => { - await docManager.install({}); + await docManager.install({ inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleInstallAllTaskMock).toHaveBeenCalledTimes(1); expect(scheduleInstallAllTaskMock).toHaveBeenCalledWith({ taskManager, logger, + inferenceId: DEFAULT_INFERENCE_ID, }); expect(waitUntilTaskCompletedMock).not.toHaveBeenCalled(); }); it('calls waitUntilTaskCompleted if wait=true', async () => { - await docManager.install({ wait: true }); + await docManager.install({ wait: true, inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleInstallAllTaskMock).toHaveBeenCalledTimes(1); expect(waitUntilTaskCompletedMock).toHaveBeenCalledTimes(1); @@ -108,7 +111,7 @@ describe('DocumentationManager', () => { kibana: { status: 'installed' }, } as Awaited>); - await docManager.install({ wait: true }); + await docManager.install({ wait: true, inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleInstallAllTaskMock).not.toHaveBeenCalled(); expect(waitUntilTaskCompletedMock).not.toHaveBeenCalled(); @@ -120,7 +123,12 @@ describe('DocumentationManager', () => { const auditLog = auditService.withoutRequest; auditService.asScoped = jest.fn(() => auditLog); - await docManager.install({ force: false, wait: false, request }); + await docManager.install({ + force: false, + wait: false, + request, + inferenceId: DEFAULT_INFERENCE_ID, + }); expect(auditLog.log).toHaveBeenCalledTimes(1); expect(auditLog.log).toHaveBeenCalledWith({ @@ -140,7 +148,7 @@ describe('DocumentationManager', () => { ); await expect( - docManager.install({ force: false, wait: false }) + docManager.install({ force: false, wait: false, inferenceId: DEFAULT_INFERENCE_ID }) ).rejects.toThrowErrorMatchingInlineSnapshot( `"Elastic documentation requires an enterprise license"` ); @@ -157,19 +165,20 @@ describe('DocumentationManager', () => { }); it('calls `scheduleEnsureUpToDateTask`', async () => { - await docManager.update({}); + await docManager.update({ inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleEnsureUpToDateTaskMock).toHaveBeenCalledTimes(1); expect(scheduleEnsureUpToDateTaskMock).toHaveBeenCalledWith({ taskManager, logger, + inferenceId: DEFAULT_INFERENCE_ID, }); expect(waitUntilTaskCompletedMock).not.toHaveBeenCalled(); }); it('calls waitUntilTaskCompleted if wait=true', async () => { - await docManager.update({ wait: true }); + await docManager.update({ wait: true, inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleEnsureUpToDateTaskMock).toHaveBeenCalledTimes(1); expect(waitUntilTaskCompletedMock).toHaveBeenCalledTimes(1); @@ -181,7 +190,7 @@ describe('DocumentationManager', () => { const auditLog = auditService.withoutRequest; auditService.asScoped = jest.fn(() => auditLog); - await docManager.update({ wait: false, request }); + await docManager.update({ wait: false, request, inferenceId: DEFAULT_INFERENCE_ID }); expect(auditLog.log).toHaveBeenCalledTimes(1); expect(auditLog.log).toHaveBeenCalledWith({ @@ -206,19 +215,20 @@ describe('DocumentationManager', () => { }); it('calls `scheduleUninstallAllTask`', async () => { - await docManager.uninstall({}); + await docManager.uninstall({ inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleUninstallAllTaskMock).toHaveBeenCalledTimes(1); expect(scheduleUninstallAllTaskMock).toHaveBeenCalledWith({ taskManager, logger, + inferenceId: DEFAULT_INFERENCE_ID, }); expect(waitUntilTaskCompletedMock).not.toHaveBeenCalled(); }); it('calls waitUntilTaskCompleted if wait=true', async () => { - await docManager.uninstall({ wait: true }); + await docManager.uninstall({ wait: true, inferenceId: DEFAULT_INFERENCE_ID }); expect(scheduleUninstallAllTaskMock).toHaveBeenCalledTimes(1); expect(waitUntilTaskCompletedMock).toHaveBeenCalledTimes(1); @@ -230,7 +240,7 @@ describe('DocumentationManager', () => { const auditLog = auditService.withoutRequest; auditService.asScoped = jest.fn(() => auditLog); - await docManager.uninstall({ wait: false, request }); + await docManager.uninstall({ wait: false, request, inferenceId: DEFAULT_INFERENCE_ID }); expect(auditLog.log).toHaveBeenCalledTimes(1); expect(auditLog.log).toHaveBeenCalledWith({ diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.ts index 40dc53e19ceea..8d7baefdd4364 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/doc_manager.ts @@ -9,6 +9,8 @@ import type { Logger } from '@kbn/logging'; import type { CoreAuditService } from '@kbn/core/server'; import { type TaskManagerStartContract, TaskStatus } from '@kbn/task-manager-plugin/server'; import type { LicensingPluginStart } from '@kbn/licensing-plugin/server'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; +import { isImpliedDefaultElserInferenceId } from '@kbn/product-doc-common/src/is_default_inference_endpoint'; import type { InstallationStatus } from '../../../common/install_status'; import type { ProductDocInstallClient } from '../doc_install_status'; import { @@ -27,6 +29,7 @@ import type { DocUninstallOptions, DocUpdateOptions, } from './types'; +import { INSTALL_ALL_TASK_ID_MULTILINGUAL } from '../../tasks/install_all'; const TEN_MIN_IN_MS = 10 * 60 * 1000; @@ -62,10 +65,11 @@ export class DocumentationManager implements DocumentationManagerAPI { this.auditService = auditService; } - async install(options: DocInstallOptions = {}): Promise { + async install(options: DocInstallOptions): Promise { const { request, force = false, wait = false } = options; + const inferenceId = options.inferenceId ?? defaultInferenceEndpoints.ELSER; - const { status } = await this.getStatus(); + const { status } = await this.getStatus({ inferenceId }); if (!force && status === 'installed') { return; } @@ -78,11 +82,14 @@ export class DocumentationManager implements DocumentationManagerAPI { const taskId = await scheduleInstallAllTask({ taskManager: this.taskManager, logger: this.logger, + inferenceId, }); if (request) { this.auditService.asScoped(request).log({ - message: `User is requesting installation of product documentation for AI Assistants. Task ID=[${taskId}]`, + message: + `User is requesting installation of product documentation for AI Assistants. Task ID=[${taskId}]` + + (inferenceId ? `| Inference ID=[${inferenceId}]` : ''), event: { action: 'product_documentation_create', category: ['database'], @@ -101,17 +108,20 @@ export class DocumentationManager implements DocumentationManagerAPI { } } - async update(options: DocUpdateOptions = {}): Promise { - const { request, wait = false } = options; + async update(options: DocUpdateOptions): Promise { + const { request, wait = false, inferenceId } = options; const taskId = await scheduleEnsureUpToDateTask({ taskManager: this.taskManager, logger: this.logger, + inferenceId, }); if (request) { this.auditService.asScoped(request).log({ - message: `User is requesting update of product documentation for AI Assistants. Task ID=[${taskId}]`, + message: + `User is requesting update of product documentation for AI Assistants. Task ID=[${taskId}]` + + (inferenceId ? `| Inference ID=[${inferenceId}]` : ''), event: { action: 'product_documentation_update', category: ['database'], @@ -130,12 +140,13 @@ export class DocumentationManager implements DocumentationManagerAPI { } } - async uninstall(options: DocUninstallOptions = {}): Promise { - const { request, wait = false } = options; + async uninstall(options: DocUninstallOptions): Promise { + const { request, wait = false, inferenceId } = options; const taskId = await scheduleUninstallAllTask({ taskManager: this.taskManager, logger: this.logger, + inferenceId, }); if (request) { @@ -159,10 +170,16 @@ export class DocumentationManager implements DocumentationManagerAPI { } } - async getStatus(): Promise { + /** + * @param inferenceId - The inference ID to get the status for. If not provided, the default ELSER inference ID will be used. + */ + async getStatus({ inferenceId }: { inferenceId: string }): Promise { + const taskId = isImpliedDefaultElserInferenceId(inferenceId) + ? INSTALL_ALL_TASK_ID + : INSTALL_ALL_TASK_ID_MULTILINGUAL; const taskStatus = await getTaskStatus({ taskManager: this.taskManager, - taskId: INSTALL_ALL_TASK_ID, + taskId, }); if (taskStatus !== 'not_scheduled') { const status = convertTaskStatus(taskStatus); @@ -171,9 +188,9 @@ export class DocumentationManager implements DocumentationManagerAPI { } } - const installStatus = await this.docInstallClient.getInstallationStatus(); + const installStatus = await this.docInstallClient.getInstallationStatus({ inferenceId }); const overallStatus = getOverallStatus(Object.values(installStatus).map((v) => v.status)); - return { status: overallStatus }; + return { status: overallStatus, installStatus }; } } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/types.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/types.ts index 5a954a5ffb0fd..82e17f05e64bc 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/types.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/doc_manager/types.ts @@ -6,7 +6,7 @@ */ import type { KibanaRequest } from '@kbn/core/server'; -import type { InstallationStatus } from '../../../common/install_status'; +import type { InstallationStatus, ProductInstallState } from '../../../common/install_status'; /** * APIs to manage the product documentation. @@ -30,8 +30,9 @@ export interface DocumentationManagerAPI { uninstall(options?: DocUninstallOptions): Promise; /** * Returns the overall installation status of the documentation. + * @param inferenceId - The inference ID to get the status for. */ - getStatus(): Promise; + getStatus({ inferenceId }: { inferenceId: string }): Promise; } /** @@ -39,6 +40,7 @@ export interface DocumentationManagerAPI { */ export interface DocGetStatusResponse { status: InstallationStatus; + installStatus?: Record; } /** @@ -61,6 +63,10 @@ export interface DocInstallOptions { * Defaults to `false` */ wait?: boolean; + /** + * If provided, the docs will be installed with the model indicated by Inference ID + */ + inferenceId: string; } /** @@ -78,6 +84,10 @@ export interface DocUninstallOptions { * Defaults to `false` */ wait?: boolean; + /** + * If provided, the docs will be uninstalled with the model indicated by Inference ID + */ + inferenceId: string; } /** @@ -95,4 +105,8 @@ export interface DocUpdateOptions { * Defaults to `false` */ wait?: boolean; + /** + * If provided, the docs will be updated with the model indicated by Inference ID + */ + inferenceId: string; } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.test.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.test.ts index ff8e2fd91dea3..7903f6bb613a7 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/package_installer.test.ts @@ -16,7 +16,7 @@ import { fetchArtifactVersionsMock, ensureDefaultElserDeployedMock, } from './package_installer.test.mocks'; - +import { cloneDeep } from 'lodash'; import { getArtifactName, getProductDocIndexName, @@ -29,6 +29,7 @@ import { installClientMock } from '../doc_install_status/service.mock'; import type { ProductInstallState } from '../../../common/install_status'; import { PackageInstaller } from './package_installer'; import { defaultInferenceEndpoints } from '@kbn/inference-common'; +import { InferenceTaskType } from '@elastic/elasticsearch/lib/api/types'; const artifactsFolder = '/lost'; const artifactRepositoryUrl = 'https://repository.com'; @@ -86,7 +87,15 @@ describe('PackageInstaller', () => { }; openZipArchiveMock.mockResolvedValue(zipArchive); - const mappings = Symbol('mappings'); + const mappings = { + properties: { + semantic: { + inference_id: '.elser', + type: 'semantic_text', + model_settings: {}, + }, + }, + }; loadMappingFileMock.mockResolvedValue(mappings); await packageInstaller.installPackage({ productName: 'kibana', productVersion: '8.16' }); @@ -114,10 +123,11 @@ describe('PackageInstaller', () => { expect(loadManifestFileMock).toHaveBeenCalledWith(zipArchive); expect(createIndexMock).toHaveBeenCalledTimes(1); + const modifiedMappings = cloneDeep(mappings); + modifiedMappings.properties.semantic.inference_id = defaultInferenceEndpoints.ELSER; expect(createIndexMock).toHaveBeenCalledWith({ - elserInferenceId: defaultInferenceEndpoints.ELSER, indexName, - mappings, + mappings: modifiedMappings, manifestVersion: TEST_FORMAT_VERSION, esClient, log: logger, @@ -125,16 +135,20 @@ describe('PackageInstaller', () => { expect(populateIndexMock).toHaveBeenCalledTimes(1); expect(populateIndexMock).toHaveBeenCalledWith({ - elserInferenceId: defaultInferenceEndpoints.ELSER, indexName, archive: zipArchive, manifestVersion: TEST_FORMAT_VERSION, + inferenceId: defaultInferenceEndpoints.ELSER, esClient, log: logger, }); expect(productDocClient.setInstallationSuccessful).toHaveBeenCalledTimes(1); - expect(productDocClient.setInstallationSuccessful).toHaveBeenCalledWith('kibana', indexName); + expect(productDocClient.setInstallationSuccessful).toHaveBeenCalledWith( + 'kibana', + indexName, + defaultInferenceEndpoints.ELSER + ); expect(zipArchive.close).toHaveBeenCalledTimes(1); @@ -166,7 +180,16 @@ describe('PackageInstaller', () => { }); await expect( - packageInstaller.installPackage({ productName: 'kibana', productVersion: '8.16' }) + packageInstaller.installPackage({ + productName: 'kibana', + productVersion: '8.16', + customInference: { + inference_id: defaultInferenceEndpoints.ELSER, + task_type: 'text_embedding' as InferenceTaskType, + service: 'elser', + service_settings: {}, + }, + }) ).rejects.toThrowError(); expect(productDocClient.setInstallationSuccessful).not.toHaveBeenCalled(); @@ -181,7 +204,8 @@ describe('PackageInstaller', () => { expect(productDocClient.setInstallationFailed).toHaveBeenCalledTimes(1); expect(productDocClient.setInstallationFailed).toHaveBeenCalledWith( 'kibana', - 'something bad' + 'something bad', + defaultInferenceEndpoints.ELSER ); }); }); @@ -195,7 +219,7 @@ describe('PackageInstaller', () => { elasticsearch: ['8.15'], }); - await packageInstaller.installAll({}); + await packageInstaller.installAll({ inferenceId: defaultInferenceEndpoints.ELSER }); expect(packageInstaller.installPackage).toHaveBeenCalledTimes(2); @@ -226,7 +250,7 @@ describe('PackageInstaller', () => { jest.spyOn(packageInstaller, 'installPackage'); - await packageInstaller.ensureUpToDate({}); + await packageInstaller.ensureUpToDate({ inferenceId: defaultInferenceEndpoints.ELSER }); expect(packageInstaller.installPackage).toHaveBeenCalledTimes(1); expect(packageInstaller.installPackage).toHaveBeenCalledWith({ @@ -249,7 +273,7 @@ describe('PackageInstaller', () => { ); expect(productDocClient.setUninstalled).toHaveBeenCalledTimes(1); - expect(productDocClient.setUninstalled).toHaveBeenCalledWith('kibana'); + expect(productDocClient.setUninstalled).toHaveBeenCalledWith('kibana', undefined); }); }); 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 2aed2063bd95a..e9f23f30626c1 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 @@ -14,6 +14,9 @@ import { type ProductName, } from '@kbn/product-doc-common'; 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 type { ProductDocInstallClient } from '../doc_install_status'; import { downloadToDisk, @@ -22,6 +25,7 @@ import { loadManifestFile, ensureDefaultElserDeployed, type ZipArchive, + ensureInferenceDeployed, } from './utils'; import { majorMinor, latestVersion } from './utils/semver'; import { @@ -30,6 +34,7 @@ import { createIndex, populateIndex, } from './steps'; +import { overrideInferenceSettings } from './steps/create_index'; interface PackageInstallerOpts { artifactsFolder: string; @@ -48,7 +53,7 @@ export class PackageInstaller { private readonly productDocClient: ProductDocInstallClient; private readonly artifactRepositoryUrl: string; private readonly currentVersion: string; - private readonly elserInferenceId?: string; + private readonly elserInferenceId: string; constructor({ artifactsFolder, @@ -68,16 +73,29 @@ export class PackageInstaller { this.elserInferenceId = elserInferenceId || defaultInferenceEndpoints.ELSER; } + private async getInferenceInfo(inferenceId?: string) { + if (!inferenceId) { + return; + } + const inferenceEndpoints = await this.esClient.inference.get({ + inference_id: inferenceId, + }); + return Array.isArray(inferenceEndpoints.endpoints) && inferenceEndpoints.endpoints.length > 0 + ? inferenceEndpoints.endpoints[0] + : undefined; + } /** * Make sure that the currently installed doc packages are up to date. * Will not upgrade products that are not already installed */ - async ensureUpToDate({}: {}) { + async ensureUpToDate(params: { inferenceId: string }) { + const { inferenceId } = params; + const inferenceInfo = await this.getInferenceInfo(inferenceId); const [repositoryVersions, installStatuses] = await Promise.all([ fetchArtifactVersions({ artifactRepositoryUrl: this.artifactRepositoryUrl, }), - this.productDocClient.getInstallationStatus(), + this.productDocClient.getInstallationStatus({ inferenceId }), ]); const toUpdate: Array<{ @@ -105,15 +123,19 @@ export class PackageInstaller { await this.installPackage({ productName, productVersion, + customInference: inferenceInfo, }); } } - async installAll({}: {}) { + async installAll(params: { inferenceId?: string } = {}) { + const { inferenceId } = params; const repositoryVersions = await fetchArtifactVersions({ artifactRepositoryUrl: this.artifactRepositoryUrl, }); const allProducts = Object.values(DocumentationProduct) as ProductName[]; + const inferenceInfo = await this.getInferenceInfo(inferenceId); + for (const productName of allProducts) { const availableVersions = repositoryVersions[productName]; if (!availableVersions || !availableVersions.length) { @@ -125,6 +147,7 @@ export class PackageInstaller { await this.installPackage({ productName, productVersion: selectedVersion, + customInference: inferenceInfo, }); } } @@ -132,32 +155,53 @@ export class PackageInstaller { async installPackage({ productName, productVersion, + customInference, }: { productName: ProductName; productVersion: string; + customInference?: InferenceInferenceEndpointInfo; }) { + const inferenceId = customInference?.inference_id ?? this.elserInferenceId; + this.log.info( - `Starting installing documentation for product [${productName}] and version [${productVersion}]` + `Starting installing documentation for product [${productName}] and version [${productVersion}] with inference ID [${inferenceId}]` ); productVersion = majorMinor(productVersion); - await this.uninstallPackage({ productName }); + await this.uninstallPackage({ productName, inferenceId }); let zipArchive: ZipArchive | undefined; try { await this.productDocClient.setInstallationStarted({ productName, productVersion, + inferenceId, }); - if (this.elserInferenceId === defaultInferenceEndpoints.ELSER) { + if (customInference && customInference?.inference_id !== this.elserInferenceId) { + 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'.` + ); + } + await ensureInferenceDeployed({ + client: this.esClient, + inferenceId, + }); + } + + if (!customInference || customInference?.inference_id === this.elserInferenceId) { await ensureDefaultElserDeployed({ client: this.esClient, }); } - const artifactFileName = getArtifactName({ productName, productVersion }); + const artifactFileName = getArtifactName({ + productName, + productVersion, + inferenceId: customInference?.inference_id ?? this.elserInferenceId, + }); const artifactUrl = `${this.artifactRepositoryUrl}/${artifactFileName}`; const artifactPath = `${this.artifactsFolder}/${artifactFileName}`; @@ -165,7 +209,6 @@ export class PackageInstaller { await downloadToDisk(artifactUrl, artifactPath); zipArchive = await openZipArchive(artifactPath); - validateArtifactArchive(zipArchive); const [manifest, mappings] = await Promise.all([ @@ -174,14 +217,16 @@ export class PackageInstaller { ]); const manifestVersion = manifest.formatVersion; - const indexName = getProductDocIndexName(productName); + const indexName = getProductDocIndexName(productName, customInference?.inference_id); + + const modifiedMappings = cloneDeep(mappings); + overrideInferenceSettings(modifiedMappings, inferenceId!); await createIndex({ indexName, - mappings, + mappings: modifiedMappings, // Mappings will be overridden by the inference ID and inference type manifestVersion, esClient: this.esClient, - elserInferenceId: this.elserInferenceId, log: this.log, }); @@ -191,27 +236,45 @@ export class PackageInstaller { archive: zipArchive, esClient: this.esClient, log: this.log, - elserInferenceId: this.elserInferenceId, + inferenceId, }); - await this.productDocClient.setInstallationSuccessful(productName, indexName); + await this.productDocClient.setInstallationSuccessful(productName, indexName, inferenceId); this.log.info( `Documentation installation successful for product [${productName}] and version [${productVersion}]` ); } catch (e) { + let message = e.message; + if (message.includes('End of central directory record signature not found.')) { + message = i18n.translate('aiInfra.productDocBase.packageInstaller.noArtifactAvailable', { + values: { + productName, + productVersion, + inferenceId, + }, + defaultMessage: + 'No documentation artifact available for product [{productName}]/[{productVersion}] for Inference ID [{inferenceId}]. Please select a different model or contact your administrator.', + }); + } this.log.error( - `Error during documentation installation of product [${productName}]/[${productVersion}] : ${e.message}` + `Error during documentation installation of product [${productName}]/[${productVersion}] : ${message}` ); - await this.productDocClient.setInstallationFailed(productName, e.message); + await this.productDocClient.setInstallationFailed(productName, message, inferenceId); throw e; } finally { zipArchive?.close(); } } - async uninstallPackage({ productName }: { productName: ProductName }) { - const indexName = getProductDocIndexName(productName); + async uninstallPackage({ + productName, + inferenceId, + }: { + productName: ProductName; + inferenceId?: string; + }) { + const indexName = getProductDocIndexName(productName, inferenceId); await this.esClient.indices.delete( { index: indexName, @@ -219,13 +282,14 @@ export class PackageInstaller { { ignore: [404] } ); - await this.productDocClient.setUninstalled(productName); + await this.productDocClient.setUninstalled(productName, inferenceId); } - async uninstallAll() { + async uninstallAll(params: { inferenceId?: string } = {}) { + const { inferenceId } = params; const allProducts = Object.values(DocumentationProduct); for (const productName of allProducts) { - await this.uninstallPackage({ productName }); + await this.uninstallPackage({ productName, inferenceId }); } } } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.test.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.test.ts index 691aeffa40a5b..23593eb1fe335 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.test.ts @@ -11,7 +11,6 @@ import type { ElasticsearchClient } from '@kbn/core/server'; import { elasticsearchServiceMock } from '@kbn/core/server/mocks'; import { LATEST_MANIFEST_FORMAT_VERSION } from '@kbn/product-doc-common'; import { createIndex } from './create_index'; -import { internalElserInferenceId } from '../../../../common/consts'; const LEGACY_SEMANTIC_TEXT_VERSION = '1.0.0'; @@ -76,7 +75,7 @@ describe('createIndex', () => { }); }); - it('rewrites the inference_id attribute of semantic_text fields in the mapping', async () => { + it('does not override the inference_id attribute of semantic_text fields in the mapping', async () => { const mappings: MappingTypeMapping = { properties: { semantic: { @@ -99,17 +98,18 @@ describe('createIndex', () => { expect(esClient.indices.create).toHaveBeenCalledWith( expect.objectContaining({ + index: '.some-index', mappings: { properties: { - semantic: { - type: 'semantic_text', - inference_id: internalElserInferenceId, - }, - bool: { - type: 'boolean', - }, + bool: { type: 'boolean' }, + semantic: { inference_id: '.elser', type: 'semantic_text' }, }, }, + settings: { + auto_expand_replicas: '0-1', + 'index.mapping.semantic_text.use_legacy_format': true, + number_of_shards: 1, + }, }) ); }); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.ts index 1e9423ae6c4fa..57892d73733f7 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/create_index.ts @@ -8,7 +8,7 @@ import type { Logger } from '@kbn/logging'; import type { ElasticsearchClient } from '@kbn/core/server'; import type { MappingTypeMapping, MappingProperty } from '@elastic/elasticsearch/lib/api/types'; -import { internalElserInferenceId } from '../../../../common/consts'; +import { isPopulatedObject } from '@kbn/ml-is-populated-object'; import { isLegacySemanticTextVersion } from '../utils'; export const createIndex = async ({ @@ -17,21 +17,17 @@ export const createIndex = async ({ manifestVersion, mappings, log, - elserInferenceId = internalElserInferenceId, }: { esClient: ElasticsearchClient; indexName: string; manifestVersion: string; mappings: MappingTypeMapping; log: Logger; - elserInferenceId?: string; }) => { log.debug(`Creating index ${indexName}`); const legacySemanticText = isLegacySemanticTextVersion(manifestVersion); - overrideInferenceId(mappings, elserInferenceId); - await esClient.indices.create({ index: indexName, mappings, @@ -43,13 +39,23 @@ export const createIndex = async ({ }); }; -const overrideInferenceId = (mappings: MappingTypeMapping, inferenceId: string) => { +export const overrideInferenceSettings = ( + mappings: MappingTypeMapping, + inferenceId: string, + modelSettingsToOverride?: object +) => { const recursiveOverride = (current: MappingTypeMapping | MappingProperty) => { - if ('type' in current && current.type === 'semantic_text') { + if (isPopulatedObject(current, ['type']) && current.type === 'semantic_text') { current.inference_id = inferenceId; + if (modelSettingsToOverride) { + // @ts-expect-error - model_settings is not typed, but exists for semantic_text field + current.model_settings = modelSettingsToOverride; + } } - if ('properties' in current && current.properties) { - for (const prop of Object.values(current.properties)) { + if (isPopulatedObject(current, ['properties'])) { + for (const prop of Object.values( + current.properties as Record + )) { recursiveOverride(prop); } } diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/populate_index.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/populate_index.ts index deeec80a11464..43a44a8b63cd6 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/populate_index.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/steps/populate_index.ts @@ -19,14 +19,14 @@ export const populateIndex = async ({ manifestVersion, archive, log, - elserInferenceId, + inferenceId = internalElserInferenceId, }: { esClient: ElasticsearchClient; indexName: string; manifestVersion: string; archive: ZipArchive; log: Logger; - elserInferenceId?: string; + inferenceId?: string; }) => { log.debug(`Starting populating index ${indexName}`); @@ -43,7 +43,7 @@ export const populateIndex = async ({ esClient, contentBuffer, legacySemanticText, - elserInferenceId, + inferenceId, }); } @@ -56,12 +56,14 @@ const indexContentFile = async ({ esClient, legacySemanticText, elserInferenceId = internalElserInferenceId, + inferenceId, }: { indexName: string; contentBuffer: Buffer; esClient: ElasticsearchClient; legacySemanticText: boolean; elserInferenceId?: string; + inferenceId?: string; }) => { const fileContent = contentBuffer.toString('utf-8'); const lines = fileContent.split('\n'); @@ -75,7 +77,7 @@ const indexContentFile = async ({ .map((doc) => rewriteInferenceId({ document: doc, - inferenceId: elserInferenceId, + inferenceId: inferenceId ?? elserInferenceId, legacySemanticText, }) ); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/ensure_default_elser_deployed.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/ensure_default_elser_deployed.ts index 14219fb003f2c..3f7714b95f85e 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/ensure_default_elser_deployed.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/ensure_default_elser_deployed.ts @@ -8,12 +8,26 @@ import type { ElasticsearchClient } from '@kbn/core/server'; import { defaultInferenceEndpoints } from '@kbn/inference-common'; -export const ensureDefaultElserDeployed = async ({ client }: { client: ElasticsearchClient }) => { +export const ensureInferenceDeployed = async ({ + client, + inferenceId, +}: { + client: ElasticsearchClient; + inferenceId?: string; +}) => { + if (!inferenceId) return; await client.inference.inference( { - inference_id: defaultInferenceEndpoints.ELSER, + inference_id: inferenceId, input: 'I just want to call the API to force the model to download and allocate', }, { requestTimeout: 10 * 60 * 1000 } ); }; + +export const ensureDefaultElserDeployed = async ({ client }: { client: ElasticsearchClient }) => { + await ensureInferenceDeployed({ + client, + inferenceId: defaultInferenceEndpoints.ELSER, + }); +}; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/index.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/index.ts index 78aa127e7ef18..a667e7890ee66 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/index.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/package_installer/utils/index.ts @@ -8,5 +8,8 @@ export { downloadToDisk } from './download'; export { openZipArchive, type ZipArchive } from './zip_archive'; export { loadManifestFile, loadMappingFile } from './archive_accessors'; -export { ensureDefaultElserDeployed } from './ensure_default_elser_deployed'; +export { + ensureDefaultElserDeployed, + ensureInferenceDeployed, +} from './ensure_default_elser_deployed'; export { isLegacySemanticTextVersion } from './manifest_versions'; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.test.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.test.ts index 9f7056d20d820..e7aec863ed635 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.test.ts @@ -11,6 +11,7 @@ import { SearchService } from './search_service'; import { getIndicesForProductNames } from './utils'; import { performSearch } from './perform_search'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; jest.mock('./perform_search'); const performSearchMock = performSearch as jest.MockedFn; @@ -32,7 +33,7 @@ describe('SearchService', () => { }); describe('#search', () => { - it('calls `performSearch` with the right parameters', async () => { + it('calls `performSearch` with the right default parameters', async () => { await service.search({ query: 'What is Kibana?', products: ['kibana'], @@ -45,7 +46,62 @@ describe('SearchService', () => { searchQuery: 'What is Kibana?', size: 42, highlights: 3, - index: getIndicesForProductNames(['kibana']), + index: getIndicesForProductNames(['kibana'], undefined), + client: esClient, + }); + }); + it('calls `performSearch` with the right default index name for ELSER inference', async () => { + await service.search({ + query: 'What is Kibana?', + products: ['kibana'], + max: 42, + highlights: 3, + inferenceId: defaultInferenceEndpoints.ELSER, + }); + + expect(performSearchMock).toHaveBeenCalledTimes(1); + expect(performSearchMock).toHaveBeenCalledWith({ + searchQuery: 'What is Kibana?', + size: 42, + highlights: 3, + index: [`.kibana_ai_product_doc_kibana`], + client: esClient, + }); + }); + + it('calls `performSearch` with the right default index name for ELSER EIS inference ID', async () => { + await service.search({ + query: 'What is Kibana?', + products: ['kibana'], + max: 42, + highlights: 3, + inferenceId: defaultInferenceEndpoints.ELSER_IN_EIS_INFERENCE_ID, + }); + + expect(performSearchMock).toHaveBeenCalledTimes(1); + expect(performSearchMock).toHaveBeenCalledWith({ + searchQuery: 'What is Kibana?', + size: 42, + highlights: 3, + index: [`.kibana_ai_product_doc_kibana`], + client: esClient, + }); + }); + it('reroutes `performSearch` to multilingual index when inference ID is E5 small', async () => { + await service.search({ + query: 'What is Kibana?', + products: ['kibana'], + max: 42, + highlights: 3, + inferenceId: defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL, + }); + + expect(performSearchMock).toHaveBeenCalledTimes(1); + expect(performSearchMock).toHaveBeenCalledWith({ + searchQuery: 'What is Kibana?', + size: 42, + highlights: 3, + index: [`.kibana_ai_product_doc_kibana-${defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL}`], client: esClient, }); }); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.ts index 5b354c0d95471..ad2ce22ccd1d4 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/search_service.ts @@ -21,13 +21,14 @@ export class SearchService { } async search(options: DocSearchOptions): Promise { - const { query, max = 3, highlights = 3, products } = options; - this.log.debug(`performing search - query=[${query}]`); + const { query, max = 3, highlights = 3, products, inferenceId } = options; + const index = getIndicesForProductNames(products, inferenceId); + this.log.debug(`performing search - query=[${query}] at index=[${index}] `); const results = await performSearch({ searchQuery: query, size: max, highlights, - index: getIndicesForProductNames(products), + index, client: this.esClient, }); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/types.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/types.ts index 910201391543e..e960b58e177c9 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/types.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/types.ts @@ -19,6 +19,8 @@ export interface DocSearchOptions { highlights?: number; /** optional list of products to filter search */ products?: ProductName[]; + /** optional inference ID to filter search */ + inferenceId?: string; } /** diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.test.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.test.ts index 0293d086d4f13..8625ce9f521f9 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.test.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.test.ts @@ -7,11 +7,12 @@ import { productDocIndexPattern, getProductDocIndexName } from '@kbn/product-doc-common'; import { getIndicesForProductNames } from './get_indices_for_product_names'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; describe('getIndicesForProductNames', () => { it('returns the index pattern when product names are not specified', () => { - expect(getIndicesForProductNames(undefined)).toEqual(productDocIndexPattern); - expect(getIndicesForProductNames([])).toEqual(productDocIndexPattern); + expect(getIndicesForProductNames(undefined, undefined)).toEqual(productDocIndexPattern); + expect(getIndicesForProductNames([], undefined)).toEqual(productDocIndexPattern); }); it('returns individual index names when product names are specified', () => { expect(getIndicesForProductNames(['kibana', 'elasticsearch'])).toEqual([ @@ -19,4 +20,31 @@ describe('getIndicesForProductNames', () => { getProductDocIndexName('elasticsearch'), ]); }); + it('returns individual index names when ELSER EIS is specified', () => { + expect(getIndicesForProductNames(['kibana', 'elasticsearch'], '.elser-v2-elastic')).toEqual([ + getProductDocIndexName('kibana'), + getProductDocIndexName('elasticsearch'), + ]); + }); + it('returns individual index names when ELSER is specified', () => { + expect( + getIndicesForProductNames(['kibana', 'elasticsearch'], defaultInferenceEndpoints.ELSER) + ).toEqual([getProductDocIndexName('kibana'), getProductDocIndexName('elasticsearch')]); + }); + + it('returns the index pattern when inferenceId is specified', () => { + expect( + getIndicesForProductNames( + ['kibana', 'elasticsearch'], + defaultInferenceEndpoints.MULTILINGUAL_E5_SMALL + ) + ).toEqual([ + '.kibana_ai_product_doc_kibana-.multilingual-e5-small-elasticsearch', + '.kibana_ai_product_doc_elasticsearch-.multilingual-e5-small-elasticsearch', + ]); + expect(getIndicesForProductNames(['kibana', 'elasticsearch'], '.anyInferenceId')).toEqual([ + '.kibana_ai_product_doc_kibana-.anyInferenceId', + '.kibana_ai_product_doc_elasticsearch-.anyInferenceId', + ]); + }); }); diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.ts index e97ed9cea3611..98ffd1bb39d6d 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/services/search/utils/get_indices_for_product_names.ts @@ -12,10 +12,11 @@ import { } from '@kbn/product-doc-common'; export const getIndicesForProductNames = ( - productNames: ProductName[] | undefined + productNames: ProductName[] | undefined, + inferenceId?: string ): string | string[] => { if (!productNames || !productNames.length) { return productDocIndexPattern; } - return productNames.map(getProductDocIndexName); + return productNames.map((productName) => getProductDocIndexName(productName, inferenceId)); }; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/ensure_up_to_date.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/ensure_up_to_date.ts index d971561914ff1..ba16d2403f290 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/ensure_up_to_date.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/ensure_up_to_date.ts @@ -10,11 +10,14 @@ import type { TaskManagerSetupContract, TaskManagerStartContract, } from '@kbn/task-manager-plugin/server'; +import { isImpliedDefaultElserInferenceId } from '@kbn/product-doc-common/src/is_default_inference_endpoint'; import type { InternalServices } from '../types'; import { isTaskCurrentlyRunningError } from './utils'; export const ENSURE_DOC_UP_TO_DATE_TASK_TYPE = 'ProductDocBase:EnsureUpToDate'; export const ENSURE_DOC_UP_TO_DATE_TASK_ID = 'ProductDocBase:EnsureUpToDate'; +export const ENSURE_DOC_UP_TO_DATE_TASK_ID_MULTILINGUAL = + 'ProductDocBase:EnsureUpToDateMultilingual'; export const registerEnsureUpToDateTaskDefinition = ({ getServices, @@ -31,8 +34,9 @@ export const registerEnsureUpToDateTaskDefinition = ({ createTaskRunner: (context) => { return { async run() { + const inferenceId = context.taskInstance?.params?.inferenceId; const { packageInstaller } = getServices(); - return packageInstaller.ensureUpToDate({}); + return packageInstaller.ensureUpToDate({ inferenceId }); }, }; }, @@ -44,27 +48,32 @@ export const registerEnsureUpToDateTaskDefinition = ({ export const scheduleEnsureUpToDateTask = async ({ taskManager, logger, + inferenceId, }: { taskManager: TaskManagerStartContract; logger: Logger; + inferenceId: string; }) => { + const taskId = isImpliedDefaultElserInferenceId(inferenceId) + ? ENSURE_DOC_UP_TO_DATE_TASK_ID + : ENSURE_DOC_UP_TO_DATE_TASK_ID_MULTILINGUAL; try { await taskManager.ensureScheduled({ - id: ENSURE_DOC_UP_TO_DATE_TASK_ID, + id: taskId, taskType: ENSURE_DOC_UP_TO_DATE_TASK_TYPE, - params: {}, + params: { inferenceId }, state: {}, scope: ['productDoc'], }); - await taskManager.runSoon(ENSURE_DOC_UP_TO_DATE_TASK_ID); + await taskManager.runSoon(taskId); - logger.info(`Task ${ENSURE_DOC_UP_TO_DATE_TASK_ID} scheduled to run soon`); + logger.info(`Task ${taskId} scheduled to run soon`); } catch (e) { if (!isTaskCurrentlyRunningError(e)) { throw e; } } - return ENSURE_DOC_UP_TO_DATE_TASK_ID; + return taskId; }; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/index.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/index.ts index 0b5833055fd8b..c57dbccedb102 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/index.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/index.ts @@ -24,6 +24,14 @@ export const registerTaskDefinitions = ({ }; export { scheduleEnsureUpToDateTask, ENSURE_DOC_UP_TO_DATE_TASK_ID } from './ensure_up_to_date'; -export { scheduleInstallAllTask, INSTALL_ALL_TASK_ID } from './install_all'; -export { scheduleUninstallAllTask, UNINSTALL_ALL_TASK_ID } from './uninstall_all'; +export { + scheduleInstallAllTask, + INSTALL_ALL_TASK_ID, + INSTALL_ALL_TASK_ID_MULTILINGUAL, +} from './install_all'; +export { + scheduleUninstallAllTask, + UNINSTALL_ALL_TASK_ID, + UNINSTALL_ALL_TASK_ID_MULTILINGUAL, +} from './uninstall_all'; export { waitUntilTaskCompleted, getTaskStatus } from './utils'; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/install_all.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/install_all.ts index 0d2cc48fb06bb..8d6fd98335f8e 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/install_all.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/install_all.ts @@ -10,11 +10,13 @@ import type { TaskManagerSetupContract, TaskManagerStartContract, } from '@kbn/task-manager-plugin/server'; +import { isImpliedDefaultElserInferenceId } from '@kbn/product-doc-common/src/is_default_inference_endpoint'; import type { InternalServices } from '../types'; import { isTaskCurrentlyRunningError } from './utils'; export const INSTALL_ALL_TASK_TYPE = 'ProductDocBase:InstallAll'; export const INSTALL_ALL_TASK_ID = 'ProductDocBase:InstallAll'; +export const INSTALL_ALL_TASK_ID_MULTILINGUAL = 'ProductDocBase:InstallAllMultilingual'; export const registerInstallAllTaskDefinition = ({ getServices, @@ -25,14 +27,15 @@ export const registerInstallAllTaskDefinition = ({ }) => { taskManager.registerTaskDefinitions({ [INSTALL_ALL_TASK_TYPE]: { - title: 'Install all product documentation artifacts', + title: `Install all product documentation artifacts ${INSTALL_ALL_TASK_TYPE}`, timeout: '10m', maxAttempts: 3, createTaskRunner: (context) => { + const inferenceId = context.taskInstance?.params?.inferenceId; return { async run() { const { packageInstaller } = getServices(); - return packageInstaller.installAll({}); + return packageInstaller.installAll({ inferenceId }); }, }; }, @@ -44,27 +47,32 @@ export const registerInstallAllTaskDefinition = ({ export const scheduleInstallAllTask = async ({ taskManager, logger, + inferenceId, }: { taskManager: TaskManagerStartContract; logger: Logger; + inferenceId: string; }) => { + const taskId = isImpliedDefaultElserInferenceId(inferenceId) + ? INSTALL_ALL_TASK_ID + : INSTALL_ALL_TASK_ID_MULTILINGUAL; try { await taskManager.ensureScheduled({ - id: INSTALL_ALL_TASK_ID, + id: taskId, taskType: INSTALL_ALL_TASK_TYPE, - params: {}, + params: { inferenceId }, state: {}, scope: ['productDoc'], }); - await taskManager.runSoon(INSTALL_ALL_TASK_ID); + await taskManager.runSoon(taskId); - logger.info(`Task ${INSTALL_ALL_TASK_ID} scheduled to run soon`); + logger.info(`Task ${taskId} scheduled to run soon`); } catch (e) { if (!isTaskCurrentlyRunningError(e)) { throw e; } } - return INSTALL_ALL_TASK_ID; + return taskId; }; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/uninstall_all.ts b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/uninstall_all.ts index 6a88fec205ddd..25cf01fc66b9f 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/uninstall_all.ts +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/server/tasks/uninstall_all.ts @@ -10,11 +10,13 @@ import type { TaskManagerSetupContract, TaskManagerStartContract, } from '@kbn/task-manager-plugin/server'; +import { isImpliedDefaultElserInferenceId } from '@kbn/product-doc-common/src/is_default_inference_endpoint'; import type { InternalServices } from '../types'; import { isTaskCurrentlyRunningError } from './utils'; export const UNINSTALL_ALL_TASK_TYPE = 'ProductDocBase:UninstallAll'; export const UNINSTALL_ALL_TASK_ID = 'ProductDocBase:UninstallAll'; +export const UNINSTALL_ALL_TASK_ID_MULTILINGUAL = 'ProductDocBase:UninstallAllMultilingual'; export const registerUninstallAllTaskDefinition = ({ getServices, @@ -25,14 +27,16 @@ export const registerUninstallAllTaskDefinition = ({ }) => { taskManager.registerTaskDefinitions({ [UNINSTALL_ALL_TASK_TYPE]: { - title: 'Uninstall all product documentation artifacts', + title: `Uninstall all product documentation artifacts ${UNINSTALL_ALL_TASK_TYPE}`, timeout: '10m', maxAttempts: 3, createTaskRunner: (context) => { return { async run() { const { packageInstaller } = getServices(); - return packageInstaller.uninstallAll(); + return packageInstaller.uninstallAll({ + inferenceId: context.taskInstance?.params?.inferenceId, + }); }, }; }, @@ -44,27 +48,35 @@ export const registerUninstallAllTaskDefinition = ({ export const scheduleUninstallAllTask = async ({ taskManager, logger, + inferenceId, }: { taskManager: TaskManagerStartContract; logger: Logger; + inferenceId: string; }) => { + // To avoid conflicts between the default ELSER model and small E5 inference IDs running at the same time, + // we use different task IDs for each inference ID. + const taskId = isImpliedDefaultElserInferenceId(inferenceId) + ? UNINSTALL_ALL_TASK_ID + : UNINSTALL_ALL_TASK_ID_MULTILINGUAL; + try { await taskManager.ensureScheduled({ - id: UNINSTALL_ALL_TASK_ID, + id: taskId, taskType: UNINSTALL_ALL_TASK_TYPE, - params: {}, + params: { inferenceId }, state: {}, scope: ['productDoc'], }); - await taskManager.runSoon(UNINSTALL_ALL_TASK_ID); + await taskManager.runSoon(taskId); - logger.info(`Task ${UNINSTALL_ALL_TASK_ID} scheduled to run soon`); + logger.info(`Task ${taskId} scheduled to run soon`); } catch (e) { if (!isTaskCurrentlyRunningError(e)) { throw e; } } - return UNINSTALL_ALL_TASK_ID; + return taskId; }; diff --git a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/tsconfig.json b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/tsconfig.json index 3b6c5d6cf88a9..51dda22279ea1 100644 --- a/x-pack/platform/plugins/shared/ai_infra/product_doc_base/tsconfig.json +++ b/x-pack/platform/plugins/shared/ai_infra/product_doc_base/tsconfig.json @@ -27,5 +27,7 @@ "@kbn/task-manager-plugin", "@kbn/inference-common", "@kbn/core-security-server", + "@kbn/ml-is-populated-object", + "@kbn/i18n" ] } diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/index.ts index b84234164f8c8..97779ec51aabd 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/index.ts +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/index.ts @@ -60,3 +60,5 @@ export const plugin = async (ctx: PluginInitializerContext { before(async () => { let statusResponse = await kibanaClient.callKibana('get', { pathname: ELASTIC_DOCS_INSTALLATION_STATUS_API_PATH, + query: { + inferenceId, + }, }); if (statusResponse.data.overall === 'installed') { logger.success('Elastic documentation is already installed'); } else { logger.info('Installing Elastic documentation'); - const installResponse = await kibanaClient.callKibana('post', { - pathname: ELASTIC_DOCS_INSTALL_ALL_API_PATH, - }); + const installResponse = await kibanaClient.callKibana( + 'post', + { + pathname: ELASTIC_DOCS_INSTALL_ALL_API_PATH, + }, + { + inferenceId, + } + ); if (!installResponse.data.installed) { logger.error('Could not install Elastic documentation'); @@ -41,6 +52,9 @@ describe('Retrieve documentation function', () => { statusResponse = await kibanaClient.callKibana('get', { pathname: ELASTIC_DOCS_INSTALLATION_STATUS_API_PATH, + query: { + inferenceId, + }, }); if (statusResponse.data.overall !== 'installed') { @@ -74,7 +88,6 @@ describe('Retrieve documentation function', () => { 'Accurately explains what Kibana Lens is and provides doc-based steps for creating a bar chart visualization', `Does not invent unsupported instructions, answers should reference what's found in the Kibana docs`, ]); - expect(result.passed).to.be(true); }); @@ -96,9 +109,15 @@ describe('Retrieve documentation function', () => { after(async () => { // Uninstall all installed documentation logger.info('Uninstalling Elastic documentation'); - const uninstallResponse = await kibanaClient.callKibana('post', { - pathname: ELASTIC_DOCS_UNINSTALL_ALL_API_PATH, - }); + const uninstallResponse = await kibanaClient.callKibana( + 'post', + { + pathname: ELASTIC_DOCS_UNINSTALL_ALL_API_PATH, + }, + { + inferenceId, + } + ); if (uninstallResponse.data.success) { logger.success('Uninstalled Elastic documentation'); 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 86389f817e950..5f628d8c93859 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 @@ -6,6 +6,8 @@ */ import { DocumentationProduct } from '@kbn/product-doc-common'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; +import { getInferenceIdFromWriteIndex } from '@kbn/observability-ai-assistant-plugin/server'; import type { FunctionRegistrationParameters } from '.'; export const RETRIEVE_DOCUMENTATION_NAME = 'retrieve_elastic_doc'; @@ -63,6 +65,12 @@ 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, @@ -70,6 +78,7 @@ export async function registerDocumentationFunction({ connectorId, request: resources.request, functionCalling: simulateFunctionCalling ? 'simulated' : 'auto', + inferenceId, }); return { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts index f8fd88b113bf4..1a075202cf3cd 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/field_maps_configuration.ts @@ -7,7 +7,6 @@ import { FieldMap } from '@kbn/data-stream-adapter'; export const ASSISTANT_ELSER_INFERENCE_ID = 'elastic-security-ai-assistant-elser2'; -export const ELASTICSEARCH_ELSER_INFERENCE_ID = '.elser-2-elasticsearch'; export const knowledgeBaseFieldMap: FieldMap = { // Base fields diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index f31f5b9fca08f..b3294a90708d0 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -38,6 +38,7 @@ import { map } from 'lodash'; import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status'; import type { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers'; import { getMlNodeCount } from '@kbn/ml-plugin/server/lib/node_utils'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { AIAssistantDataClient, AIAssistantDataClientParams } from '..'; import { GetElser } from '../../types'; import { @@ -68,10 +69,7 @@ import { loadSecurityLabs, getSecurityLabsDocsCount, } from '../../lib/langchain/content_loaders/security_labs_loader'; -import { - ASSISTANT_ELSER_INFERENCE_ID, - ELASTICSEARCH_ELSER_INFERENCE_ID, -} from './field_maps_configuration'; +import { ASSISTANT_ELSER_INFERENCE_ID } from './field_maps_configuration'; import { BulkOperationError } from '../../lib/data_stream/documents_data_writer'; import { AUDIT_OUTCOME, KnowledgeBaseAuditAction, knowledgeBaseAuditEvent } from './audit_events'; import { findDocuments } from '../find'; @@ -849,7 +847,7 @@ export const getInferenceEndpointId = async ({ esClient }: { esClient: Elasticse } // Fallback to the default inference endpoint - return ELASTICSEARCH_ELSER_INFERENCE_ID; + return defaultInferenceEndpoints.ELSER; }; /** diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts index c6405e0685cf5..949a4be70e663 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/helpers.ts @@ -13,6 +13,7 @@ import { DeleteByQueryRequest } from '@elastic/elasticsearch/lib/api/types'; import { i18n } from '@kbn/i18n'; import { ProductDocBaseStartContract } from '@kbn/product-doc-base-plugin/server'; import type { Logger } from '@kbn/logging'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { getResourceName } from '.'; import { knowledgeBaseIngestPipeline } from '../ai_assistant_data_clients/knowledge_base/ingest_pipeline'; import { GetElser } from '../types'; @@ -154,12 +155,17 @@ export const ensureProductDocumentationInstalled = async ({ logger: Logger; }) => { try { - const { status } = await productDocManager.getStatus(); + const { status } = await productDocManager.getStatus({ + inferenceId: defaultInferenceEndpoints.ELSER, + }); if (status !== 'installed') { logger.debug(`Installing product documentation for AIAssistantService`); setIsProductDocumentationInProgress(true); try { - await productDocManager.install({ wait: true }); + await productDocManager.install({ + wait: true, + inferenceId: defaultInferenceEndpoints.ELSER, + }); logger.debug(`Successfully installed product documentation for AIAssistantService`); } catch (e) { logger.warn(`Failed to install product documentation for AIAssistantService: ${e.message}`); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 349c22c590ab8..185817d0afe6f 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -27,6 +27,7 @@ import { omit, some } from 'lodash'; import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status'; import { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers'; import { IRuleDataClient } from '@kbn/rule-registry-plugin/server'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { alertSummaryFieldsFieldMap } from '../ai_assistant_data_clients/alert_summary/field_maps_configuration'; import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration'; import { defendInsightsFieldMap } from '../lib/defend_insights/persistence/field_maps_configuration'; @@ -49,7 +50,6 @@ import { assistantAnonymizationFieldsFieldMap } from '../ai_assistant_data_clien import { AIAssistantDataClient } from '../ai_assistant_data_clients'; import { ASSISTANT_ELSER_INFERENCE_ID, - ELASTICSEARCH_ELSER_INFERENCE_ID, knowledgeBaseFieldMap, } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; import { @@ -290,7 +290,7 @@ export class AIAssistantService { type: 'semantic_text', array: false, required: false, - ...(targetInferenceEndpointId !== ELASTICSEARCH_ELSER_INFERENCE_ID + ...(targetInferenceEndpointId !== defaultInferenceEndpoints.ELSER ? { inference_id: targetInferenceEndpointId } : {}), }, @@ -374,7 +374,7 @@ export class AIAssistantService { if (isUsingDedicatedInferenceEndpoint) { this.knowledgeBaseDataStream = await this.rolloverDataStream( ASSISTANT_ELSER_INFERENCE_ID, - ELASTICSEARCH_ELSER_INFERENCE_ID + defaultInferenceEndpoints.ELSER ); } else { // We need to make sure that the data stream is created with the correct mappings @@ -542,7 +542,9 @@ export class AIAssistantService { } public async getProductDocumentationStatus(): Promise { - const status = await this.productDocManager?.getStatus(); + const status = await this.productDocManager?.getStatus({ + inferenceId: defaultInferenceEndpoints.ELSER, + }); if (!status) { return 'uninstalled'; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/config_schema.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/config_schema.ts index 4c4273f00b95b..1539a4a6f15d0 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/config_schema.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/config_schema.ts @@ -5,14 +5,14 @@ * 2.0. */ import { schema } from '@kbn/config-schema'; -import { ELASTICSEARCH_ELSER_INFERENCE_ID } from './ai_assistant_data_clients/knowledge_base/field_maps_configuration'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; export interface ConfigSchema { elserInferenceId: string; responseTimeout: number; } export const configSchema = schema.object({ - elserInferenceId: schema.string({ defaultValue: ELASTICSEARCH_ELSER_INFERENCE_ID }), + elserInferenceId: schema.string({ defaultValue: defaultInferenceEndpoints.ELSER }), responseTimeout: schema.number({ defaultValue: 10 * 60 * 1000, // 10 minutes }), diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts index 6ea4c7f6a9a32..2526d67e5ce20 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.test.ts @@ -25,7 +25,7 @@ import { import { licensingMock } from '@kbn/licensing-plugin/server/mocks'; import { appendAssistantMessageToConversation, langChainExecute } from './helpers'; import { getPrompt } from '../lib/prompt'; -import { ELASTICSEARCH_ELSER_INFERENCE_ID } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; const license = licensingMock.createLicenseMock(); const actionsClient = actionsClientMock.create(); @@ -132,7 +132,7 @@ const mockResponse = { error: jest.fn().mockImplementation((x) => x), }; const mockConfig = { - elserInferenceId: ELASTICSEARCH_ELSER_INFERENCE_ID, + elserInferenceId: defaultInferenceEndpoints.ELSER, responseTimeout: 1000, }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/register_routes.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/register_routes.test.ts index 5e757fe9916f0..86d1f4ddf7046 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/register_routes.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/register_routes.test.ts @@ -52,7 +52,7 @@ import { deleteAttackDiscoverySchedulesRoute } from './attack_discovery/schedule import { findAttackDiscoverySchedulesRoute } from './attack_discovery/schedules/find'; import { disableAttackDiscoverySchedulesRoute } from './attack_discovery/schedules/disable'; import { enableAttackDiscoverySchedulesRoute } from './attack_discovery/schedules/enable'; -import { ELASTICSEARCH_ELSER_INFERENCE_ID } from '../ai_assistant_data_clients/knowledge_base/field_maps_configuration'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; jest.mock('./alert_summary/find_route'); const findAlertSummaryRouteMock = findAlertSummaryRoute as jest.Mock; @@ -136,7 +136,7 @@ const enableAttackDiscoverySchedulesRouteMock = enableAttackDiscoverySchedulesRo describe('registerRoutes', () => { const loggerMock = loggingSystemMock.createLogger(); let server: ReturnType; - const config = { elserInferenceId: ELASTICSEARCH_ELSER_INFERENCE_ID, responseTimeout: 60000 }; + const config = { elserInferenceId: defaultInferenceEndpoints.ELSER, responseTimeout: 60000 }; beforeEach(async () => { jest.clearAllMocks(); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/tsconfig.json b/x-pack/solutions/security/plugins/elastic_assistant/tsconfig.json index 465c049c85f05..e7a0c039dc218 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/tsconfig.json +++ b/x-pack/solutions/security/plugins/elastic_assistant/tsconfig.json @@ -88,7 +88,8 @@ "@kbn/ai-security-labs-content", "@kbn/inference-langchain", "@kbn/security-solution-features", - "@kbn/core-http-server-mocks" + "@kbn/core-http-server-mocks", + "@kbn/inference-common" ], "exclude": [ "target/**/*", diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts index 4f408075260eb..cad4c12b45bf7 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts @@ -22,6 +22,8 @@ import type { } from '@kbn/elastic-assistant-common'; import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +const DEFAULT_INFERENCE_ID = '.elser-2-elasticsearch'; + describe('ProductDocumentationTool', () => { const chain = {} as RetrievalQAChain; const esClient = { @@ -97,6 +99,7 @@ describe('ProductDocumentationTool', () => { connectorId: 'fake-connector', request, functionCalling: 'auto', + inferenceId: DEFAULT_INFERENCE_ID, }); }); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts index 5192508603993..7ea44ff96dd4e 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts @@ -16,6 +16,7 @@ import { import type { ContentReferencesStore } from '@kbn/elastic-assistant-common'; import type { RetrieveDocumentationResultDoc } from '@kbn/llm-tasks-plugin/server'; import type { Require } from '@kbn/elastic-assistant-plugin/server/types'; +import { defaultInferenceEndpoints } from '@kbn/inference-common'; import { APP_UI_ID } from '../../../../common'; export type ProductDocumentationToolParams = Require< @@ -53,6 +54,7 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = { connectorId, request, functionCalling: 'auto', + inferenceId: defaultInferenceEndpoints.ELSER, }); const enrichedDocuments = response.documents.map(enrichDocument(contentReferencesStore)); diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts index bc273ad9584a9..e4f8490b47925 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/complete/functions/retrieve_elastic_doc.spec.ts @@ -15,6 +15,8 @@ import { chatComplete } from '../../utils/conversation'; import type { DeploymentAgnosticFtrProviderContext } from '../../../../../ftr_provider_context'; import { installProductDoc, uninstallProductDoc } from '../../utils/product_doc_base'; +const DEFAULT_INFERENCE_ID = '.elser-2-elasticsearch'; + export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) { const log = getService('log'); const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi'); @@ -90,7 +92,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon connectorId = await observabilityAIAssistantAPIClient.createProxyActionConnector({ port: llmProxy.getPort(), }); - await installProductDoc(supertest); + await installProductDoc(supertest, DEFAULT_INFERENCE_ID); void llmProxy.interceptWithFunctionRequest({ name: 'retrieve_elastic_doc', @@ -115,7 +117,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon }); after(async () => { - await uninstallProductDoc(supertest); + await uninstallProductDoc(supertest, DEFAULT_INFERENCE_ID); llmProxy.close(); await observabilityAIAssistantAPIClient.deleteActionConnector({ actionId: connectorId, diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/product_doc_base.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/product_doc_base.ts index 3fd8913d3d9bb..97cdf016ee785 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/product_doc_base.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/product_doc_base.ts @@ -12,20 +12,26 @@ import { import type SuperTest from 'supertest'; -export async function installProductDoc(supertest: SuperTest.Agent) { +export async function installProductDoc(supertest: SuperTest.Agent, inferenceId: string) { 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') + .send({ + inferenceId, + }) .expect(200); } -export async function uninstallProductDoc(supertest: SuperTest.Agent) { +export async function uninstallProductDoc(supertest: SuperTest.Agent, inferenceId: string) { 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') + .send({ + inferenceId, + }) .expect(200); }