From bb6be5c828445522a8abb1ba94c7033b43acb31e Mon Sep 17 00:00:00 2001 From: Lisa Cawley Date: Mon, 24 Jun 2024 16:18:48 -0700 Subject: [PATCH 1/2] [OAS] Add operationId values for APM APIs (#186840) --- .../plugins/observability_solution/apm/docs/openapi/apm.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/x-pack/plugins/observability_solution/apm/docs/openapi/apm.yaml b/x-pack/plugins/observability_solution/apm/docs/openapi/apm.yaml index 60c8a74d75b88..d37137302fd21 100644 --- a/x-pack/plugins/observability_solution/apm/docs/openapi/apm.yaml +++ b/x-pack/plugins/observability_solution/apm/docs/openapi/apm.yaml @@ -15,6 +15,7 @@ paths: post: summary: Create an APM agent key description: Create a new agent key for APM. + operationId: createAgentKey tags: - APM agent keys requestBody: @@ -56,6 +57,7 @@ paths: get: summary: Search for annotations description: Search for annotations related to a specific service. + operationId: getAnnotation tags: - APM annotations parameters: @@ -110,6 +112,7 @@ paths: post: summary: Create a service annotation description: Create a new annotation for a specific service. + operationId: createAnnotation tags: - APM annotations parameters: From 59036542ab51eb3db6a87fae6f282635f5a3ebdf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=B8ren=20Louv-Jansen?= Date: Tue, 25 Jun 2024 03:05:06 +0200 Subject: [PATCH 2/2] [Obs AI Assistant] Extract `recallFromConnectors` from KB service (#186796) Minor refactor to extract `recallFromConnectors` from the KB service. This change is made to make it easier to unit test parts of it, and reduce the size of the KB service. --- .../server/functions/elasticsearch.ts | 5 +- .../functions/get_dataset_info/index.ts | 6 +- .../service/knowledge_base_service/index.ts | 122 +---------------- .../recall_from_connectors.ts | 128 ++++++++++++++++++ .../observability_ai_assistant/tsconfig.json | 2 + 5 files changed, 138 insertions(+), 125 deletions(-) create mode 100644 x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/recall_from_connectors.ts diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/elasticsearch.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/elasticsearch.ts index 61a8b6adf3ed3..81a9cffc6d033 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/elasticsearch.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/elasticsearch.ts @@ -38,9 +38,8 @@ export function registerElasticsearchFunction({ }, }, async ({ arguments: { method, path, body } }) => { - const response = await ( - await resources.context.core - ).elasticsearch.client.asCurrentUser.transport.request({ + const esClient = (await resources.context.core).elasticsearch.client; + const response = esClient.asCurrentUser.transport.request({ method, path, body, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/get_dataset_info/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/get_dataset_info/index.ts index bac7963cecbdf..4b9128ed549f3 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/get_dataset_info/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/get_dataset_info/index.ts @@ -41,13 +41,13 @@ export function registerGetDatasetInfoFunction({ async ({ arguments: { index }, messages, chat }, signal) => { const coreContext = await resources.context.core; - const esClient = coreContext.elasticsearch.client.asCurrentUser; + const esClient = coreContext.elasticsearch.client; const savedObjectsClient = coreContext.savedObjects.client; let indices: string[] = []; try { - const body = await esClient.indices.resolveIndex({ + const body = await esClient.asCurrentUser.indices.resolveIndex({ name: index === '' ? '*' : index, expand_wildcards: 'open', }); @@ -81,7 +81,7 @@ export function registerGetDatasetInfoFunction({ const relevantFieldNames = await getRelevantFieldNames({ index, messages, - esClient, + esClient: esClient.asCurrentUser, dataViews: await resources.plugins.dataViews.start(), savedObjectsClient, signal, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts index da040e9c2bd06..c59e27d8c82a3 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts @@ -11,15 +11,15 @@ import type { Logger } from '@kbn/logging'; import type { TaskManagerStartContract } from '@kbn/task-manager-plugin/server'; import pLimit from 'p-limit'; import pRetry from 'p-retry'; -import { isEmpty, map, orderBy } from 'lodash'; +import { map, orderBy } from 'lodash'; import { encode } from 'gpt-tokenizer'; import { MlTrainedModelDeploymentNodesStats } from '@elastic/elasticsearch/lib/api/types'; -import { aiAssistantSearchConnectorIndexPattern } from '../../../common'; import { INDEX_QUEUED_DOCUMENTS_TASK_ID, INDEX_QUEUED_DOCUMENTS_TASK_TYPE } from '..'; import { KnowledgeBaseEntry, KnowledgeBaseEntryRole, UserInstruction } from '../../../common/types'; import type { ObservabilityAIAssistantResourceNames } from '../types'; import { getAccessQuery } from '../util/get_access_query'; import { getCategoryQuery } from '../util/get_category_query'; +import { recallFromConnectors } from './recall_from_connectors'; interface Dependencies { esClient: { asInternalUser: ElasticsearchClient }; @@ -348,122 +348,6 @@ export class KnowledgeBaseService { })); } - private async getConnectorIndices( - esClient: { asCurrentUser: ElasticsearchClient }, - uiSettingsClient: IUiSettingsClient - ) { - // improve performance by running this in parallel with the `uiSettingsClient` request - const responsePromise = esClient.asCurrentUser.transport.request({ - method: 'GET', - path: '_connector', - querystring: { - filter_path: 'results.index_name', - }, - }); - - const customSearchConnectorIndex = await uiSettingsClient.get( - aiAssistantSearchConnectorIndexPattern - ); - - if (customSearchConnectorIndex) { - return customSearchConnectorIndex.split(','); - } - - const response = (await responsePromise) as { results?: Array<{ index_name: string }> }; - const connectorIndices = response.results?.map((result) => result.index_name); - - // preserve backwards compatibility with 8.14 (may not be needed in the future) - if (isEmpty(connectorIndices)) { - return ['search-*']; - } - - return connectorIndices; - } - - private async recallFromConnectors({ - queries, - esClient, - uiSettingsClient, - modelId, - }: { - queries: Array<{ text: string; boost?: number }>; - esClient: { asCurrentUser: ElasticsearchClient }; - uiSettingsClient: IUiSettingsClient; - modelId: string; - }): Promise { - const ML_INFERENCE_PREFIX = 'ml.inference.'; - - const connectorIndices = await this.getConnectorIndices(esClient, uiSettingsClient); - - const fieldCaps = await esClient.asCurrentUser.fieldCaps({ - index: connectorIndices, - fields: `${ML_INFERENCE_PREFIX}*`, - allow_no_indices: true, - types: ['sparse_vector'], - filters: '-metadata,-parent', - }); - - const fieldsWithVectors = Object.keys(fieldCaps.fields).map((field) => - field.replace('_expanded.predicted_value', '').replace(ML_INFERENCE_PREFIX, '') - ); - - if (!fieldsWithVectors.length) { - return []; - } - - const esQueries = fieldsWithVectors.flatMap((field) => { - const vectorField = `${ML_INFERENCE_PREFIX}${field}_expanded.predicted_value`; - const modelField = `${ML_INFERENCE_PREFIX}${field}_expanded.model_id`; - - return queries.map(({ text, boost = 1 }) => { - return { - bool: { - should: [ - { - text_expansion: { - [vectorField]: { - model_text: text, - model_id: modelId, - boost, - }, - }, - }, - ], - filter: [ - { - term: { - [modelField]: modelId, - }, - }, - ], - }, - }; - }); - }); - - const response = await esClient.asCurrentUser.search({ - index: connectorIndices, - query: { - bool: { - should: esQueries, - }, - }, - size: 20, - _source: { - exclude: ['_*', 'ml*'], - }, - }); - - const results = response.hits.hits.map((hit) => ({ - text: JSON.stringify(hit._source), - score: hit._score!, - is_correction: false, - id: hit._id, - })); - - return results; - } - recall = async ({ user, queries, @@ -499,7 +383,7 @@ export class KnowledgeBaseService { } throw error; }), - this.recallFromConnectors({ + recallFromConnectors({ esClient, uiSettingsClient, queries, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/recall_from_connectors.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/recall_from_connectors.ts new file mode 100644 index 0000000000000..e042a9422721e --- /dev/null +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/recall_from_connectors.ts @@ -0,0 +1,128 @@ +/* + * 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 { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { IUiSettingsClient } from '@kbn/core-ui-settings-server'; +import { isEmpty } from 'lodash'; +import { RecalledEntry } from '.'; +import { aiAssistantSearchConnectorIndexPattern } from '../../../common'; + +export async function recallFromConnectors({ + queries, + esClient, + uiSettingsClient, + modelId, +}: { + queries: Array<{ text: string; boost?: number }>; + esClient: { asCurrentUser: ElasticsearchClient }; + uiSettingsClient: IUiSettingsClient; + modelId: string; +}): Promise { + const ML_INFERENCE_PREFIX = 'ml.inference.'; + + const connectorIndices = await getConnectorIndices(esClient, uiSettingsClient); + + const fieldCaps = await esClient.asCurrentUser.fieldCaps({ + index: connectorIndices, + fields: `${ML_INFERENCE_PREFIX}*`, + allow_no_indices: true, + types: ['sparse_vector'], + filters: '-metadata,-parent', + }); + + const fieldsWithVectors = Object.keys(fieldCaps.fields).map((field) => + field.replace('_expanded.predicted_value', '').replace(ML_INFERENCE_PREFIX, '') + ); + + if (!fieldsWithVectors.length) { + return []; + } + + const esQueries = fieldsWithVectors.flatMap((field) => { + const vectorField = `${ML_INFERENCE_PREFIX}${field}_expanded.predicted_value`; + const modelField = `${ML_INFERENCE_PREFIX}${field}_expanded.model_id`; + + return queries.map(({ text, boost = 1 }) => { + return { + bool: { + should: [ + { + text_expansion: { + [vectorField]: { + model_text: text, + model_id: modelId, + boost, + }, + }, + }, + ], + filter: [ + { + term: { + [modelField]: modelId, + }, + }, + ], + }, + }; + }); + }); + + const response = await esClient.asCurrentUser.search({ + index: connectorIndices, + query: { + bool: { + should: esQueries, + }, + }, + size: 20, + _source: { + exclude: ['_*', 'ml*'], + }, + }); + + const results = response.hits.hits.map((hit) => ({ + text: JSON.stringify(hit._source), + score: hit._score!, + is_correction: false, + id: hit._id, + })); + + return results; +} + +async function getConnectorIndices( + esClient: { asCurrentUser: ElasticsearchClient }, + uiSettingsClient: IUiSettingsClient +) { + // improve performance by running this in parallel with the `uiSettingsClient` request + const responsePromise = esClient.asCurrentUser.transport.request({ + method: 'GET', + path: '_connector', + querystring: { + filter_path: 'results.index_name', + }, + }); + + const customSearchConnectorIndex = await uiSettingsClient.get( + aiAssistantSearchConnectorIndexPattern + ); + + if (customSearchConnectorIndex) { + return customSearchConnectorIndex.split(','); + } + + const response = (await responsePromise) as { results?: Array<{ index_name: string }> }; + const connectorIndices = response.results?.map((result) => result.index_name); + + // preserve backwards compatibility with 8.14 (may not be needed in the future) + if (isEmpty(connectorIndices)) { + return ['search-*']; + } + + return connectorIndices; +} diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json index aa26acbb6154a..4626a3716759a 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json @@ -38,6 +38,8 @@ "@kbn/features-plugin", "@kbn/cloud-plugin", "@kbn/serverless", + "@kbn/core-elasticsearch-server", + "@kbn/core-ui-settings-server", ], "exclude": ["target/**/*"] }