From e0cbaba2ad304c6c52e10248252b283cb71ee624 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 11 Jun 2025 10:26:19 +0200 Subject: [PATCH 01/17] Add first set of tools --- .../shared/onechat/onechat-common/index.ts | 2 + .../onechat/onechat-common/tools/constants.ts | 29 ++++ .../onechat/onechat-common/tools/index.ts | 1 + .../onechat/onechat-server/src/tools.ts | 2 +- .../plugins/shared/onechat/server/plugin.ts | 3 + .../shared/onechat/server/tools/index.ts | 8 + .../onechat/server/tools/register_tools.ts | 32 ++++ .../tools/retrieval/get_document_by_id.ts | 73 +++++++++ .../tools/retrieval/get_index_mapping.ts | 48 ++++++ .../onechat/server/tools/retrieval/index.ts | 13 ++ .../server/tools/retrieval/list_indices.ts | 62 ++++++++ .../tools/retrieval/rerank_documents.ts | 147 ++++++++++++++++++ .../server/tools/retrieval/search_dsl.ts | 71 +++++++++ .../server/tools/retrieval/search_fulltext.ts | 89 +++++++++++ 14 files changed, 579 insertions(+), 1 deletion(-) create mode 100644 x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/index.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/index.ts b/x-pack/platform/packages/shared/onechat/onechat-common/index.ts index 1873dfe5a5a78..90d494923b123 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/index.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/index.ts @@ -23,6 +23,8 @@ export { createBuiltinToolId, builtinToolProviderId, unknownToolProviderId, + OnechatToolIds, + OnechatToolTags, } from './tools'; export { OnechatErrorCode, diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts new file mode 100644 index 0000000000000..92f3aa1ca26cf --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts @@ -0,0 +1,29 @@ +/* + * 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. + */ + +/** + * Ids of built-in onechat tools + */ +export const OnechatToolIds = { + listIndices: 'list_indices', + getIndexMapping: 'get_index_mapping', + getDocumentById: 'get_document_by_id', + searchDsl: 'search_dsl', + searchFulltext: 'search_fulltext', + rerankDocuments: 'rerank_documents', + generateEsql: 'generate_esql', +}; + +/** + * Common set of tags used for platform tools. + */ +export const OnechatToolTags = { + /** + * Tag associated to tools related to data retrieval + */ + retrieval: 'retrieval', +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts index cd12a06d0dace..1e766df5580f6 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts @@ -23,3 +23,4 @@ export { builtinToolProviderId, unknownToolProviderId, } from './tools'; +export { OnechatToolIds, OnechatToolTags } from './constants'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts b/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts index d78c2d8035008..d067c58a7b9db 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts @@ -22,7 +22,7 @@ import type { ToolEventEmitter } from './events'; /** * Subset of {@link ToolDescriptorMeta} that can be defined during tool registration. */ -export type RegisteredToolMeta = Partial>; +export type RegisteredToolMeta = Partial>; /** * Onechat tool, as registered by built-in tool providers. diff --git a/x-pack/platform/plugins/shared/onechat/server/plugin.ts b/x-pack/platform/plugins/shared/onechat/server/plugin.ts index a72629bb7f286..1232515493bdd 100644 --- a/x-pack/platform/plugins/shared/onechat/server/plugin.ts +++ b/x-pack/platform/plugins/shared/onechat/server/plugin.ts @@ -17,6 +17,7 @@ import type { import { registerRoutes } from './routes'; import { ServiceManager } from './services'; import { registerFeatures } from './features'; +import { registerTools } from './tools'; export class OnechatPlugin implements @@ -47,6 +48,8 @@ export class OnechatPlugin registerFeatures({ features: pluginsSetup.features }); + registerTools({ tools: serviceSetups.tools }); + const router = coreSetup.http.createRouter(); registerRoutes({ router, diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/index.ts b/x-pack/platform/plugins/shared/onechat/server/tools/index.ts new file mode 100644 index 0000000000000..b2ba3025a3166 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/index.ts @@ -0,0 +1,8 @@ +/* + * 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. + */ + +export { registerTools } from './register_tools'; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts new file mode 100644 index 0000000000000..ef84b6fbb4cde --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts @@ -0,0 +1,32 @@ +/* + * 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 type { RegisteredTool } from '@kbn/onechat-server'; +import type { ToolsServiceSetup } from '../services/tools'; +import { + listIndicesTool, + getIndexMappingsTool, + getDocumentByIdTool, + searchFulltextTool, + searchDslTool, + rerankDocumentsTool, +} from './retrieval'; + +export const registerTools = ({ tools: registry }: { tools: ToolsServiceSetup }) => { + const tools: Array> = [ + listIndicesTool(), + getIndexMappingsTool(), + getDocumentByIdTool(), + searchFulltextTool(), + searchDslTool(), + rerankDocumentsTool(), + ]; + + tools.forEach((tool) => { + registry.register(tool); + }); +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts new file mode 100644 index 0000000000000..3b70fd8c151d1 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts @@ -0,0 +1,73 @@ +/* + * 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 { z } from '@kbn/zod'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const getDocumentByIdSchema = z.object({ + id: z.string().describe('ID of the document to retrieve'), + index: z.string().describe('Index to retrieve the document from'), +}); + +export type GetDocumentByIdResult = + | { + id: string; + index: string; + found: true; + _source: unknown; + } + | { + id: string; + index: string; + found: false; + }; + +export const getDocumentByIdTool = (): RegisteredTool< + typeof getDocumentByIdSchema, + GetDocumentByIdResult +> => { + return { + id: OnechatToolIds.getDocumentById, + description: 'Retrieve the full content of a document based on its ID and index name.', + schema: getDocumentByIdSchema, + handler: async ({ id, index }, { esClient }) => { + return getDocumentById({ id, index, esClient: esClient.asCurrentUser }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const getDocumentById = async ({ + id, + index, + esClient, +}: { + id: string; + index: string; + esClient: ElasticsearchClient; +}): Promise => { + const { body: response, statusCode } = await esClient.get( + { + id, + index, + }, + { ignore: [404], meta: true } + ); + if (statusCode === 404) { + return { id, index, found: false }; + } + return { + id, + index, + found: true, + _source: response._source ?? {}, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts new file mode 100644 index 0000000000000..6ea3c3b6a5837 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts @@ -0,0 +1,48 @@ +/* + * 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 type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; +import { z } from '@kbn/zod'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const getIndexMappingsSchema = z.object({ + indices: z.array(z.string()).min(1).describe('List of indices to retrieve mappings for.'), +}); + +export type GetIndexMappingsResult = Record; + +export const getIndexMappingsTool = (): RegisteredTool< + typeof getIndexMappingsSchema, + GetIndexMappingsResult +> => { + return { + id: OnechatToolIds.getIndexMapping, + description: 'Retrieve mappings for the specified index or indices.', + schema: getIndexMappingsSchema, + handler: async ({ indices }, { esClient }) => { + return getIndexMappings({ indices, esClient: esClient.asCurrentUser }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const getIndexMappings = async ({ + indices, + esClient, +}: { + indices: string[]; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.indices.getMapping({ + index: indices, + }); + return response; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts new file mode 100644 index 0000000000000..8435882743394 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts @@ -0,0 +1,13 @@ +/* + * 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. + */ + +export { getDocumentByIdTool } from './get_document_by_id'; +export { getIndexMappingsTool } from './get_index_mapping'; +export { listIndicesTool } from './list_indices'; +export { rerankDocumentsTool } from './rerank_documents'; +export { searchDslTool } from './search_dsl'; +export { searchFulltextTool } from './search_fulltext'; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts new file mode 100644 index 0000000000000..cb32ad9cbba61 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts @@ -0,0 +1,62 @@ +/* + * 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 { z } from '@kbn/zod'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const listIndicesSchema = z.object({ + pattern: z.string().optional().describe('optional pattern to filter indices by. Defaults to *'), +}); + +export interface ListIndexInfo { + index: string; + status: string; + health: string; + uuid: string; + docsCount: number; + primaries: number; + replicas: number; +} + +export const listIndicesTool = (): RegisteredTool => { + return { + id: OnechatToolIds.listIndices, + description: 'List the indices in the Elasticsearch cluster the current user has access to.', + schema: listIndicesSchema, + handler: async ({ pattern = '*' }, { esClient }) => { + return listIndices({ pattern, esClient: esClient.asCurrentUser }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const listIndices = async ({ + pattern = '*', + esClient, +}: { + pattern?: string; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.cat.indices({ + index: pattern, + format: 'json', + }); + + return response.map(({ index, status, health, uuid, 'docs.count': docsCount, pri, rep }) => ({ + index: index ?? 'unknown', + status: status ?? 'unknown', + health: health ?? 'unknown', + uuid: uuid ?? 'unknown', + docsCount: parseInt(docsCount ?? '0', 10), + primaries: parseInt(pri ?? '1', 10), + replicas: parseInt(rep ?? '0', 10), + })); +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts new file mode 100644 index 0000000000000..87cf4a593732e --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts @@ -0,0 +1,147 @@ +/* + * 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 { z } from '@kbn/zod'; +import { BaseMessageLike } from '@langchain/core/messages'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const rerankDocumentsSchema = z.object({ + query: z.string().describe('Text query to rerank snippets by.'), + documents: z + .array( + z.object({ + id: z.string().describe('ID of the document.'), + index: z.string().optional().describe('Index the document is from, if applicable.'), + snippet: z.string().describe('Text snippet to use for reranking.'), + }) + ) + .min(1) + .describe('Documents to rerank'), +}); + +export const rerankDocumentsTool = (): RegisteredTool => { + return { + id: OnechatToolIds.rerankDocuments, + description: 'Score and rerank documents based their relevance against a text query.', + schema: rerankDocumentsSchema, + handler: async ({ query, documents }, { modelProvider }) => { + const { chatModel } = await modelProvider.getDefaultModel(); + + const rerankDocs = documents.map((doc) => ({ + id: doc.id, + content: doc.snippet, + })); + + const analysisModel = chatModel.withStructuredOutput( + z.object({ + ratings: z + .array( + z.object({ + id: z.string().describe('ID of the document'), + grade: z.number().describe('Score of the document, between 0 and 10'), + reason: z + .string() + .optional() + .describe('Optional reason for the rating. Keep it short and concise.'), + }) + ) + .describe('the ratings, one per document using the "{id, grade, reason}" format.'), + }) + ); + + const { ratings } = await analysisModel.invoke( + getAnalysisPrompt({ query, documents: rerankDocs }) + ); + + const ratingMap = ratings.reduce((acc, rating) => { + acc[rating.id] = { + grade: rating.grade, + }; + return acc; + }, {} as Record); + + const documentsWithRatings = documents.map((doc) => ({ + ...doc, + rating: ratingMap[doc.id] ?? 'n/a', + })); + + return documentsWithRatings; + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +interface RerankingDoc { + id: string; + content: unknown; +} + +export const getAnalysisPrompt = ({ + query, + documents, +}: { + query: string; + documents: RerankingDoc[]; +}): BaseMessageLike[] => { + const resultEntry = (document: RerankingDoc): string => { + return ` + ### Document (ID: ${document.id}) + + **Document ID:** "${document.id}" + + **Content:** + \`\`\` + ${JSON.stringify(document.content, null, 2)} + \`\`\` + `; + }; + + return [ + [ + 'system', + ` + ## Current task: Relevance Analysis + + Your task is to evaluate at set of documents in relation to the user’s query, + and assign a relevance rating from 0 to 10 using the following criteria: + - **0:** The document is completely irrelevant. + - **5:** The document is somewhat related and might be useful. + - **8:** The document is very relevant and contains useful information. + - **10:** The document is absolutely crucial for answering the query. + + **Instructions:** + - **Independent Ratings:** Rate each document independently based solely on its relevance to the provided query. + - **Format:** Return your ratings as a JSON object with a "ratings" array, where each element follows the \`"{id}|{grade}"\` format. Example: \`{"ratings": ["0|7", "1|5", "2|10"]}\`. + - **Document IDs:** Use the document IDs provided in this prompt, not any IDs contained in the document content. + - **Optional Comments:** You may include an optional \`"comment"\` field with additional remarks on your ratings. + + ## Input + + ## Input + + You will receive: + 1. The search query from the user. + 2. A list of documents, each with an assigned document ID, that were retrieved in the previous step. + `, + ], + [ + 'human', + ` + ## Input + + **Search Query:**: "${query}" + + ## Documents + + ${documents.map(resultEntry).join('\n')} + `, + ], + ]; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts new file mode 100644 index 0000000000000..6b95443d08d98 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts @@ -0,0 +1,71 @@ +/* + * 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 { z } from '@kbn/zod'; +import type { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const searchDslSchema = z.object({ + query: z.any().describe('Elasticsearch DSL query to run (string or JSON object)'), + index: z.string().describe('Index to search against'), + size: z.number().optional().default(5).describe('Number of documents to return. Defaults to 5.'), +}); + +export interface SearchDslResult { + id: string; + index: string; + source: unknown; +} + +export interface SearchDslResponse { + results: SearchDslResult[]; +} + +export const searchDslTool = (): RegisteredTool => { + return { + id: OnechatToolIds.searchDsl, + description: 'Run a DSL search query on one index and return matching documents.', + schema: searchDslSchema, + handler: async ({ query, index, size }, { esClient }) => { + const parsedQuery = typeof query === 'string' ? JSON.parse(query) : query ?? {}; + return searchDsl({ query: parsedQuery, index, size, esClient: esClient.asCurrentUser }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const searchDsl = async ({ + query, + index, + size, + esClient, +}: { + query: QueryDslQueryContainer; + index: string; + size: number; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.search({ + index, + size, + query, + }); + + const results = response.hits.hits.map((hit) => { + return { + id: hit._id!, + index: hit._index!, + source: hit._source ?? {}, + }; + }); + + return { results }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts new file mode 100644 index 0000000000000..da9f8a7766037 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts @@ -0,0 +1,89 @@ +/* + * 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 { z } from '@kbn/zod'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const fulltextSearchSchema = z.object({ + term: z.string().describe('Term to search for'), + field: z.string().describe('Field to perform fulltext search on'), + index: z.string().describe('Index to search against'), + size: z + .number() + .optional() + .default(10) + .describe('Number of documents to return. Defaults to 10.'), +}); + +export interface SearchFulltextResult { + id: string; + index: string; + highlight: string[]; +} + +export interface SearchFulltextResponse { + results: SearchFulltextResult[]; +} + +export const searchFulltextTool = (): RegisteredTool< + typeof fulltextSearchSchema, + SearchFulltextResponse +> => { + return { + id: OnechatToolIds.searchFulltext, + description: 'Find documents based on a simple fulltext search.', + schema: fulltextSearchSchema, + handler: async ({ term, field, index, size }, { esClient }) => { + return searchFulltext({ term, field, index, size, esClient: esClient.asCurrentUser }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const searchFulltext = async ({ + term, + field, + index, + size, + esClient, +}: { + term: string; + field: string; + index: string; + size: number; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.search({ + index, + size, + query: { + match: { + [field]: term, + }, + }, + highlight: { + number_of_fragments: 5, + fields: { + [field]: {}, + }, + }, + }); + + const results = response.hits.hits.map((hit) => { + return { + id: hit._id!, + index: hit._index!, + highlight: hit.highlight?.[field] || [hit._source[field]], + }; + }); + + return { results }; +}; From 76b86c03e224de1b4459fdbbf4ded40aeb91e1a8 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 11 Jun 2025 11:02:17 +0200 Subject: [PATCH 02/17] refactor reranking tool --- .../tools/retrieval/rerank_documents.ts | 117 +++++++++++++----- 1 file changed, 83 insertions(+), 34 deletions(-) diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts index 87cf4a593732e..2b188590c3ae3 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts @@ -7,6 +7,7 @@ import { z } from '@kbn/zod'; import { BaseMessageLike } from '@langchain/core/messages'; +import type { InferenceChatModel } from '@kbn/inference-langchain'; import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; @@ -24,7 +25,21 @@ const rerankDocumentsSchema = z.object({ .describe('Documents to rerank'), }); -export const rerankDocumentsTool = (): RegisteredTool => { +interface DocumentWithRating { + id: string; + index?: string; + snippet: string; + rating: number; +} + +interface RerankResponse { + documents: DocumentWithRating[]; +} + +export const rerankDocumentsTool = (): RegisteredTool< + typeof rerankDocumentsSchema, + RerankResponse +> => { return { id: OnechatToolIds.rerankDocuments, description: 'Score and rerank documents based their relevance against a text query.', @@ -37,40 +52,25 @@ export const rerankDocumentsTool = (): RegisteredTool((doc) => { + const matchingDoc = documents.find((d) => d.id === doc.id)!; + return { + ...matchingDoc, + rating: doc.rating, + }; }) - ); - - const { ratings } = await analysisModel.invoke( - getAnalysisPrompt({ query, documents: rerankDocs }) - ); - - const ratingMap = ratings.reduce((acc, rating) => { - acc[rating.id] = { - grade: rating.grade, - }; - return acc; - }, {} as Record); - - const documentsWithRatings = documents.map((doc) => ({ - ...doc, - rating: ratingMap[doc.id] ?? 'n/a', - })); + .sort((a, b) => b.rating - a.rating); - return documentsWithRatings; + return { + documents: rerankedDocs, + }; }, meta: { tags: [OnechatToolTags.retrieval], @@ -78,12 +78,61 @@ export const rerankDocumentsTool = (): RegisteredTool => { + const analysisModel = chatModel.withStructuredOutput( + z.object({ + ratings: z + .array( + z.object({ + id: z.string().describe('ID of the document'), + grade: z.number().describe('Score of the document, between 0 and 10'), + reason: z + .string() + .optional() + .describe('Optional reason for the rating. Keep it short and concise.'), + }) + ) + .describe('the ratings, one per document using the "{id, grade, reason}" format.'), + }) + ); + + const { ratings } = await analysisModel.invoke(getRerankingPrompt({ query, documents })); + + const ratingMap = ratings.reduce((acc, rating) => { + acc[rating.id] = { + grade: rating.grade, + }; + return acc; + }, {} as Record); + + const documentsWithRatings = documents.map((doc) => ({ + ...doc, + rating: ratingMap[doc.id] ?? 0, + })); + + return documentsWithRatings; +}; + interface RerankingDoc { id: string; content: unknown; } -export const getAnalysisPrompt = ({ +interface RerankedDoc { + id: string; + content: unknown; + rating: number; +} + +export const getRerankingPrompt = ({ query, documents, }: { From 4a592dc6c97f67f680ce9e999a42f33294b373db Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 11 Jun 2025 11:31:51 +0200 Subject: [PATCH 03/17] tweak system prompt to use the new default tools --- .../agents/conversational/system_prompt.ts | 27 ++++++++++++++++--- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts index b5b27c66605e4..67a65ea7318da 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts @@ -5,16 +5,35 @@ * 2.0. */ -import { BaseMessage, BaseMessageLike } from '@langchain/core/messages'; +import type { BaseMessage, BaseMessageLike } from '@langchain/core/messages'; +import { OnechatToolIds } from '@kbn/onechat-common'; -export const defaultSystemPrompt = - 'You are a helpful chat assistant from the Elasticsearch company.'; +export const defaultSystemPrompt = ` + You are a helpful chat assistant from the Elasticsearch company. + + You have a set of tools at your disposal that can be used to help you answering questions. + In particular, you have tools to access the Elasticsearch cluster on behalf of the user, to search and retrieve documents + they have access to. + + - Never infer an index name from the user's input. Instead, use the ${OnechatToolIds.listIndices} tool + to list the indices in the Elasticsearch cluster the current user has access to. + E.g if the user asks "Can you find documents in the alerts index", Don't assume the index name is "alerts", + and use the ${OnechatToolIds.listIndices} instead to retrieve the list of indices and identify the correct one. + + - Once you have identified the correct index, use the ${OnechatToolIds.getIndexMapping} tool to retrieve its mappings, + as you will need it to call any search tool. + + - When doing fulltext search, prefer the ${OnechatToolIds.searchFulltext} tool over the ${OnechatToolIds.searchDsl} one + when possible. + + - Search tools return highlights of the documents that match the query. The full content of a document can be retrieved + using the ${OnechatToolIds.getDocumentById} tool. + `; const getFullSystemPrompt = (systemPrompt: string) => { return `${systemPrompt} ### Additional info - - You have tools at your disposal that you can use - The current date is: ${new Date().toISOString()} - You can use markdown format to structure your response `; From 83b800959f7bf747d882690163ad033006b0bddd Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Thu, 12 Jun 2025 20:53:49 +0200 Subject: [PATCH 04/17] WIP researcher agent --- .../onechat/onechat-common/tools/constants.ts | 4 + .../onechat/onechat-server/src/tools.ts | 5 + .../server/services/agents/research/graph.ts | 215 ++++++++++++++++++ .../server/services/agents/research/index.ts | 8 + .../services/agents/research/prompts.ts | 138 +++++++++++ .../agents/research/researcher_as_tool.ts | 44 ++++ .../agents/research/run_researcher_agent.ts | 116 ++++++++++ .../server/services/agents/research/utils.ts | 27 +++ .../server/services/runner/run_tool.ts | 3 +- .../onechat/server/tools/register_tools.ts | 6 + .../server/tools/retrieval/execute_esql.ts | 41 ++++ .../onechat/server/tools/retrieval/index.ts | 2 + .../server/tools/retrieval/list_indices.ts | 7 +- .../server/tools/retrieval/nl_to_esql.ts | 152 +++++++++++++ 14 files changed, 766 insertions(+), 2 deletions(-) create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/research/index.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/research/prompts.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/research/run_researcher_agent.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts index 92f3aa1ca26cf..2b2b41afefd6f 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts @@ -9,6 +9,9 @@ * Ids of built-in onechat tools */ export const OnechatToolIds = { + indexExplorer: 'index_explorer', + + /// old listIndices: 'list_indices', getIndexMapping: 'get_index_mapping', getDocumentById: 'get_document_by_id', @@ -16,6 +19,7 @@ export const OnechatToolIds = { searchFulltext: 'search_fulltext', rerankDocuments: 'rerank_documents', generateEsql: 'generate_esql', + executeEsql: 'execute_esql', }; /** diff --git a/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts b/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts index d067c58a7b9db..6a9811080312d 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts @@ -7,6 +7,7 @@ import type { z, ZodObject } from '@kbn/zod'; import type { MaybePromise } from '@kbn/utility-types'; +import type { Logger } from '@kbn/logging'; import type { IScopedClusterClient } from '@kbn/core-elasticsearch-server'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { @@ -137,6 +138,10 @@ export interface ToolHandlerContext { * Event emitter that can be used to emits custom events */ events: ToolEventEmitter; + /** + * Logger scoped to this execution + */ + logger: Logger; } /** diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts new file mode 100644 index 0000000000000..ac3f81d4dd0e3 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts @@ -0,0 +1,215 @@ +/* + * 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 { z } from '@kbn/zod'; +import { StateGraph, Annotation } from '@langchain/langgraph'; +import { BaseMessage } from '@langchain/core/messages'; +import { ToolNode } from '@langchain/langgraph/prebuilt'; +import type { StructuredTool } from '@langchain/core/tools'; +import type { Logger } from '@kbn/core/server'; +import { InferenceChatModel } from '@kbn/inference-langchain'; +import { getReflectionPrompt, getExecutionPrompt, getAnswerPrompt } from './prompts'; +import { extractToolResults } from './utils'; +import { getToolCalls, extractTextContent } from '../conversational/utils/from_langchain_messages'; + +// tool_choice: toolName + +export interface PlannedAction { + knowledgeGap: string; +} + +export interface ExecutedAction { + knowledgeGap: string; + toolName: string; + arguments: any; + response: any; +} + +// +// process queue -> create knowledge entries -> reason +// + +// tools: +// - index explorer +// - fulltext search +// - get_document_by_id +// - ES|QL? + +interface ReflectionResult { + isSufficient: boolean; + knowledgeGaps: string[]; + reasoning?: string; +} + +export const createAgentGraph = async ({ + chatModel, + tools, + systemPrompt, +}: { + chatModel: InferenceChatModel; + tools: StructuredTool[]; + systemPrompt?: string; + logger: Logger; +}) => { + const StateAnnotation = Annotation.Root({ + // inputs + initialQuery: Annotation(), // the search query + cycleBudget: Annotation(), // budget in number of cycles - TODO + // internal state + actionsQueue: Annotation({ + reducer: (state, actions) => { + return actions ?? state; + }, + default: () => [], + }), + processedActions: Annotation({ + reducer: (current, next) => { + return [...current, ...next]; + }, + default: () => [], + }), + lastReflectionResult: Annotation(), + // outputs + generatedAnswer: Annotation(), + }); + + /** + * Initialize the flow by adding a first index explorer call to the action queue. + */ + const initialize = async (state: typeof StateAnnotation.State) => { + const firstAction: PlannedAction = { + knowledgeGap: state.initialQuery, + }; + return { + actionsQueue: [firstAction], + }; + }; + + const processQueueItem = async (state: typeof StateAnnotation.State) => { + const [nextItem, ...queue] = state.actionsQueue; + + console.log('*** processQueueItem - nextItem: ', nextItem); + + const toolNode = new ToolNode(tools); + const executionModel = chatModel.bindTools(tools); + + const response = await executionModel.invoke( + getExecutionPrompt({ + nextAction: nextItem, + executedActions: state.processedActions, + }) + ); + const toolCalls = getToolCalls(response); + + console.log('*** processQueueItem - toolCalls: ', toolCalls); + + const toolMessages = await toolNode.invoke([response]); + const toolResults = extractToolResults(toolMessages); + + const processedActions: ExecutedAction[] = []; + processedActions.push({ + ...nextItem, + toolName: toolCalls[0].toolId.toolId, + arguments: toolCalls[0].args, + response: toolResults[0].result, + }); + + return { + actionsQueue: queue, + processedActions: [processedActions], + }; + }; + + const evaluateQueue = async (state: typeof StateAnnotation.State) => { + const { actionsQueue } = state; + if (actionsQueue.length) { + return 'process_queue_item'; + } + return 'reflection'; + }; + + const reflection = async (state: typeof StateAnnotation.State) => { + const reflectModel = chatModel.withStructuredOutput( + z.object({ + isSufficient: z + .boolean() + .describe('Whether the provided info are sufficient to answer the user question'), + knowledgeGaps: z + .array(z.string()) + .describe('A description of what information is missing or needs clarification'), + reasoning: z + .string() + .optional() + .describe( + 'Optional reasoning on why the provided info are sufficient or not. Can be used as scratch pad for thoughts.' + ), + }) + ); + + const response: ReflectionResult = await reflectModel.invoke( + getReflectionPrompt({ + userQuery: state.initialQuery, + summaries: state.processedActions, + }) + ); + + console.log('*** reflection response: ', response); + + return { + lastReflectionResult: response, + actionsQueue: [ + ...state.actionsQueue, + response.knowledgeGaps.map((gap) => ({ knowledgeGap: gap })), + ], + }; + }; + + const evaluateReflection = async (state: typeof StateAnnotation.State) => { + if (state.lastReflectionResult.isSufficient) { + return 'answer'; + } + return 'process_queue_item'; + }; + + const answer = async (state: typeof StateAnnotation.State) => { + const response = await chatModel.invoke( + getAnswerPrompt({ + userQuery: state.initialQuery, + executedActions: state.processedActions, + }) + ); + + const generatedAnswer = extractTextContent(response); + + console.log('*** answer - response: ', generatedAnswer); + + return { + generatedAnswer, + }; + }; + + // note: the node names are used in the event convertion logic, they should *not* be changed + const graph = new StateGraph(StateAnnotation) + .addNode('initialize', initialize) + .addNode('process_queue_item', processQueueItem) + .addNode('reflection', reflection) + .addNode('answer', answer) + .addEdge('__start__', 'initialize') + .addEdge('initialize', 'process_queue_item') + .addConditionalEdges('process_queue_item', evaluateQueue, { + process_queue_item: 'process_queue_item', + reflection: 'reflection', + }) + .addConditionalEdges('reflection', evaluateReflection, { + process_queue_item: 'process_queue_item', + answer: 'answer', + }) + .addEdge('answer', '__end__') + .compile(); + + return graph; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/index.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/index.ts new file mode 100644 index 0000000000000..45be5d76abcb8 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/index.ts @@ -0,0 +1,8 @@ +/* + * 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. + */ + +export { researcherTool } from './researcher_as_tool'; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/prompts.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/prompts.ts new file mode 100644 index 0000000000000..8cddecb20598f --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/prompts.ts @@ -0,0 +1,138 @@ +/* + * 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 type { BaseMessageLike } from '@langchain/core/messages'; +import type { PlannedAction, ExecutedAction } from './graph'; + +export const getExecutionPrompt = ({ + nextAction, + executedActions, +}: { + nextAction: PlannedAction; + executedActions: ExecutedAction[]; +}): BaseMessageLike[] => { + return [ + [ + 'system', + `You are an expert research assistant from the Elasticsearch company. + + Instructions: + You will be with a goal, and a list of already executed actions. With those information, + please choose which tool to call to reach the goal. + + Requirements: + - You *must* call a tool + - Be attentive, as some tools may require information from the previously executed action. For + example, search tools usually require to target an index, which can be retrieved using the index explorer tool. + - Be careful to not call the same tool twice with the same parameters. Check the action history. + + Action history and current goal will be provided in the next user message. + `, + ], + [ + 'user', + ` + ### Current goal: + + Trying to find information about: "${nextAction.knowledgeGap}" + + ### Action history: + + ${executedActions.map((action) => JSON.stringify(action, undefined, 2)).join('\n')} + `, + ], + ]; +}; + +export const getReflectionPrompt = ({ + userQuery, + summaries, +}: { + userQuery: string; + summaries: any[]; +}): BaseMessageLike[] => { + return [ + [ + 'system', + `You are an expert research assistant from the Elasticsearch company analyzing summaries about "${userQuery}". + + Instructions: + Identify knowledge gaps or areas that need deeper exploration and generate the corresponding follow-up queries. (1 or multiple). + - If provided summaries are sufficient to answer the user's question, don't generate a follow-up query. + - Focus on technical details, implementation specifics, or emerging trends that weren't fully covered. + + Requirements: + - Ensure the follow-up query is self-contained and includes necessary context for web search. + + Output Format: + - Format your response as a JSON object with these exact keys: + - "is_sufficient": true or false + - "knowledge_gap": Describe what information is missing or needs clarification + - "follow_up_queries": Write a specific question to address this gap + + Example 1: if information are sufficient: + \`\`\`json + { + "isSufficient": true, + "knowledgeGaps": [], + } + \`\`\` + `, + ], + [ + 'user', + ` + ### Summaries: + + ${summaries.map((summary) => JSON.stringify(summary, undefined, 2)).join('\n')} + `, + ], + ]; +}; + +export const getAnswerPrompt = ({ + userQuery, + executedActions, +}: { + userQuery: string; + executedActions: ExecutedAction[]; +}): BaseMessageLike[] => { + return [ + [ + 'system', + `Generate a high-quality answer to the user's question based on the provided summaries. + + Instructions: + - The current date is ${new Date().toISOString()}. + - You are the final step of a multi-step research process, don't mention that you are the final step. + - You have access to all the information gathered from the previous steps. + - You have access to the user's question. + - Generate a high-quality answer to the user's question based on the provided summaries and the user's question. + + User Context: + - {research_topic} + + Summaries: + {summaries} + `, + ], + [ + 'user', + ` + ### User question + + "${userQuery}" + + ### Gathered information + + \`\`\`json + ${JSON.stringify(executedActions, undefined, 2)} + \`\`\` + `, + ], + ]; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts new file mode 100644 index 0000000000000..a2c286d1aa6b4 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts @@ -0,0 +1,44 @@ +/* + * 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 { z } from '@kbn/zod'; +import type { RegisteredTool } from '@kbn/onechat-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { runSearchAgent } from './run_researcher_agent'; + +const researcherSchema = z.object({ + instructions: z.string().describe('Instructions for the researcher'), +}); + +export interface ResearcherResponse { + answer: string; +} + +export const researcherTool = (): RegisteredTool => { + return { + id: 'researcher', + description: 'An agentic researcher agent to perform search tasks', + schema: researcherSchema, + handler: async ({ instructions }, { toolProvider, request, modelProvider, runner, logger }) => { + const searchAgentResult = await runSearchAgent( + { + instructions, + toolProvider, + }, + { request, modelProvider, runner, logger } + ); + + return { + answer: searchAgentResult.answer, + }; + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/run_researcher_agent.ts new file mode 100644 index 0000000000000..9cea9d51d1522 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/run_researcher_agent.ts @@ -0,0 +1,116 @@ +/* + * 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 { v4 as uuidv4 } from 'uuid'; +import { Observable, from, filter, shareReplay, firstValueFrom, map, lastValueFrom } from 'rxjs'; +import type { Logger } from '@kbn/logging'; +import { StreamEvent } from '@langchain/core/tracers/log_stream'; +import type { KibanaRequest } from '@kbn/core-http-server'; +import { ChatAgentEvent, isRoundCompleteEvent } from '@kbn/onechat-common'; +import type { ModelProvider, ScopedRunner, ToolProvider } from '@kbn/onechat-server'; +import { + providerToLangchainTools, + toLangchainTool, + conversationLangchainMessages, +} from '../conversational/utils'; +import { createAgentGraph } from './graph'; +import { convertGraphEvents, addRoundCompleteEvent } from '../conversational/convert_graph_events'; + +export interface RunSearchAgentContext { + logger: Logger; + request: KibanaRequest; + modelProvider: ModelProvider; + runner: ScopedRunner; +} + +export interface RunSearchAgentParams { + /** + * The search instructions + */ + instructions: string; + /** + * Top level tool provider to use to retrieve internal tools + */ + toolProvider: ToolProvider; + /** + * Handler to react to the agent's events. + */ + onEvent?: (event: ChatAgentEvent) => void; +} + +export interface RunSearchAgentResponse { + answer: string; +} + +export type RunChatAgentFn = ( + params: RunSearchAgentParams, + context: RunSearchAgentContext +) => Promise; + +const agentGraphName = 'researcher-agent'; + +const noopOnEvent = () => {}; + +/** + * Create the handler function for the default onechat agent. + */ +export const runSearchAgent: RunChatAgentFn = async ( + { instructions, toolProvider, onEvent = noopOnEvent }, + { logger, request, modelProvider } +) => { + const model = await modelProvider.getDefaultModel(); + const langchainTools = await providerToLangchainTools({ request, toolProvider, logger }); + const agentGraph = await createAgentGraph({ + logger, + chatModel: model.chatModel, + tools: langchainTools, + systemPrompt: '', + }); + + const eventStream = agentGraph.streamEvents( + { initialQuery: instructions }, + { + version: 'v2', + runName: agentGraphName, + metadata: { + graphName: agentGraphName, + // runId, + }, + recursionLimit: 10, + callbacks: [], + } + ); + + const events$ = from(eventStream).pipe( + filter(isStreamEvent), + // convertGraphEvents({ graphName: agentGraphName, runName: agentGraphName }), + // addRoundCompleteEvent({ userInput: instructions }), + shareReplay() + ); + + events$.subscribe(onEvent); + + await lastValueFrom(events$); + + // return await extractRound(events$); + return { + answer: 'hello', + }; +}; + +export const extractRound = async (events$: Observable) => { + return await firstValueFrom( + events$.pipe( + filter(isRoundCompleteEvent), + map((event) => event.data.round) + ) + ); +}; + +const isStreamEvent = (input: any): input is StreamEvent => { + return 'event' in input; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts new file mode 100644 index 0000000000000..56d69e495d186 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts @@ -0,0 +1,27 @@ +/* + * 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 { BaseMessage, isToolMessage } from '@langchain/core/messages'; +import { extractTextContent } from '../conversational/utils/from_langchain_messages'; + +interface ToolResult { + toolCallId: string; + result: string; +} + +export const extractToolResults = (messages: BaseMessage[]): ToolResult[] => { + const results: ToolResult[] = []; + for (const message of messages) { + if (isToolMessage(message)) { + results.push({ + toolCallId: message.tool_call_id, + result: extractTextContent(message), + }); + } + } + return results; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/runner/run_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/runner/run_tool.ts index 1fbcc67b1f5ca..77a51ca90055e 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/runner/run_tool.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/runner/run_tool.ts @@ -47,10 +47,11 @@ export const createToolHandlerContext = >({ manager: RunnerManager; }): ToolHandlerContext => { const { onEvent } = toolExecutionParams; - const { request, defaultConnectorId, elasticsearch, modelProviderFactory, toolsService } = + const { request, defaultConnectorId, elasticsearch, modelProviderFactory, toolsService, logger } = manager.deps; return { request, + logger, esClient: elasticsearch.client.asScoped(request), modelProvider: modelProviderFactory({ request, defaultConnectorId }), runner: manager.getRunner(), diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts index ef84b6fbb4cde..e872e194a9940 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts @@ -14,7 +14,10 @@ import { searchFulltextTool, searchDslTool, rerankDocumentsTool, + nlToEsqlTool, + executeEsqlTool, } from './retrieval'; +import { researcherTool } from '../services/agents/research'; export const registerTools = ({ tools: registry }: { tools: ToolsServiceSetup }) => { const tools: Array> = [ @@ -24,6 +27,9 @@ export const registerTools = ({ tools: registry }: { tools: ToolsServiceSetup }) searchFulltextTool(), searchDslTool(), rerankDocumentsTool(), + nlToEsqlTool(), + executeEsqlTool(), + researcherTool(), ]; tools.forEach((tool) => { diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts new file mode 100644 index 0000000000000..489cfe54b5543 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts @@ -0,0 +1,41 @@ +/* + * 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 { z } from '@kbn/zod'; +import type { EsqlEsqlColumnInfo, FieldValue } from '@elastic/elasticsearch/lib/api/types'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; + +const executeEsqlToolSchema = z.object({ + query: z.string().describe('The ES|QL query to execute'), +}); + +export interface ExecuteEsqlResponse { + columns: EsqlEsqlColumnInfo[]; + values: FieldValue[][]; +} + +export const executeEsqlTool = (): RegisteredTool< + typeof executeEsqlToolSchema, + ExecuteEsqlResponse +> => { + return { + id: OnechatToolIds.executeEsql, + description: 'Execute an ES|QL query and return the results.', + schema: executeEsqlToolSchema, + handler: async ({ query }, { esClient }) => { + const response = await esClient.asCurrentUser.esql.query({ query, drop_null_columns: true }); + return { + columns: response.columns, + values: response.values, + }; + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts index 8435882743394..50b11634c6289 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts @@ -11,3 +11,5 @@ export { listIndicesTool } from './list_indices'; export { rerankDocumentsTool } from './rerank_documents'; export { searchDslTool } from './search_dsl'; export { searchFulltextTool } from './search_fulltext'; +export { nlToEsqlTool } from './nl_to_esql'; +export { executeEsqlTool } from './execute_esql'; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts index cb32ad9cbba61..c39f2885e9406 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts @@ -11,7 +11,12 @@ import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; const listIndicesSchema = z.object({ - pattern: z.string().optional().describe('optional pattern to filter indices by. Defaults to *'), + pattern: z + .string() + .optional() + .describe( + '(optional) pattern to filter indices by. Defaults to *. Leave empty to list all indices (recommended)' + ), }); export interface ListIndexInfo { diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts new file mode 100644 index 0000000000000..f5486b54aaf0c --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts @@ -0,0 +1,152 @@ +/* + * 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 { BaseMessageLike } from '@langchain/core/messages'; +import { z } from '@kbn/zod'; +import { filter, toArray, firstValueFrom } from 'rxjs'; +import { isChatCompletionMessageEvent, isChatCompletionEvent } from '@kbn/inference-common'; +import { naturalLanguageToEsql } from '@kbn/inference-plugin/server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { INLINE_ESQL_QUERY_REGEX } from '@kbn/inference-plugin/common/tasks/nl_to_esql/constants'; +import type { RegisteredTool } from '@kbn/onechat-server'; +import { listIndices, ListIndexInfo } from './list_indices'; +import { getIndexMappings } from './get_index_mapping'; + +const nlToEsqlToolSchema = z.object({ + query: z.string().describe('The query to generate an ES|QL query from.'), + context: z + .string() + .optional() + .describe('(optional) Additional context that can be used to generate the ES|QL query'), +}); + +export interface NlToEsqlResponse { + answer: string; + queries: string[]; +} + +export const nlToEsqlTool = (): RegisteredTool => { + return { + id: OnechatToolIds.generateEsql, + description: 'Generate an ES|QL query from a natural language query.', + schema: nlToEsqlToolSchema, + handler: async ({ query, context }, { esClient, modelProvider }) => { + const { chatModel, inferenceClient } = await modelProvider.getDefaultModel(); + const indexInfo = await listIndices({ esClient: esClient.asCurrentUser, pattern: '*' }); + + const indexSelectionModel = chatModel.withStructuredOutput( + z.object({ + indices: z + .array( + z.object({ + name: z.string().describe('name of the index'), + reason: z + .string() + .optional() + .describe('(optional) reason why the index is relevant'), + }) + ) + .describe('the index, or indices, that should be used to generate the ES|QL query'), + }) + ); + + const { indices: selectedIndices } = await indexSelectionModel.invoke( + getIndexSelectionPrompt({ query, context, indices: indexInfo }) + ); + + const indexMappings = await getIndexMappings({ + indices: selectedIndices.map((index) => { + return index.name; + }), + esClient: esClient.asCurrentUser, + }); + + // console.log('selectedIndices: ', selectedIndices); + // console.log('indexMappings: ', indexMappings); + + const esqlEvents$ = naturalLanguageToEsql({ + // @ts-expect-error using a scoped inference client + connectorId: undefined, + client: inferenceClient, + logger: { debug: () => undefined }, + input: ` + Generate an ES|QL query for the following: + + *User query*: ${query}, + *Additional context*"${context} + + *Indices: ${selectedIndices}* + + *Index mappings: ${indexMappings}* + `, + }); + + const messages = await firstValueFrom( + esqlEvents$.pipe( + filter(isChatCompletionEvent), + filter(isChatCompletionMessageEvent), + toArray() + ) + ); + + const fullContent = messages.map((message) => message.content).join('\n'); + const esqlQueries = extractEsqlQueries(fullContent); + + return { + answer: fullContent, + queries: esqlQueries, + }; + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const getIndexSelectionPrompt = ({ + query, + context, + indices, +}: { + query: string; + context?: string; + indices: ListIndexInfo[]; +}): BaseMessageLike[] => { + const resultEntry = (document: ListIndexInfo): string => { + return ` + - **${document.index}** + `; + }; + + return [ + [ + 'system', + ` + ## Current task: Index identification + + Given a user query and additional context, identify the relevant indices that should be searched + for documents that contain relevant information for the user query.`, + ], + [ + 'human', + ` + ## Input + + **Search Query:**: "${query}" + **Additional context:**: "${context ?? 'N/A'}" + + ## List of indices + + ${indices.map(resultEntry).join('\n')} + `, + ], + ]; +}; + +const extractEsqlQueries = (message: string): string[] => { + return Array.from(message.matchAll(INLINE_ESQL_QUERY_REGEX)).map(([match, query]) => query); +}; From 097c2c07efa071fc07440de56f9fc77b9c820707 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Fri, 13 Jun 2025 15:08:11 +0200 Subject: [PATCH 05/17] WIP smart tools --- .../agents/conversational/system_prompt.ts | 9 +- .../server/services/chat/chat_service.ts | 8 +- .../onechat/server/tools/register_tools.ts | 10 +- .../server/tools/retrieval/execute_esql.ts | 21 ++- .../tools/retrieval/get_document_by_id.ts | 4 +- .../tools/retrieval/get_index_mapping.ts | 51 ++++++- .../onechat/server/tools/retrieval/index.ts | 1 + .../server/tools/retrieval/index_explorer.ts | 134 ++++++++++++++++++ .../server/tools/retrieval/nl_to_esql.ts | 125 ++++++---------- .../server/tools/retrieval/search_fulltext.ts | 97 +++++++++++-- .../tools/retrieval/utils/flatten_fields.ts | 49 +++++++ 11 files changed, 390 insertions(+), 119 deletions(-) create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts index 67a65ea7318da..93ac45917c8c5 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts @@ -15,13 +15,8 @@ export const defaultSystemPrompt = ` In particular, you have tools to access the Elasticsearch cluster on behalf of the user, to search and retrieve documents they have access to. - - Never infer an index name from the user's input. Instead, use the ${OnechatToolIds.listIndices} tool - to list the indices in the Elasticsearch cluster the current user has access to. - E.g if the user asks "Can you find documents in the alerts index", Don't assume the index name is "alerts", - and use the ${OnechatToolIds.listIndices} instead to retrieve the list of indices and identify the correct one. - - - Once you have identified the correct index, use the ${OnechatToolIds.getIndexMapping} tool to retrieve its mappings, - as you will need it to call any search tool. + - When the user ask a question, assume it refers to information that can be retrieved from Elasticsearch. + For example if the user asks "What are my latest alerts", assume you need to search the cluster for documents. - When doing fulltext search, prefer the ${OnechatToolIds.searchFulltext} tool over the ${OnechatToolIds.searchDsl} one when possible. diff --git a/x-pack/platform/plugins/shared/onechat/server/services/chat/chat_service.ts b/x-pack/platform/plugins/shared/onechat/server/services/chat/chat_service.ts index e8441161abb92..c1e02eabfc4c5 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/chat/chat_service.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/chat/chat_service.ts @@ -197,7 +197,8 @@ class ChatServiceImpl implements ChatService { statusCode: 500, }) ); - }) + }), + shareReplay() ); }) ); @@ -321,8 +322,9 @@ const getExecutionEvents$ = ({ ); return () => {}; - }).pipe(shareReplay()); - }) + }); + }), + shareReplay(1) ); }; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts index e872e194a9940..2ec6465bbaf0c 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts @@ -16,20 +16,22 @@ import { rerankDocumentsTool, nlToEsqlTool, executeEsqlTool, + indexExplorerTool, } from './retrieval'; import { researcherTool } from '../services/agents/research'; export const registerTools = ({ tools: registry }: { tools: ToolsServiceSetup }) => { const tools: Array> = [ - listIndicesTool(), - getIndexMappingsTool(), + // listIndicesTool(), + // getIndexMappingsTool(), getDocumentByIdTool(), searchFulltextTool(), - searchDslTool(), - rerankDocumentsTool(), + // searchDslTool(), + // rerankDocumentsTool(), nlToEsqlTool(), executeEsqlTool(), researcherTool(), + indexExplorerTool(), ]; tools.forEach((tool) => { diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts index 489cfe54b5543..671b2be13061e 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts @@ -7,6 +7,7 @@ import { z } from '@kbn/zod'; import type { EsqlEsqlColumnInfo, FieldValue } from '@elastic/elasticsearch/lib/api/types'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; @@ -28,14 +29,24 @@ export const executeEsqlTool = (): RegisteredTool< description: 'Execute an ES|QL query and return the results.', schema: executeEsqlToolSchema, handler: async ({ query }, { esClient }) => { - const response = await esClient.asCurrentUser.esql.query({ query, drop_null_columns: true }); - return { - columns: response.columns, - values: response.values, - }; + return executeEsql({ query, esClient: esClient.asCurrentUser }); }, meta: { tags: [OnechatToolTags.retrieval], }, }; }; + +export const executeEsql = async ({ + query, + esClient, +}: { + query: string; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.esql.query({ query, drop_null_columns: true }); + return { + columns: response.columns, + values: response.values, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts index 3b70fd8c151d1..486167cf7025f 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts @@ -12,7 +12,7 @@ import type { RegisteredTool } from '@kbn/onechat-server'; const getDocumentByIdSchema = z.object({ id: z.string().describe('ID of the document to retrieve'), - index: z.string().describe('Index to retrieve the document from'), + index: z.string().describe('Name of the index to retrieve the document from'), }); export type GetDocumentByIdResult = @@ -34,7 +34,7 @@ export const getDocumentByIdTool = (): RegisteredTool< > => { return { id: OnechatToolIds.getDocumentById, - description: 'Retrieve the full content of a document based on its ID and index name.', + description: 'Retrieve the full content (source) of a document based on its ID and index name.', schema: getDocumentByIdSchema, handler: async ({ id, index }, { esClient }) => { return getDocumentById({ id, index, esClient: esClient.asCurrentUser }); diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts index 6ea3c3b6a5837..d9818755fdb00 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts @@ -15,7 +15,11 @@ const getIndexMappingsSchema = z.object({ indices: z.array(z.string()).min(1).describe('List of indices to retrieve mappings for.'), }); -export type GetIndexMappingsResult = Record; +export interface GetIndexMappingEntry { + mappings: MappingTypeMapping; +} + +export type GetIndexMappingsResult = Record; export const getIndexMappingsTool = (): RegisteredTool< typeof getIndexMappingsSchema, @@ -44,5 +48,48 @@ export const getIndexMappings = async ({ const response = await esClient.indices.getMapping({ index: indices, }); - return response; + + return Object.entries(response).reduce((res, [indexName, mappingRes]) => { + res[indexName] = { mappings: cleanupMapping(mappingRes.mappings) }; + return res; + }, {} as GetIndexMappingsResult); +}; + +/** + * Remove non-relevant mapping information such as `ignore_above` to reduce overall token length of response + * @param mapping + */ +const cleanupMapping = (mapping: MappingTypeMapping): MappingTypeMapping => { + const recurseKeys = ['properties', 'fields']; + const fieldsToKeep = ['type', 'dynamic', '_meta', 'enabled']; + + function recursiveCleanup(obj: Record): Record { + if (Array.isArray(obj)) { + return obj.map((item) => recursiveCleanup(item)); + } else if (obj !== null && typeof obj === 'object') { + const cleaned: Record = {}; + + for (const key of Object.keys(obj)) { + if (recurseKeys.includes(key)) { + const value = obj[key]; + if (value !== null && typeof value === 'object') { + // For properties/fields: preserve all keys inside + const subCleaned: Record = {}; + for (const fieldName of Object.keys(value)) { + subCleaned[fieldName] = recursiveCleanup(value[fieldName]); + } + cleaned[key] = subCleaned; + } + } else if (fieldsToKeep.includes(key)) { + cleaned[key] = recursiveCleanup(obj[key]); + } + } + + return cleaned; + } else { + return obj; + } + } + + return recursiveCleanup(mapping); }; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts index 50b11634c6289..b231af7dda245 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts @@ -13,3 +13,4 @@ export { searchDslTool } from './search_dsl'; export { searchFulltextTool } from './search_fulltext'; export { nlToEsqlTool } from './nl_to_esql'; export { executeEsqlTool } from './execute_esql'; +export { indexExplorerTool } from './index_explorer'; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts new file mode 100644 index 0000000000000..02dc42797b69f --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts @@ -0,0 +1,134 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; + +import { z } from '@kbn/zod'; +import { BaseMessageLike } from '@langchain/core/messages'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import type { RegisteredTool, ScopedModel } from '@kbn/onechat-server'; +import { listIndices } from './list_indices'; +import { getIndexMappings } from './get_index_mapping'; + +const indexExplorerSchema = z.object({ + query: z.string().describe('A natural language query to infer which indices to use.'), + limit: z + .number() + .optional() + .describe('(optional) Limit the max number of indices to return. Defaults to 1.'), + indexPattern: z + .string() + .optional() + .describe('(optional) Index pattern to filter indices by. Defaults to *.'), +}); + +export interface RelevantIndex { + indexName: string; + mappings: MappingTypeMapping; + reason: string; +} + +export interface IndexExplorerResponse { + indices: RelevantIndex[]; +} + +export const indexExplorerTool = (): RegisteredTool< + typeof indexExplorerSchema, + IndexExplorerResponse +> => { + return { + id: OnechatToolIds.indexExplorer, + description: `List relevant indices and corresponding mappings based on a natural language query. + + The 'indexPattern' parameter can be used to filter indices by a specific pattern, e.g. 'foo*'. + This should *only* be used if you know what you're doing (e.g. if the user explicitly specified a pattern). + Otherwise, leave it empty to list all indices. + + *Example:* + User: "Show me my latest alerts" + You: call tool 'indexExplorer' with { query: 'indices containing alerts' } + Tool result: [{ indexName: '.alerts', mappings: {...} }] + `, + schema: indexExplorerSchema, + handler: async ({ query, indexPattern = '*', limit = 1 }, { esClient, modelProvider }) => { + const model = await modelProvider.getDefaultModel(); + return indexExplorer({ query, indexPattern, limit, esClient: esClient.asCurrentUser, model }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; + +export const indexExplorer = async ({ + query, + indexPattern = '*', + limit = 1, + esClient, + model, +}: { + query: string; + indexPattern?: string; + limit?: number; + esClient: ElasticsearchClient; + model: ScopedModel; +}): Promise => { + const { chatModel } = model; + + const allIndices = await listIndices({ + pattern: indexPattern, + esClient, + }); + + const indexSelectorModel = chatModel.withStructuredOutput( + z.object({ + indices: z.array( + z.object({ + indexName: z.string().describe('name of the index'), + reason: z.string().describe('brief explanation of why this index could be relevant'), + }) + ), + }) + ); + + const indexSelectorPrompt: BaseMessageLike[] = [ + [ + 'user', + `You are an AI assistant for the Elasticsearch company. + based on a natural language query from the user, your task is to select up to ${limit} most relevant indices from a list of indices. + + *The query is:* ${query} + + *List of indices:* + ${allIndices.map((index) => `- ${index.index}`).join('\n')} + + Based on those information, please return most relevant indices with your reasoning. + Remember, you should select at maximum ${limit} indices. + `, + ], + ]; + + const { indices: selectedIndices } = await indexSelectorModel.invoke(indexSelectorPrompt); + + const mappings = await getIndexMappings({ + indices: selectedIndices.map((index) => index.indexName), + esClient, + }); + + const relevantIndices: RelevantIndex[] = selectedIndices.map( + ({ indexName, reason }) => { + return { + indexName, + reason, + mappings: mappings[indexName].mappings, + }; + } + ); + + return { indices: relevantIndices }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts index f5486b54aaf0c..ae296d3cede2e 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts @@ -5,6 +5,7 @@ * 2.0. */ +import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; import { BaseMessageLike } from '@langchain/core/messages'; import { z } from '@kbn/zod'; import { filter, toArray, firstValueFrom } from 'rxjs'; @@ -15,13 +16,20 @@ import { INLINE_ESQL_QUERY_REGEX } from '@kbn/inference-plugin/common/tasks/nl_t import type { RegisteredTool } from '@kbn/onechat-server'; import { listIndices, ListIndexInfo } from './list_indices'; import { getIndexMappings } from './get_index_mapping'; +import { indexExplorer } from './index_explorer'; const nlToEsqlToolSchema = z.object({ query: z.string().describe('The query to generate an ES|QL query from.'), + index: z + .string() + .optional() + .describe( + '(optional) Index to search against. If not provided, will use the index explorer to find the best index to use.' + ), context: z .string() .optional() - .describe('(optional) Additional context that can be used to generate the ES|QL query'), + .describe('(optional) Additional context that could be useful to generate the ES|QL query'), }); export interface NlToEsqlResponse { @@ -34,54 +42,49 @@ export const nlToEsqlTool = (): RegisteredTool { - const { chatModel, inferenceClient } = await modelProvider.getDefaultModel(); - const indexInfo = await listIndices({ esClient: esClient.asCurrentUser, pattern: '*' }); - - const indexSelectionModel = chatModel.withStructuredOutput( - z.object({ - indices: z - .array( - z.object({ - name: z.string().describe('name of the index'), - reason: z - .string() - .optional() - .describe('(optional) reason why the index is relevant'), - }) - ) - .describe('the index, or indices, that should be used to generate the ES|QL query'), - }) - ); - - const { indices: selectedIndices } = await indexSelectionModel.invoke( - getIndexSelectionPrompt({ query, context, indices: indexInfo }) - ); - - const indexMappings = await getIndexMappings({ - indices: selectedIndices.map((index) => { - return index.name; - }), - esClient: esClient.asCurrentUser, - }); - - // console.log('selectedIndices: ', selectedIndices); - // console.log('indexMappings: ', indexMappings); + handler: async ({ query, index, context }, { esClient, modelProvider }) => { + const model = await modelProvider.getDefaultModel(); + + let selectedIndex = index; + let mappings: MappingTypeMapping; + + if (index) { + selectedIndex = index; + const indexMappings = await getIndexMappings({ + indices: [index], + esClient: esClient.asCurrentUser, + }); + mappings = indexMappings[index].mappings; + } else { + const { + indices: [firstIndex], + } = await indexExplorer({ + query, + esClient: esClient.asCurrentUser, + limit: 1, + model, + }); + selectedIndex = firstIndex.indexName; + mappings = firstIndex.mappings; + } const esqlEvents$ = naturalLanguageToEsql({ // @ts-expect-error using a scoped inference client connectorId: undefined, - client: inferenceClient, + client: model.inferenceClient, logger: { debug: () => undefined }, input: ` - Generate an ES|QL query for the following: + Your task is to generate an ES|QL query. - *User query*: ${query}, - *Additional context*"${context} + - User query: "${query}", + - Additional context: "${context ?? 'N/A'} + - Index to use: "${selectedIndex}" + - Mapping of this index: + \`\`\`json + ${JSON.stringify(mappings, undefined, 2)} + \`\`\` - *Indices: ${selectedIndices}* - - *Index mappings: ${indexMappings}* + Given those info, please generate an ES|QL query to address the user request `, }); @@ -107,46 +110,6 @@ export const nlToEsqlTool = (): RegisteredTool { - const resultEntry = (document: ListIndexInfo): string => { - return ` - - **${document.index}** - `; - }; - - return [ - [ - 'system', - ` - ## Current task: Index identification - - Given a user query and additional context, identify the relevant indices that should be searched - for documents that contain relevant information for the user query.`, - ], - [ - 'human', - ` - ## Input - - **Search Query:**: "${query}" - **Additional context:**: "${context ?? 'N/A'}" - - ## List of indices - - ${indices.map(resultEntry).join('\n')} - `, - ], - ]; -}; - const extractEsqlQueries = (message: string): string[] => { return Array.from(message.matchAll(INLINE_ESQL_QUERY_REGEX)).map(([match, query]) => query); }; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts index da9f8a7766037..20fabd0be5cb6 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts @@ -9,11 +9,24 @@ import { z } from '@kbn/zod'; import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; +import { indexExplorer } from './index_explorer'; +import { getIndexMappings } from './get_index_mapping'; +import { flattenFields } from './utils/flatten_fields'; const fulltextSearchSchema = z.object({ term: z.string().describe('Term to search for'), - field: z.string().describe('Field to perform fulltext search on'), - index: z.string().describe('Index to search against'), + index: z + .string() + .optional() + .describe( + '(optional) Index to search against. If not provided, will use index explorer to find the best index to use.' + ), + fields: z + .array(z.string()) + .optional() + .describe( + '(optional) Fields to perform fulltext search on. If not provided, will use all searchable fields.' + ), size: z .number() .optional() @@ -24,23 +37,66 @@ const fulltextSearchSchema = z.object({ export interface SearchFulltextResult { id: string; index: string; - highlight: string[]; + highlights: string[]; } export interface SearchFulltextResponse { results: SearchFulltextResult[]; } +// TODO: rename to relevance search export const searchFulltextTool = (): RegisteredTool< typeof fulltextSearchSchema, SearchFulltextResponse > => { return { id: OnechatToolIds.searchFulltext, - description: 'Find documents based on a simple fulltext search.', + description: `Find relevant documents in an index based on a simple fulltext search. + + - The 'index' parameter can be used to specify which index to search against. If not provided, the tool will use the index explorer to find the best index to use. + - The 'fields' parameter can be used to specify which fields to search on. If not provided, the tool will use all searchable fields. + + It is perfectly fine not to not specify both 'index' and 'fields'. Those should only be used when you already know about the index and fields you want to search on, + e.g if the user explicitly specified them.`, schema: fulltextSearchSchema, - handler: async ({ term, field, index, size }, { esClient }) => { - return searchFulltext({ term, field, index, size, esClient: esClient.asCurrentUser }); + handler: async ({ term, index, fields = [], size }, { esClient, modelProvider }) => { + const model = await modelProvider.getDefaultModel(); + + let selectedIndex = index; + let selectedFields = fields; + + if (!selectedIndex) { + const { indices } = await indexExplorer({ + query: term, + esClient: esClient.asCurrentUser, + model, + }); + if (indices.length === 0) { + return { results: [] }; + } + selectedIndex = indices[0].indexName; + } + + if (!fields.length) { + const mappings = await getIndexMappings({ + indices: [selectedIndex], + esClient: esClient.asCurrentUser, + }); + + const flattenedFields = flattenFields(mappings[selectedIndex]); + + selectedFields = flattenedFields + .filter((field) => field.type === 'text' || field.type === 'semantic_text') + .map((field) => field.path); + } + + return searchFulltext({ + term, + fields: selectedFields, + index: selectedIndex, + size, + esClient: esClient.asCurrentUser, + }); }, meta: { tags: [OnechatToolTags.retrieval], @@ -50,13 +106,13 @@ export const searchFulltextTool = (): RegisteredTool< export const searchFulltext = async ({ term, - field, + fields, index, size, esClient, }: { term: string; - field: string; + fields: string[]; index: string; size: number; esClient: ElasticsearchClient; @@ -64,16 +120,24 @@ export const searchFulltext = async ({ const response = await esClient.search({ index, size, - query: { - match: { - [field]: term, + retriever: { + rrf: { + retrievers: fields.map((field) => { + return { + standard: { + query: { + match: { + [field]: term, + }, + }, + }, + }; + }), }, }, highlight: { number_of_fragments: 5, - fields: { - [field]: {}, - }, + fields: fields.reduce((memo, field) => ({ ...memo, [field]: {} }), {}), }, }); @@ -81,7 +145,10 @@ export const searchFulltext = async ({ return { id: hit._id!, index: hit._index!, - highlight: hit.highlight?.[field] || [hit._source[field]], + highlights: Object.entries(hit.highlight ?? {}).reduce((acc, [field, highlights]) => { + acc.push(...highlights); + return acc; + }, [] as string[]), }; }); diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts new file mode 100644 index 0000000000000..77230f9900225 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts @@ -0,0 +1,49 @@ +/* + * 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 type { MappingTypeMapping, MappingProperty } from '@elastic/elasticsearch/lib/api/types'; + +export type FieldType = Extract['type']; + +export interface MappingField { + path: string; + type: FieldType; +} + +export interface MappingProperties { + [key: string]: { + type?: string; // Leaf field (e.g., "text", "keyword", etc.) + properties?: MappingProperties; // Nested object fields + }; +} + +export const flattenFields = ({ mappings }: { mappings: MappingTypeMapping }): MappingField[] => { + const properties: MappingProperties = mappings.properties ?? {}; + + function extractFields(obj: MappingProperties, prefix = ''): MappingField[] { + let fields: MappingField[] = []; + + for (const [key, value] of Object.entries(obj)) { + const fieldPath = prefix ? `${prefix}.${key}` : key; + + if (value.type) { + // If it's a leaf field, add it + fields.push({ + type: value.type as FieldType, + path: fieldPath, + }); + } else if (value.properties) { + // If it's an object, go deeper + fields = fields.concat(extractFields(value.properties, fieldPath)); + } + } + + return fields; + } + + return extractFields(properties); +}; From 408ba49aa9d5ab9ce7e413c2785b01b58aac461c Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Mon, 16 Jun 2025 08:33:33 +0200 Subject: [PATCH 06/17] wip --- .../onechat/onechat-common/tools/constants.ts | 1 + .../server/tools/retrieval/nl_to_esql.ts | 112 ++++++++++-------- .../server/tools/retrieval/search_dsl.ts | 81 ++++++------- .../server/tools/retrieval/utils/esql.ts | 28 +++++ 4 files changed, 129 insertions(+), 93 deletions(-) create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts index 2b2b41afefd6f..282da2178ff77 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts @@ -10,6 +10,7 @@ */ export const OnechatToolIds = { indexExplorer: 'index_explorer', + // relevanceSearch: 'relevance_search', /// old listIndices: 'list_indices', diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts index ae296d3cede2e..86a0a2a9688eb 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts @@ -6,15 +6,14 @@ */ import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; -import { BaseMessageLike } from '@langchain/core/messages'; import { z } from '@kbn/zod'; import { filter, toArray, firstValueFrom } from 'rxjs'; import { isChatCompletionMessageEvent, isChatCompletionEvent } from '@kbn/inference-common'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { naturalLanguageToEsql } from '@kbn/inference-plugin/server'; import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; import { INLINE_ESQL_QUERY_REGEX } from '@kbn/inference-plugin/common/tasks/nl_to_esql/constants'; -import type { RegisteredTool } from '@kbn/onechat-server'; -import { listIndices, ListIndexInfo } from './list_indices'; +import type { RegisteredTool, ScopedModel } from '@kbn/onechat-server'; import { getIndexMappings } from './get_index_mapping'; import { indexExplorer } from './index_explorer'; @@ -44,36 +43,62 @@ export const nlToEsqlTool = (): RegisteredTool { const model = await modelProvider.getDefaultModel(); + return generateEsql({ + query, + context, + index, + model, + esClient: esClient.asCurrentUser, + }); + }, + meta: { + tags: [OnechatToolTags.retrieval], + }, + }; +}; - let selectedIndex = index; - let mappings: MappingTypeMapping; +export const generateEsql = async ({ + query, + context, + index, + model, + esClient, +}: { + query: string; + context?: string; + index?: string; + model: ScopedModel; + esClient: ElasticsearchClient; +}): Promise => { + let selectedIndex: string | undefined; + let mappings: MappingTypeMapping; - if (index) { - selectedIndex = index; - const indexMappings = await getIndexMappings({ - indices: [index], - esClient: esClient.asCurrentUser, - }); - mappings = indexMappings[index].mappings; - } else { - const { - indices: [firstIndex], - } = await indexExplorer({ - query, - esClient: esClient.asCurrentUser, - limit: 1, - model, - }); - selectedIndex = firstIndex.indexName; - mappings = firstIndex.mappings; - } + if (index) { + selectedIndex = index; + const indexMappings = await getIndexMappings({ + indices: [index], + esClient, + }); + mappings = indexMappings[index].mappings; + } else { + const { + indices: [firstIndex], + } = await indexExplorer({ + query, + esClient, + limit: 1, + model, + }); + selectedIndex = firstIndex.indexName; + mappings = firstIndex.mappings; + } - const esqlEvents$ = naturalLanguageToEsql({ - // @ts-expect-error using a scoped inference client - connectorId: undefined, - client: model.inferenceClient, - logger: { debug: () => undefined }, - input: ` + const esqlEvents$ = naturalLanguageToEsql({ + // @ts-expect-error using a scoped inference client + connectorId: undefined, + client: model.inferenceClient, + logger: { debug: () => undefined }, + input: ` Your task is to generate an ES|QL query. - User query: "${query}", @@ -86,27 +111,18 @@ export const nlToEsqlTool = (): RegisteredTool message.content).join('\n'); - const esqlQueries = extractEsqlQueries(fullContent); + const fullContent = messages.map((message) => message.content).join('\n'); + const esqlQueries = extractEsqlQueries(fullContent); - return { - answer: fullContent, - queries: esqlQueries, - }; - }, - meta: { - tags: [OnechatToolTags.retrieval], - }, + return { + answer: fullContent, + queries: esqlQueries, }; }; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts index 6b95443d08d98..0c0ac83603d3f 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts @@ -6,66 +6,57 @@ */ import { z } from '@kbn/zod'; -import type { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; +import { generateEsql } from './nl_to_esql'; +import { executeEsql, ExecuteEsqlResponse } from './execute_esql'; +// smart search const searchDslSchema = z.object({ - query: z.any().describe('Elasticsearch DSL query to run (string or JSON object)'), - index: z.string().describe('Index to search against'), - size: z.number().optional().default(5).describe('Number of documents to return. Defaults to 5.'), + query: z.string().describe('A natural language query expressing the search request'), + index: z + .string() + .optional() + .describe( + '(optional) Index to search against. If not provided, will use the index explorer to find the best index to use.' + ), + context: z + .string() + .optional() + .describe('(optional) Additional context that could be useful to perform the search'), }); -export interface SearchDslResult { - id: string; - index: string; - source: unknown; -} - -export interface SearchDslResponse { - results: SearchDslResult[]; -} +export type SearchDslResponse = ExecuteEsqlResponse | { success: false; reason: string }; export const searchDslTool = (): RegisteredTool => { return { id: OnechatToolIds.searchDsl, description: 'Run a DSL search query on one index and return matching documents.', schema: searchDslSchema, - handler: async ({ query, index, size }, { esClient }) => { - const parsedQuery = typeof query === 'string' ? JSON.parse(query) : query ?? {}; - return searchDsl({ query: parsedQuery, index, size, esClient: esClient.asCurrentUser }); + handler: async ({ query, index, context }, { esClient, modelProvider }) => { + const model = await modelProvider.getDefaultModel(); + + const generateResponse = await generateEsql({ + query, + context, + index, + model, + esClient: esClient.asCurrentUser, + }); + + if (generateResponse.queries.length < 1) { + return { success: false, reason: 'No query was generated' }; + } + + const executeResponse = await executeEsql({ + query: generateResponse.queries[0], + esClient: esClient.asCurrentUser, + }); + + return executeResponse; }, meta: { tags: [OnechatToolTags.retrieval], }, }; }; - -export const searchDsl = async ({ - query, - index, - size, - esClient, -}: { - query: QueryDslQueryContainer; - index: string; - size: number; - esClient: ElasticsearchClient; -}): Promise => { - const response = await esClient.search({ - index, - size, - query, - }); - - const results = response.hits.hits.map((hit) => { - return { - id: hit._id!, - index: hit._index!, - source: hit._source ?? {}, - }; - }); - - return { results }; -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts new file mode 100644 index 0000000000000..dec1d2259f334 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts @@ -0,0 +1,28 @@ +/* + * 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 type { EsqlEsqlColumnInfo, FieldValue } from '@elastic/elasticsearch/lib/api/types'; + +export interface EsqlResponse { + columns: EsqlEsqlColumnInfo[]; + values: FieldValue[][]; +} + +export const esqlResponseToJson = (esql: EsqlResponse): Array> => { + const results: Array> = []; + + const { columns, values } = esql; + for (const item of values) { + const entry: Record = {}; + for (let i = 0; i < columns.length; i++) { + entry[columns[i].name] = item[i]; + } + results.push(entry); + } + + return results; +}; From bb64fb6a73d4dd7d7b8407304465a313fac66198 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Mon, 16 Jun 2025 14:17:17 +0200 Subject: [PATCH 07/17] move tool implementation to dedicated package --- package.json | 1 + tsconfig.base.json | 2 + .../shared/onechat/onechat-common/index.ts | 4 +- .../onechat/onechat-common/tools/constants.ts | 12 +- .../onechat/onechat-common/tools/index.ts | 2 +- .../onechat/onechat-genai-utils/README.md | 3 + .../onechat/onechat-genai-utils/index.ts | 34 +++ .../onechat-genai-utils/jest.config.js | 12 ++ .../onechat/onechat-genai-utils/kibana.jsonc | 7 + .../onechat/onechat-genai-utils/package.json | 6 + .../tools/generate_esql.ts} | 53 +---- .../onechat-genai-utils/tools/index.ts | 11 + .../tools/index_explorer.ts | 117 +++++++++++ .../onechat-genai-utils/tools/nl_search.ts | 44 ++++ .../tools/relevance_search.ts | 66 ++++++ .../tools/steps/execute_esql.ts | 31 +++ .../tools/steps/get_documents.ts | 50 +++++ .../tools/steps/get_mappings.ts | 37 ++++ .../onechat-genai-utils/tools/steps/index.ts | 20 ++ .../tools/steps/list_indices.ts | 41 ++++ .../tools/steps/perform_match_search.ts | 69 ++++++ .../onechat-genai-utils/tools}/utils/esql.ts | 15 +- .../onechat-genai-utils/tools/utils/index.ts | 9 + .../tools/utils/mappings.ts | 91 ++++++++ .../onechat/onechat-genai-utils/tsconfig.json | 17 ++ .../agents/conversational/system_prompt.ts | 12 +- .../agents/research/researcher_as_tool.ts | 4 +- .../onechat/server/tools/register_tools.ts | 24 +-- .../server/tools/retrieval/execute_esql.ts | 33 +-- .../server/tools/retrieval/generate_esql.ts | 46 ++++ .../tools/retrieval/get_document_by_id.ts | 48 +---- .../tools/retrieval/get_index_mapping.ts | 71 +------ .../onechat/server/tools/retrieval/index.ts | 9 +- .../server/tools/retrieval/index_explorer.ts | 93 +-------- .../server/tools/retrieval/list_indices.ts | 41 +--- .../retrieval/{search_dsl.ts => nl_search.ts} | 31 +-- .../tools/retrieval/relevance_search.ts | 73 +++++++ .../tools/retrieval/rerank_documents.ts | 196 ------------------ .../server/tools/retrieval/search_fulltext.ts | 156 -------------- .../tools/retrieval/utils/flatten_fields.ts | 49 ----- yarn.lock | 4 + 41 files changed, 868 insertions(+), 776 deletions(-) create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/README.md create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/index.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/jest.config.js create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/kibana.jsonc create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json rename x-pack/platform/{plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts => packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts} (59%) create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/nl_search.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/relevance_search.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/execute_esql.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_documents.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_mappings.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/index.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/list_indices.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/perform_match_search.ts rename x-pack/platform/{plugins/shared/onechat/server/tools/retrieval => packages/shared/onechat/onechat-genai-utils/tools}/utils/esql.ts (61%) create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/index.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/mappings.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts rename x-pack/platform/plugins/shared/onechat/server/tools/retrieval/{search_dsl.ts => nl_search.ts} (59%) create mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/relevance_search.ts delete mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts delete mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts delete mode 100644 x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts diff --git a/package.json b/package.json index b9fc429a4dcd7..68ca13193dcfc 100644 --- a/package.json +++ b/package.json @@ -737,6 +737,7 @@ "@kbn/oidc-provider-plugin": "link:x-pack/test/security_api_integration/plugins/oidc_provider", "@kbn/onechat-browser": "link:x-pack/platform/packages/shared/onechat/onechat-browser", "@kbn/onechat-common": "link:x-pack/platform/packages/shared/onechat/onechat-common", + "@kbn/onechat-genai-utils": "link:x-pack/platform/packages/shared/onechat/onechat-genai-utils", "@kbn/onechat-plugin": "link:x-pack/platform/plugins/shared/onechat", "@kbn/onechat-server": "link:x-pack/platform/packages/shared/onechat/onechat-server", "@kbn/open-telemetry-instrumented-plugin": "link:src/platform/test/common/plugins/otel_metrics", diff --git a/tsconfig.base.json b/tsconfig.base.json index bd61c76c04047..4f8018d74b377 100644 --- a/tsconfig.base.json +++ b/tsconfig.base.json @@ -1404,6 +1404,8 @@ "@kbn/onechat-browser/*": ["x-pack/platform/packages/shared/onechat/onechat-browser/*"], "@kbn/onechat-common": ["x-pack/platform/packages/shared/onechat/onechat-common"], "@kbn/onechat-common/*": ["x-pack/platform/packages/shared/onechat/onechat-common/*"], + "@kbn/onechat-genai-utils": ["x-pack/platform/packages/shared/onechat/onechat-genai-utils"], + "@kbn/onechat-genai-utils/*": ["x-pack/platform/packages/shared/onechat/onechat-genai-utils/*"], "@kbn/onechat-plugin": ["x-pack/platform/plugins/shared/onechat"], "@kbn/onechat-plugin/*": ["x-pack/platform/plugins/shared/onechat/*"], "@kbn/onechat-server": ["x-pack/platform/packages/shared/onechat/onechat-server"], diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/index.ts b/x-pack/platform/packages/shared/onechat/onechat-common/index.ts index 90d494923b123..ccee4d8490269 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/index.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/index.ts @@ -23,8 +23,8 @@ export { createBuiltinToolId, builtinToolProviderId, unknownToolProviderId, - OnechatToolIds, - OnechatToolTags, + BuiltinToolIds, + BuiltinTags, } from './tools'; export { OnechatErrorCode, diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts index 282da2178ff77..b6ff77b126d61 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts @@ -8,17 +8,13 @@ /** * Ids of built-in onechat tools */ -export const OnechatToolIds = { +export const BuiltinToolIds = { indexExplorer: 'index_explorer', - // relevanceSearch: 'relevance_search', - - /// old + relevanceSearch: 'relevance_search', + naturalLanguageSearch: 'nl_search', listIndices: 'list_indices', getIndexMapping: 'get_index_mapping', getDocumentById: 'get_document_by_id', - searchDsl: 'search_dsl', - searchFulltext: 'search_fulltext', - rerankDocuments: 'rerank_documents', generateEsql: 'generate_esql', executeEsql: 'execute_esql', }; @@ -26,7 +22,7 @@ export const OnechatToolIds = { /** * Common set of tags used for platform tools. */ -export const OnechatToolTags = { +export const BuiltinTags = { /** * Tag associated to tools related to data retrieval */ diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts index 1e766df5580f6..740dfec99b4da 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/index.ts @@ -23,4 +23,4 @@ export { builtinToolProviderId, unknownToolProviderId, } from './tools'; -export { OnechatToolIds, OnechatToolTags } from './constants'; +export { BuiltinToolIds, BuiltinTags } from './constants'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/README.md b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/README.md new file mode 100644 index 0000000000000..c44aa96866a98 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/README.md @@ -0,0 +1,3 @@ +# @kbn/onechat-genai-utils + +Empty package generated by @kbn/generate diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/index.ts new file mode 100644 index 0000000000000..06c5513de6fb1 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/index.ts @@ -0,0 +1,34 @@ +/* + * 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. + */ + +export { + esqlResponseToJson, + flattenMappings, + cleanupMapping, + type MappingField, +} from './tools/utils'; +export { + getDocumentById, + type GetDocumentByIdResult, + getIndexMappings, + type GetIndexMappingEntry, + type GetIndexMappingsResult, + executeEsql, + type EsqlResponse, + listIndices, + type ListIndexInfo, +} from './tools/steps'; +export { + indexExplorer, + type IndexExplorerResponse, + generateEsql, + type GenerateEsqlResponse, + relevanceSearch, + type RelevanceSearchResponse, + naturalLanguageSearch, + type NaturalLanguageSearchResponse, +} from './tools'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/jest.config.js b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/jest.config.js new file mode 100644 index 0000000000000..1f93c726bd06d --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/jest.config.js @@ -0,0 +1,12 @@ +/* + * 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. + */ + +module.exports = { + preset: '@kbn/test/jest_node', + rootDir: '../../../../../..', + roots: ['/x-pack/platform/packages/shared/onechat/onechat-genai-utils'], +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/kibana.jsonc b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/kibana.jsonc new file mode 100644 index 0000000000000..d65030355f31d --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/kibana.jsonc @@ -0,0 +1,7 @@ +{ + "type": "shared-server", + "id": "@kbn/onechat-genai-utils", + "owner": "@elastic/workchat-eng", + "group": "platform", + "visibility": "shared" +} diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json new file mode 100644 index 0000000000000..220e3faded6f5 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json @@ -0,0 +1,6 @@ +{ + "name": "@kbn/onechat-genai-utils", + "private": true, + "version": "1.0.0", + "license": "Elastic License 2.0 OR AGPL-3.0-only OR SSPL-1.0" +} \ No newline at end of file diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts similarity index 59% rename from x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts rename to x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts index 86a0a2a9688eb..2b95805c81911 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_to_esql.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts @@ -6,57 +6,20 @@ */ import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; -import { z } from '@kbn/zod'; import { filter, toArray, firstValueFrom } from 'rxjs'; -import { isChatCompletionMessageEvent, isChatCompletionEvent } from '@kbn/inference-common'; import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { isChatCompletionMessageEvent, isChatCompletionEvent } from '@kbn/inference-common'; import { naturalLanguageToEsql } from '@kbn/inference-plugin/server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; -import { INLINE_ESQL_QUERY_REGEX } from '@kbn/inference-plugin/common/tasks/nl_to_esql/constants'; -import type { RegisteredTool, ScopedModel } from '@kbn/onechat-server'; -import { getIndexMappings } from './get_index_mapping'; +import type { ScopedModel } from '@kbn/onechat-server'; import { indexExplorer } from './index_explorer'; +import { getIndexMappings } from './steps/get_mappings'; +import { extractEsqlQueries } from './utils/esql'; -const nlToEsqlToolSchema = z.object({ - query: z.string().describe('The query to generate an ES|QL query from.'), - index: z - .string() - .optional() - .describe( - '(optional) Index to search against. If not provided, will use the index explorer to find the best index to use.' - ), - context: z - .string() - .optional() - .describe('(optional) Additional context that could be useful to generate the ES|QL query'), -}); - -export interface NlToEsqlResponse { +export interface GenerateEsqlResponse { answer: string; queries: string[]; } -export const nlToEsqlTool = (): RegisteredTool => { - return { - id: OnechatToolIds.generateEsql, - description: 'Generate an ES|QL query from a natural language query.', - schema: nlToEsqlToolSchema, - handler: async ({ query, index, context }, { esClient, modelProvider }) => { - const model = await modelProvider.getDefaultModel(); - return generateEsql({ - query, - context, - index, - model, - esClient: esClient.asCurrentUser, - }); - }, - meta: { - tags: [OnechatToolTags.retrieval], - }, - }; -}; - export const generateEsql = async ({ query, context, @@ -69,7 +32,7 @@ export const generateEsql = async ({ index?: string; model: ScopedModel; esClient: ElasticsearchClient; -}): Promise => { +}): Promise => { let selectedIndex: string | undefined; let mappings: MappingTypeMapping; @@ -125,7 +88,3 @@ export const generateEsql = async ({ queries: esqlQueries, }; }; - -const extractEsqlQueries = (message: string): string[] => { - return Array.from(message.matchAll(INLINE_ESQL_QUERY_REGEX)).map(([match, query]) => query); -}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index.ts new file mode 100644 index 0000000000000..b5624a13a5981 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index.ts @@ -0,0 +1,11 @@ +/* + * 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. + */ + +export { indexExplorer, type IndexExplorerResponse } from './index_explorer'; +export { generateEsql, type GenerateEsqlResponse } from './generate_esql'; +export { relevanceSearch, type RelevanceSearchResponse } from './relevance_search'; +export { naturalLanguageSearch, type NaturalLanguageSearchResponse } from './nl_search'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts new file mode 100644 index 0000000000000..a1715315b5308 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts @@ -0,0 +1,117 @@ +/* + * 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 type { BaseMessageLike } from '@langchain/core/messages'; +import { z } from '@kbn/zod'; +import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; +import type { ScopedModel } from '@kbn/onechat-server'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { ListIndexInfo, listIndices } from './utils/listings'; +import { getIndexMappings } from './steps/get_mappings'; + +export interface RelevantIndex { + indexName: string; + mappings: MappingTypeMapping; + reason: string; +} + +export interface IndexExplorerResponse { + indices: RelevantIndex[]; +} + +export const indexExplorer = async ({ + query, + indexPattern = '*', + limit = 1, + esClient, + model, +}: { + query: string; + indexPattern?: string; + limit?: number; + esClient: ElasticsearchClient; + model: ScopedModel; +}): Promise => { + const allIndices = await listIndices({ + pattern: indexPattern, + esClient, + }); + + const selectedIndices = await selectIndices({ + indices: allIndices, + query, + model, + limit, + }); + + const mappings = await getIndexMappings({ + indices: selectedIndices.map((index) => index.indexName), + esClient, + }); + + const relevantIndices: RelevantIndex[] = selectedIndices.map( + ({ indexName, reason }) => { + return { + indexName, + reason, + mappings: mappings[indexName].mappings, + }; + } + ); + + return { indices: relevantIndices }; +}; + +export interface SelectedIndex { + indexName: string; + reason: string; +} + +const selectIndices = async ({ + indices, + query, + model, + limit = 1, +}: { + indices: ListIndexInfo[]; + query: string; + model: ScopedModel; + limit?: number; +}): Promise => { + const { chatModel } = model; + const indexSelectorModel = chatModel.withStructuredOutput( + z.object({ + indices: z.array( + z.object({ + indexName: z.string().describe('name of the index'), + reason: z.string().describe('brief explanation of why this index could be relevant'), + }) + ), + }) + ); + + const indexSelectorPrompt: BaseMessageLike[] = [ + [ + 'user', + `You are an AI assistant for the Elasticsearch company. + based on a natural language query from the user, your task is to select up to ${limit} most relevant indices from a list of indices. + + *The query is:* ${query} + + *List of indices:* + ${indices.map((index) => `- ${index.index}`).join('\n')} + + Based on those information, please return most relevant indices with your reasoning. + Remember, you should select at maximum ${limit} indices. + `, + ], + ]; + + const { indices: selectedIndices } = await indexSelectorModel.invoke(indexSelectorPrompt); + + return selectedIndices; +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/nl_search.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/nl_search.ts new file mode 100644 index 0000000000000..5b08677c452e7 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/nl_search.ts @@ -0,0 +1,44 @@ +/* + * 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 type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { ScopedModel } from '@kbn/onechat-server'; +import { executeEsql, EsqlResponse } from './steps/execute_esql'; +import { generateEsql } from './generate_esql'; + +export type NaturalLanguageSearchResponse = EsqlResponse | { success: false; reason: string }; + +export const naturalLanguageSearch = async ({ + query, + context, + index, + model, + esClient, +}: { + query: string; + context?: string; + index?: string; + model: ScopedModel; + esClient: ElasticsearchClient; +}): Promise => { + const generateResponse = await generateEsql({ + query, + context, + index, + model, + esClient, + }); + + if (generateResponse.queries.length < 1) { + return { success: false, reason: 'No query was generated' }; + } + + return await executeEsql({ + query: generateResponse.queries[0], + esClient, + }); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/relevance_search.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/relevance_search.ts new file mode 100644 index 0000000000000..b3441b430472f --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/relevance_search.ts @@ -0,0 +1,66 @@ +/* + * 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 type { ScopedModel } from '@kbn/onechat-server'; +import { indexExplorer } from './index_explorer'; +import { flattenMappings } from './utils'; +import { getIndexMappings, performMatchSearch, PerformMatchSearchResponse } from './steps'; + +export type RelevanceSearchResponse = PerformMatchSearchResponse; + +export const relevanceSearch = async ({ + term, + index, + fields = [], + size = 10, + model, + esClient, +}: { + term: string; + index?: string; + fields?: string[]; + size?: number; + model: ScopedModel; + esClient: ElasticsearchClient; +}): Promise => { + let selectedIndex = index; + let selectedFields = fields; + + if (!selectedIndex) { + const { indices } = await indexExplorer({ + query: term, + esClient, + model, + }); + if (indices.length === 0) { + return { results: [] }; + } + selectedIndex = indices[0].indexName; + } + + if (!fields.length) { + const mappings = await getIndexMappings({ + indices: [selectedIndex], + esClient, + }); + + const flattenedFields = flattenMappings(mappings[selectedIndex]); + + selectedFields = flattenedFields + .filter((field) => field.type === 'text' || field.type === 'semantic_text') + .map((field) => field.path); + } + + return performMatchSearch({ + term, + fields: selectedFields, + index: selectedIndex, + size, + esClient, + }); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/execute_esql.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/execute_esql.ts new file mode 100644 index 0000000000000..137cf6624c356 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/execute_esql.ts @@ -0,0 +1,31 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { EsqlEsqlColumnInfo, FieldValue } from '@elastic/elasticsearch/lib/api/types'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; + +export interface EsqlResponse { + columns: EsqlEsqlColumnInfo[]; + values: FieldValue[][]; +} + +/** + * Execute an ES|QL query and returns the response. + */ +export const executeEsql = async ({ + query, + esClient, +}: { + query: string; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.esql.query({ query, drop_null_columns: true }); + return { + columns: response.columns, + values: response.values, + }; +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_documents.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_documents.ts new file mode 100644 index 0000000000000..fbe76f46e8a05 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_documents.ts @@ -0,0 +1,50 @@ +/* + * 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'; + +export interface GetDocumentByIdSuccess { + id: string; + index: string; + found: true; + _source: unknown; +} + +export interface GetDocumentByIdFailure { + id: string; + index: string; + found: false; +} + +export type GetDocumentByIdResult = GetDocumentByIdSuccess | GetDocumentByIdFailure; + +export const getDocumentById = async ({ + id, + index, + esClient, +}: { + id: string; + index: string; + esClient: ElasticsearchClient; +}): Promise => { + const { body: response, statusCode } = await esClient.get( + { + id, + index, + }, + { ignore: [404], meta: true } + ); + if (statusCode === 404) { + return { id, index, found: false }; + } + return { + id, + index, + found: true, + _source: response._source ?? {}, + }; +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_mappings.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_mappings.ts new file mode 100644 index 0000000000000..e7d2ba5580dc3 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/get_mappings.ts @@ -0,0 +1,37 @@ +/* + * 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 type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; +import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { cleanupMapping } from '../utils'; + +export interface GetIndexMappingEntry { + mappings: MappingTypeMapping; +} + +export type GetIndexMappingsResult = Record; + +export const getIndexMappings = async ({ + indices, + cleanup = true, + esClient, +}: { + indices: string[]; + cleanup?: boolean; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.indices.getMapping({ + index: indices, + }); + + return Object.entries(response).reduce((res, [indexName, mappingRes]) => { + res[indexName] = { + mappings: cleanup ? cleanupMapping(mappingRes.mappings) : mappingRes.mappings, + }; + return res; + }, {} as GetIndexMappingsResult); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/index.ts new file mode 100644 index 0000000000000..1380403476d1c --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/index.ts @@ -0,0 +1,20 @@ +/* + * 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. + */ + +export { getDocumentById, type GetDocumentByIdResult } from './get_documents'; +export { + getIndexMappings, + type GetIndexMappingEntry, + type GetIndexMappingsResult, +} from './get_mappings'; +export { + performMatchSearch, + type PerformMatchSearchResponse, + type MatchResult, +} from './perform_match_search'; +export { executeEsql, type EsqlResponse } from './execute_esql'; +export { listIndices, type ListIndexInfo } from './list_indices'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/list_indices.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/list_indices.ts new file mode 100644 index 0000000000000..c17bc9970720d --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/list_indices.ts @@ -0,0 +1,41 @@ +/* + * 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'; + +export interface ListIndexInfo { + index: string; + status: string; + health: string; + uuid: string; + docsCount: number; + primaries: number; + replicas: number; +} + +export const listIndices = async ({ + pattern = '*', + esClient, +}: { + pattern?: string; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.cat.indices({ + index: pattern, + format: 'json', + }); + + return response.map(({ index, status, health, uuid, 'docs.count': docsCount, pri, rep }) => ({ + index: index ?? 'unknown', + status: status ?? 'unknown', + health: health ?? 'unknown', + uuid: uuid ?? 'unknown', + docsCount: parseInt(docsCount ?? '0', 10), + primaries: parseInt(pri ?? '1', 10), + replicas: parseInt(rep ?? '0', 10), + })); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/perform_match_search.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/perform_match_search.ts new file mode 100644 index 0000000000000..18816bfc89653 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/steps/perform_match_search.ts @@ -0,0 +1,69 @@ +/* + * 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'; + +export interface MatchResult { + id: string; + index: string; + highlights: string[]; +} + +export interface PerformMatchSearchResponse { + results: MatchResult[]; +} + +export const performMatchSearch = async ({ + term, + fields, + index, + size, + esClient, +}: { + term: string; + fields: string[]; + index: string; + size: number; + esClient: ElasticsearchClient; +}): Promise => { + const response = await esClient.search({ + index, + size, + retriever: { + rrf: { + retrievers: fields.map((field) => { + return { + standard: { + query: { + match: { + [field]: term, + }, + }, + }, + }; + }), + }, + }, + highlight: { + number_of_fragments: 5, + fields: fields.reduce((memo, field) => ({ ...memo, [field]: {} }), {}), + }, + }); + + const results = response.hits.hits.map((hit) => { + return { + id: hit._id!, + index: hit._index!, + highlights: Object.entries(hit.highlight ?? {}).reduce((acc, [field, highlights]) => { + acc.push(...highlights); + return acc; + }, [] as string[]), + }; + }); + + return { results }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql.ts similarity index 61% rename from x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts rename to x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql.ts index dec1d2259f334..88ca462d51a7f 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/esql.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql.ts @@ -5,13 +5,12 @@ * 2.0. */ -import type { EsqlEsqlColumnInfo, FieldValue } from '@elastic/elasticsearch/lib/api/types'; - -export interface EsqlResponse { - columns: EsqlEsqlColumnInfo[]; - values: FieldValue[][]; -} +import { INLINE_ESQL_QUERY_REGEX } from '@kbn/inference-plugin/common/tasks/nl_to_esql/constants'; +import type { EsqlResponse } from '../steps/execute_esql'; +/** + * Converts an ES|QL /_query columnar response to a JSON representation + */ export const esqlResponseToJson = (esql: EsqlResponse): Array> => { const results: Array> = []; @@ -26,3 +25,7 @@ export const esqlResponseToJson = (esql: EsqlResponse): Array { + return Array.from(message.matchAll(INLINE_ESQL_QUERY_REGEX)).map(([match, query]) => query); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/index.ts new file mode 100644 index 0000000000000..da4437c57368e --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/index.ts @@ -0,0 +1,9 @@ +/* + * 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. + */ + +export { esqlResponseToJson, extractEsqlQueries } from './esql'; +export { flattenMappings, cleanupMapping, type MappingField } from './mappings'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/mappings.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/mappings.ts new file mode 100644 index 0000000000000..6ef8fc8b8afa4 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/mappings.ts @@ -0,0 +1,91 @@ +/* + * 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 type { MappingTypeMapping, MappingProperty } from '@elastic/elasticsearch/lib/api/types'; + +export type FieldType = Extract['type']; + +export interface MappingField { + path: string; + type: FieldType; +} + +interface MappingProperties { + [key: string]: { + type?: string; // Leaf field (e.g., "text", "keyword", etc.) + properties?: MappingProperties; // Nested object fields + }; +} + +/** + * Returns a flattened representation of the mappings, with all fields at the top level. + */ +export const flattenMappings = ({ mappings }: { mappings: MappingTypeMapping }): MappingField[] => { + const properties: MappingProperties = mappings.properties ?? {}; + + function extractFields(obj: MappingProperties, prefix = ''): MappingField[] { + let fields: MappingField[] = []; + + for (const [key, value] of Object.entries(obj)) { + const fieldPath = prefix ? `${prefix}.${key}` : key; + + if (value.type) { + // If it's a leaf field, add it + fields.push({ + type: value.type as FieldType, + path: fieldPath, + }); + } else if (value.properties) { + // If it's an object, go deeper + fields = fields.concat(extractFields(value.properties, fieldPath)); + } + } + + return fields; + } + + return extractFields(properties); +}; + +/** + * Remove non-relevant mapping information such as `ignore_above` to reduce overall token length of response + * @param mapping + */ +export const cleanupMapping = (mapping: MappingTypeMapping): MappingTypeMapping => { + const recurseKeys = ['properties', 'fields']; + const fieldsToKeep = ['type', 'dynamic', '_meta', 'enabled']; + + function recursiveCleanup(obj: Record): Record { + if (Array.isArray(obj)) { + return obj.map((item) => recursiveCleanup(item)); + } else if (obj !== null && typeof obj === 'object') { + const cleaned: Record = {}; + + for (const key of Object.keys(obj)) { + if (recurseKeys.includes(key)) { + const value = obj[key]; + if (value !== null && typeof value === 'object') { + // For properties/fields: preserve all keys inside + const subCleaned: Record = {}; + for (const fieldName of Object.keys(value)) { + subCleaned[fieldName] = recursiveCleanup(value[fieldName]); + } + cleaned[key] = subCleaned; + } + } else if (fieldsToKeep.includes(key)) { + cleaned[key] = recursiveCleanup(obj[key]); + } + } + + return cleaned; + } else { + return obj; + } + } + + return recursiveCleanup(mapping); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json new file mode 100644 index 0000000000000..63f0b5ff33faa --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json @@ -0,0 +1,17 @@ +{ + "extends": "../../../../../../tsconfig.base.json", + "compilerOptions": { + "outDir": "target/types", + "types": [ + "jest", + "node" + ] + }, + "include": [ + "**/*.ts", + ], + "exclude": [ + "target/**/*" + ], + "kbn_references": [] +} diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts index 93ac45917c8c5..35a569c628278 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts @@ -6,7 +6,7 @@ */ import type { BaseMessage, BaseMessageLike } from '@langchain/core/messages'; -import { OnechatToolIds } from '@kbn/onechat-common'; +import { BuiltinToolIds } from '@kbn/onechat-common'; export const defaultSystemPrompt = ` You are a helpful chat assistant from the Elasticsearch company. @@ -18,11 +18,13 @@ export const defaultSystemPrompt = ` - When the user ask a question, assume it refers to information that can be retrieved from Elasticsearch. For example if the user asks "What are my latest alerts", assume you need to search the cluster for documents. - - When doing fulltext search, prefer the ${OnechatToolIds.searchFulltext} tool over the ${OnechatToolIds.searchDsl} one - when possible. + - Your two main search tools are "${BuiltinToolIds.relevanceSearch}" and "${BuiltinToolIds.naturalLanguageSearch}" + - When doing fulltext search, prefer the "${BuiltinToolIds.relevanceSearch}" tool as it performs better for plain fulltext searches. + - For more advanced queries, use the "${BuiltinToolIds.naturalLanguageSearch}" tool. - - Search tools return highlights of the documents that match the query. The full content of a document can be retrieved - using the ${OnechatToolIds.getDocumentById} tool. + - Never call the "${BuiltinToolIds.executeEsql}" tool without a valid ES|QL query generated by the "${BuiltinToolIds.generateEsql}" tool. + - More generally, only use the ES|QL tools ("${BuiltinToolIds.executeEsql}" and "${BuiltinToolIds.generateEsql}") if the user explicitly asks + to either generate or execute an ES|QL query. Prefer the "${BuiltinToolIds.naturalLanguageSearch}" otherwise. `; const getFullSystemPrompt = (systemPrompt: string) => { diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts index a2c286d1aa6b4..be1863f8e7e2e 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts @@ -7,7 +7,7 @@ import { z } from '@kbn/zod'; import type { RegisteredTool } from '@kbn/onechat-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { runSearchAgent } from './run_researcher_agent'; @@ -38,7 +38,7 @@ export const researcherTool = (): RegisteredTool { const tools: Array> = [ - // listIndicesTool(), - // getIndexMappingsTool(), getDocumentByIdTool(), - searchFulltextTool(), - // searchDslTool(), - // rerankDocumentsTool(), - nlToEsqlTool(), executeEsqlTool(), - researcherTool(), + naturalLanguageSearchTool(), + generateEsqlTool(), + relevanceSearchTool(), + getIndexMappingsTool(), + listIndicesTool(), indexExplorerTool(), + researcherTool(), ]; tools.forEach((tool) => { diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts index 671b2be13061e..48630bd154728 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.ts @@ -6,47 +6,24 @@ */ import { z } from '@kbn/zod'; -import type { EsqlEsqlColumnInfo, FieldValue } from '@elastic/elasticsearch/lib/api/types'; -import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; +import { executeEsql, EsqlResponse } from '@kbn/onechat-genai-utils'; const executeEsqlToolSchema = z.object({ query: z.string().describe('The ES|QL query to execute'), }); -export interface ExecuteEsqlResponse { - columns: EsqlEsqlColumnInfo[]; - values: FieldValue[][]; -} - -export const executeEsqlTool = (): RegisteredTool< - typeof executeEsqlToolSchema, - ExecuteEsqlResponse -> => { +export const executeEsqlTool = (): RegisteredTool => { return { - id: OnechatToolIds.executeEsql, + id: BuiltinToolIds.executeEsql, description: 'Execute an ES|QL query and return the results.', schema: executeEsqlToolSchema, handler: async ({ query }, { esClient }) => { return executeEsql({ query, esClient: esClient.asCurrentUser }); }, meta: { - tags: [OnechatToolTags.retrieval], + tags: [BuiltinTags.retrieval], }, }; }; - -export const executeEsql = async ({ - query, - esClient, -}: { - query: string; - esClient: ElasticsearchClient; -}): Promise => { - const response = await esClient.esql.query({ query, drop_null_columns: true }); - return { - columns: response.columns, - values: response.values, - }; -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts new file mode 100644 index 0000000000000..9200b2f25b47b --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts @@ -0,0 +1,46 @@ +/* + * 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 { z } from '@kbn/zod'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; +import { generateEsql, GenerateEsqlResponse } from '@kbn/onechat-genai-utils'; + +const nlToEsqlToolSchema = z.object({ + query: z.string().describe('The query to generate an ES|QL query from.'), + index: z + .string() + .optional() + .describe( + '(optional) Index to search against. If not provided, will use the index explorer to find the best index to use.' + ), + context: z + .string() + .optional() + .describe('(optional) Additional context that could be useful to generate the ES|QL query'), +}); + +export const generateEsqlTool = (): RegisteredTool => { + return { + id: BuiltinToolIds.generateEsql, + description: 'Generate an ES|QL query from a natural language query.', + schema: nlToEsqlToolSchema, + handler: async ({ query, index, context }, { esClient, modelProvider }) => { + const model = await modelProvider.getDefaultModel(); + return generateEsql({ + query, + context, + index, + model, + esClient: esClient.asCurrentUser, + }); + }, + meta: { + tags: [BuiltinTags.retrieval], + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts index 486167cf7025f..c44b1e6c7296e 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts @@ -6,68 +6,28 @@ */ import { z } from '@kbn/zod'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; +import { getDocumentById, GetDocumentByIdResult } from '@kbn/onechat-genai-utils'; const getDocumentByIdSchema = z.object({ id: z.string().describe('ID of the document to retrieve'), index: z.string().describe('Name of the index to retrieve the document from'), }); -export type GetDocumentByIdResult = - | { - id: string; - index: string; - found: true; - _source: unknown; - } - | { - id: string; - index: string; - found: false; - }; - export const getDocumentByIdTool = (): RegisteredTool< typeof getDocumentByIdSchema, GetDocumentByIdResult > => { return { - id: OnechatToolIds.getDocumentById, + id: BuiltinToolIds.getDocumentById, description: 'Retrieve the full content (source) of a document based on its ID and index name.', schema: getDocumentByIdSchema, handler: async ({ id, index }, { esClient }) => { return getDocumentById({ id, index, esClient: esClient.asCurrentUser }); }, meta: { - tags: [OnechatToolTags.retrieval], - }, - }; -}; - -export const getDocumentById = async ({ - id, - index, - esClient, -}: { - id: string; - index: string; - esClient: ElasticsearchClient; -}): Promise => { - const { body: response, statusCode } = await esClient.get( - { - id, - index, + tags: [BuiltinTags.retrieval], }, - { ignore: [404], meta: true } - ); - if (statusCode === 404) { - return { id, index, found: false }; - } - return { - id, - index, - found: true, - _source: response._source ?? {}, }; }; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts index d9818755fdb00..9beb2b69f5bee 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.ts @@ -5,91 +5,28 @@ * 2.0. */ -import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; import { z } from '@kbn/zod'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; +import { getIndexMappings, GetIndexMappingsResult } from '@kbn/onechat-genai-utils'; const getIndexMappingsSchema = z.object({ indices: z.array(z.string()).min(1).describe('List of indices to retrieve mappings for.'), }); -export interface GetIndexMappingEntry { - mappings: MappingTypeMapping; -} - -export type GetIndexMappingsResult = Record; - export const getIndexMappingsTool = (): RegisteredTool< typeof getIndexMappingsSchema, GetIndexMappingsResult > => { return { - id: OnechatToolIds.getIndexMapping, + id: BuiltinToolIds.getIndexMapping, description: 'Retrieve mappings for the specified index or indices.', schema: getIndexMappingsSchema, handler: async ({ indices }, { esClient }) => { return getIndexMappings({ indices, esClient: esClient.asCurrentUser }); }, meta: { - tags: [OnechatToolTags.retrieval], + tags: [BuiltinTags.retrieval], }, }; }; - -export const getIndexMappings = async ({ - indices, - esClient, -}: { - indices: string[]; - esClient: ElasticsearchClient; -}): Promise => { - const response = await esClient.indices.getMapping({ - index: indices, - }); - - return Object.entries(response).reduce((res, [indexName, mappingRes]) => { - res[indexName] = { mappings: cleanupMapping(mappingRes.mappings) }; - return res; - }, {} as GetIndexMappingsResult); -}; - -/** - * Remove non-relevant mapping information such as `ignore_above` to reduce overall token length of response - * @param mapping - */ -const cleanupMapping = (mapping: MappingTypeMapping): MappingTypeMapping => { - const recurseKeys = ['properties', 'fields']; - const fieldsToKeep = ['type', 'dynamic', '_meta', 'enabled']; - - function recursiveCleanup(obj: Record): Record { - if (Array.isArray(obj)) { - return obj.map((item) => recursiveCleanup(item)); - } else if (obj !== null && typeof obj === 'object') { - const cleaned: Record = {}; - - for (const key of Object.keys(obj)) { - if (recurseKeys.includes(key)) { - const value = obj[key]; - if (value !== null && typeof value === 'object') { - // For properties/fields: preserve all keys inside - const subCleaned: Record = {}; - for (const fieldName of Object.keys(value)) { - subCleaned[fieldName] = recursiveCleanup(value[fieldName]); - } - cleaned[key] = subCleaned; - } - } else if (fieldsToKeep.includes(key)) { - cleaned[key] = recursiveCleanup(obj[key]); - } - } - - return cleaned; - } else { - return obj; - } - } - - return recursiveCleanup(mapping); -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts index b231af7dda245..859da7c162dcf 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts @@ -8,9 +8,8 @@ export { getDocumentByIdTool } from './get_document_by_id'; export { getIndexMappingsTool } from './get_index_mapping'; export { listIndicesTool } from './list_indices'; -export { rerankDocumentsTool } from './rerank_documents'; -export { searchDslTool } from './search_dsl'; -export { searchFulltextTool } from './search_fulltext'; -export { nlToEsqlTool } from './nl_to_esql'; -export { executeEsqlTool } from './execute_esql'; export { indexExplorerTool } from './index_explorer'; +export { generateEsqlTool } from './generate_esql'; +export { executeEsqlTool } from './execute_esql'; +export { naturalLanguageSearchTool } from './nl_search'; +export { relevanceSearchTool } from './relevance_search'; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts index 02dc42797b69f..34d290ecd2b86 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts @@ -5,15 +5,10 @@ * 2.0. */ -import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; - import { z } from '@kbn/zod'; -import { BaseMessageLike } from '@langchain/core/messages'; -import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; -import type { RegisteredTool, ScopedModel } from '@kbn/onechat-server'; -import { listIndices } from './list_indices'; -import { getIndexMappings } from './get_index_mapping'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; +import { indexExplorer, IndexExplorerResponse } from '@kbn/onechat-genai-utils'; const indexExplorerSchema = z.object({ query: z.string().describe('A natural language query to infer which indices to use.'), @@ -27,22 +22,12 @@ const indexExplorerSchema = z.object({ .describe('(optional) Index pattern to filter indices by. Defaults to *.'), }); -export interface RelevantIndex { - indexName: string; - mappings: MappingTypeMapping; - reason: string; -} - -export interface IndexExplorerResponse { - indices: RelevantIndex[]; -} - export const indexExplorerTool = (): RegisteredTool< typeof indexExplorerSchema, IndexExplorerResponse > => { return { - id: OnechatToolIds.indexExplorer, + id: BuiltinToolIds.indexExplorer, description: `List relevant indices and corresponding mappings based on a natural language query. The 'indexPattern' parameter can be used to filter indices by a specific pattern, e.g. 'foo*'. @@ -60,75 +45,7 @@ export const indexExplorerTool = (): RegisteredTool< return indexExplorer({ query, indexPattern, limit, esClient: esClient.asCurrentUser, model }); }, meta: { - tags: [OnechatToolTags.retrieval], + tags: [BuiltinTags.retrieval], }, }; }; - -export const indexExplorer = async ({ - query, - indexPattern = '*', - limit = 1, - esClient, - model, -}: { - query: string; - indexPattern?: string; - limit?: number; - esClient: ElasticsearchClient; - model: ScopedModel; -}): Promise => { - const { chatModel } = model; - - const allIndices = await listIndices({ - pattern: indexPattern, - esClient, - }); - - const indexSelectorModel = chatModel.withStructuredOutput( - z.object({ - indices: z.array( - z.object({ - indexName: z.string().describe('name of the index'), - reason: z.string().describe('brief explanation of why this index could be relevant'), - }) - ), - }) - ); - - const indexSelectorPrompt: BaseMessageLike[] = [ - [ - 'user', - `You are an AI assistant for the Elasticsearch company. - based on a natural language query from the user, your task is to select up to ${limit} most relevant indices from a list of indices. - - *The query is:* ${query} - - *List of indices:* - ${allIndices.map((index) => `- ${index.index}`).join('\n')} - - Based on those information, please return most relevant indices with your reasoning. - Remember, you should select at maximum ${limit} indices. - `, - ], - ]; - - const { indices: selectedIndices } = await indexSelectorModel.invoke(indexSelectorPrompt); - - const mappings = await getIndexMappings({ - indices: selectedIndices.map((index) => index.indexName), - esClient, - }); - - const relevantIndices: RelevantIndex[] = selectedIndices.map( - ({ indexName, reason }) => { - return { - indexName, - reason, - mappings: mappings[indexName].mappings, - }; - } - ); - - return { indices: relevantIndices }; -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts index c39f2885e9406..279409670328d 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.ts @@ -6,9 +6,9 @@ */ import { z } from '@kbn/zod'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; +import { listIndices, ListIndexInfo } from '@kbn/onechat-genai-utils'; const listIndicesSchema = z.object({ pattern: z @@ -19,49 +19,16 @@ const listIndicesSchema = z.object({ ), }); -export interface ListIndexInfo { - index: string; - status: string; - health: string; - uuid: string; - docsCount: number; - primaries: number; - replicas: number; -} - export const listIndicesTool = (): RegisteredTool => { return { - id: OnechatToolIds.listIndices, + id: BuiltinToolIds.listIndices, description: 'List the indices in the Elasticsearch cluster the current user has access to.', schema: listIndicesSchema, handler: async ({ pattern = '*' }, { esClient }) => { return listIndices({ pattern, esClient: esClient.asCurrentUser }); }, meta: { - tags: [OnechatToolTags.retrieval], + tags: [BuiltinTags.retrieval], }, }; }; - -export const listIndices = async ({ - pattern = '*', - esClient, -}: { - pattern?: string; - esClient: ElasticsearchClient; -}): Promise => { - const response = await esClient.cat.indices({ - index: pattern, - format: 'json', - }); - - return response.map(({ index, status, health, uuid, 'docs.count': docsCount, pri, rep }) => ({ - index: index ?? 'unknown', - status: status ?? 'unknown', - health: health ?? 'unknown', - uuid: uuid ?? 'unknown', - docsCount: parseInt(docsCount ?? '0', 10), - primaries: parseInt(pri ?? '1', 10), - replicas: parseInt(rep ?? '0', 10), - })); -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_search.ts similarity index 59% rename from x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts rename to x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_search.ts index 0c0ac83603d3f..a875b42c1808d 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_dsl.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_search.ts @@ -6,12 +6,10 @@ */ import { z } from '@kbn/zod'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import type { RegisteredTool } from '@kbn/onechat-server'; -import { generateEsql } from './nl_to_esql'; -import { executeEsql, ExecuteEsqlResponse } from './execute_esql'; +import { naturalLanguageSearch, NaturalLanguageSearchResponse } from '@kbn/onechat-genai-utils'; -// smart search const searchDslSchema = z.object({ query: z.string().describe('A natural language query expressing the search request'), index: z @@ -26,37 +24,26 @@ const searchDslSchema = z.object({ .describe('(optional) Additional context that could be useful to perform the search'), }); -export type SearchDslResponse = ExecuteEsqlResponse | { success: false; reason: string }; - -export const searchDslTool = (): RegisteredTool => { +export const naturalLanguageSearchTool = (): RegisteredTool< + typeof searchDslSchema, + NaturalLanguageSearchResponse +> => { return { - id: OnechatToolIds.searchDsl, + id: BuiltinToolIds.naturalLanguageSearch, description: 'Run a DSL search query on one index and return matching documents.', schema: searchDslSchema, handler: async ({ query, index, context }, { esClient, modelProvider }) => { const model = await modelProvider.getDefaultModel(); - - const generateResponse = await generateEsql({ + return naturalLanguageSearch({ query, context, index, model, esClient: esClient.asCurrentUser, }); - - if (generateResponse.queries.length < 1) { - return { success: false, reason: 'No query was generated' }; - } - - const executeResponse = await executeEsql({ - query: generateResponse.queries[0], - esClient: esClient.asCurrentUser, - }); - - return executeResponse; }, meta: { - tags: [OnechatToolTags.retrieval], + tags: [BuiltinTags.retrieval], }, }; }; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/relevance_search.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/relevance_search.ts new file mode 100644 index 0000000000000..d174172127281 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/relevance_search.ts @@ -0,0 +1,73 @@ +/* + * 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 { z } from '@kbn/zod'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; +import { relevanceSearch } from '@kbn/onechat-genai-utils'; + +const relevanceSearchSchema = z.object({ + term: z.string().describe('Term to search for'), + index: z + .string() + .optional() + .describe( + '(optional) Index to search against. If not provided, will use index explorer to find the best index to use.' + ), + fields: z + .array(z.string()) + .optional() + .describe( + '(optional) Fields to perform fulltext search on. If not provided, will use all searchable fields.' + ), + size: z + .number() + .optional() + .default(10) + .describe('Number of documents to return. Defaults to 10.'), +}); + +export interface SearchFulltextResult { + id: string; + index: string; + highlights: string[]; +} + +export interface SearchFulltextResponse { + results: SearchFulltextResult[]; +} + +export const relevanceSearchTool = (): RegisteredTool< + typeof relevanceSearchSchema, + SearchFulltextResponse +> => { + return { + id: BuiltinToolIds.relevanceSearch, + description: `Find relevant documents in an index based on a simple fulltext search. + + - The 'index' parameter can be used to specify which index to search against. If not provided, the tool will use the index explorer to find the best index to use. + - The 'fields' parameter can be used to specify which fields to search on. If not provided, the tool will use all searchable fields. + + It is perfectly fine not to not specify both 'index' and 'fields'. Those should only be used when you already know about the index and fields you want to search on, + e.g if the user explicitly specified them.`, + schema: relevanceSearchSchema, + handler: async ({ term, index, fields = [], size }, { esClient, modelProvider }) => { + const model = await modelProvider.getDefaultModel(); + return relevanceSearch({ + term, + index, + fields, + size, + model, + esClient: esClient.asCurrentUser, + }); + }, + meta: { + tags: [BuiltinTags.retrieval], + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts deleted file mode 100644 index 2b188590c3ae3..0000000000000 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/rerank_documents.ts +++ /dev/null @@ -1,196 +0,0 @@ -/* - * 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 { z } from '@kbn/zod'; -import { BaseMessageLike } from '@langchain/core/messages'; -import type { InferenceChatModel } from '@kbn/inference-langchain'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; -import type { RegisteredTool } from '@kbn/onechat-server'; - -const rerankDocumentsSchema = z.object({ - query: z.string().describe('Text query to rerank snippets by.'), - documents: z - .array( - z.object({ - id: z.string().describe('ID of the document.'), - index: z.string().optional().describe('Index the document is from, if applicable.'), - snippet: z.string().describe('Text snippet to use for reranking.'), - }) - ) - .min(1) - .describe('Documents to rerank'), -}); - -interface DocumentWithRating { - id: string; - index?: string; - snippet: string; - rating: number; -} - -interface RerankResponse { - documents: DocumentWithRating[]; -} - -export const rerankDocumentsTool = (): RegisteredTool< - typeof rerankDocumentsSchema, - RerankResponse -> => { - return { - id: OnechatToolIds.rerankDocuments, - description: 'Score and rerank documents based their relevance against a text query.', - schema: rerankDocumentsSchema, - handler: async ({ query, documents }, { modelProvider }) => { - const { chatModel } = await modelProvider.getDefaultModel(); - - const rerankDocs = documents.map((doc) => ({ - id: doc.id, - content: doc.snippet, - })); - - const docsWithRatings = await rerankDocuments({ - query, - documents: rerankDocs, - chatModel, - }); - - const rerankedDocs = docsWithRatings - .map((doc) => { - const matchingDoc = documents.find((d) => d.id === doc.id)!; - return { - ...matchingDoc, - rating: doc.rating, - }; - }) - .sort((a, b) => b.rating - a.rating); - - return { - documents: rerankedDocs, - }; - }, - meta: { - tags: [OnechatToolTags.retrieval], - }, - }; -}; - -export const rerankDocuments = async ({ - query, - documents, - chatModel, -}: { - query: string; - documents: RerankingDoc[]; - chatModel: InferenceChatModel; -}): Promise => { - const analysisModel = chatModel.withStructuredOutput( - z.object({ - ratings: z - .array( - z.object({ - id: z.string().describe('ID of the document'), - grade: z.number().describe('Score of the document, between 0 and 10'), - reason: z - .string() - .optional() - .describe('Optional reason for the rating. Keep it short and concise.'), - }) - ) - .describe('the ratings, one per document using the "{id, grade, reason}" format.'), - }) - ); - - const { ratings } = await analysisModel.invoke(getRerankingPrompt({ query, documents })); - - const ratingMap = ratings.reduce((acc, rating) => { - acc[rating.id] = { - grade: rating.grade, - }; - return acc; - }, {} as Record); - - const documentsWithRatings = documents.map((doc) => ({ - ...doc, - rating: ratingMap[doc.id] ?? 0, - })); - - return documentsWithRatings; -}; - -interface RerankingDoc { - id: string; - content: unknown; -} - -interface RerankedDoc { - id: string; - content: unknown; - rating: number; -} - -export const getRerankingPrompt = ({ - query, - documents, -}: { - query: string; - documents: RerankingDoc[]; -}): BaseMessageLike[] => { - const resultEntry = (document: RerankingDoc): string => { - return ` - ### Document (ID: ${document.id}) - - **Document ID:** "${document.id}" - - **Content:** - \`\`\` - ${JSON.stringify(document.content, null, 2)} - \`\`\` - `; - }; - - return [ - [ - 'system', - ` - ## Current task: Relevance Analysis - - Your task is to evaluate at set of documents in relation to the user’s query, - and assign a relevance rating from 0 to 10 using the following criteria: - - **0:** The document is completely irrelevant. - - **5:** The document is somewhat related and might be useful. - - **8:** The document is very relevant and contains useful information. - - **10:** The document is absolutely crucial for answering the query. - - **Instructions:** - - **Independent Ratings:** Rate each document independently based solely on its relevance to the provided query. - - **Format:** Return your ratings as a JSON object with a "ratings" array, where each element follows the \`"{id}|{grade}"\` format. Example: \`{"ratings": ["0|7", "1|5", "2|10"]}\`. - - **Document IDs:** Use the document IDs provided in this prompt, not any IDs contained in the document content. - - **Optional Comments:** You may include an optional \`"comment"\` field with additional remarks on your ratings. - - ## Input - - ## Input - - You will receive: - 1. The search query from the user. - 2. A list of documents, each with an assigned document ID, that were retrieved in the previous step. - `, - ], - [ - 'human', - ` - ## Input - - **Search Query:**: "${query}" - - ## Documents - - ${documents.map(resultEntry).join('\n')} - `, - ], - ]; -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts deleted file mode 100644 index 20fabd0be5cb6..0000000000000 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/search_fulltext.ts +++ /dev/null @@ -1,156 +0,0 @@ -/* - * 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 { z } from '@kbn/zod'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { OnechatToolIds, OnechatToolTags } from '@kbn/onechat-common'; -import type { RegisteredTool } from '@kbn/onechat-server'; -import { indexExplorer } from './index_explorer'; -import { getIndexMappings } from './get_index_mapping'; -import { flattenFields } from './utils/flatten_fields'; - -const fulltextSearchSchema = z.object({ - term: z.string().describe('Term to search for'), - index: z - .string() - .optional() - .describe( - '(optional) Index to search against. If not provided, will use index explorer to find the best index to use.' - ), - fields: z - .array(z.string()) - .optional() - .describe( - '(optional) Fields to perform fulltext search on. If not provided, will use all searchable fields.' - ), - size: z - .number() - .optional() - .default(10) - .describe('Number of documents to return. Defaults to 10.'), -}); - -export interface SearchFulltextResult { - id: string; - index: string; - highlights: string[]; -} - -export interface SearchFulltextResponse { - results: SearchFulltextResult[]; -} - -// TODO: rename to relevance search -export const searchFulltextTool = (): RegisteredTool< - typeof fulltextSearchSchema, - SearchFulltextResponse -> => { - return { - id: OnechatToolIds.searchFulltext, - description: `Find relevant documents in an index based on a simple fulltext search. - - - The 'index' parameter can be used to specify which index to search against. If not provided, the tool will use the index explorer to find the best index to use. - - The 'fields' parameter can be used to specify which fields to search on. If not provided, the tool will use all searchable fields. - - It is perfectly fine not to not specify both 'index' and 'fields'. Those should only be used when you already know about the index and fields you want to search on, - e.g if the user explicitly specified them.`, - schema: fulltextSearchSchema, - handler: async ({ term, index, fields = [], size }, { esClient, modelProvider }) => { - const model = await modelProvider.getDefaultModel(); - - let selectedIndex = index; - let selectedFields = fields; - - if (!selectedIndex) { - const { indices } = await indexExplorer({ - query: term, - esClient: esClient.asCurrentUser, - model, - }); - if (indices.length === 0) { - return { results: [] }; - } - selectedIndex = indices[0].indexName; - } - - if (!fields.length) { - const mappings = await getIndexMappings({ - indices: [selectedIndex], - esClient: esClient.asCurrentUser, - }); - - const flattenedFields = flattenFields(mappings[selectedIndex]); - - selectedFields = flattenedFields - .filter((field) => field.type === 'text' || field.type === 'semantic_text') - .map((field) => field.path); - } - - return searchFulltext({ - term, - fields: selectedFields, - index: selectedIndex, - size, - esClient: esClient.asCurrentUser, - }); - }, - meta: { - tags: [OnechatToolTags.retrieval], - }, - }; -}; - -export const searchFulltext = async ({ - term, - fields, - index, - size, - esClient, -}: { - term: string; - fields: string[]; - index: string; - size: number; - esClient: ElasticsearchClient; -}): Promise => { - const response = await esClient.search({ - index, - size, - retriever: { - rrf: { - retrievers: fields.map((field) => { - return { - standard: { - query: { - match: { - [field]: term, - }, - }, - }, - }; - }), - }, - }, - highlight: { - number_of_fragments: 5, - fields: fields.reduce((memo, field) => ({ ...memo, [field]: {} }), {}), - }, - }); - - const results = response.hits.hits.map((hit) => { - return { - id: hit._id!, - index: hit._index!, - highlights: Object.entries(hit.highlight ?? {}).reduce((acc, [field, highlights]) => { - acc.push(...highlights); - return acc; - }, [] as string[]), - }; - }); - - return { results }; -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts deleted file mode 100644 index 77230f9900225..0000000000000 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/utils/flatten_fields.ts +++ /dev/null @@ -1,49 +0,0 @@ -/* - * 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 type { MappingTypeMapping, MappingProperty } from '@elastic/elasticsearch/lib/api/types'; - -export type FieldType = Extract['type']; - -export interface MappingField { - path: string; - type: FieldType; -} - -export interface MappingProperties { - [key: string]: { - type?: string; // Leaf field (e.g., "text", "keyword", etc.) - properties?: MappingProperties; // Nested object fields - }; -} - -export const flattenFields = ({ mappings }: { mappings: MappingTypeMapping }): MappingField[] => { - const properties: MappingProperties = mappings.properties ?? {}; - - function extractFields(obj: MappingProperties, prefix = ''): MappingField[] { - let fields: MappingField[] = []; - - for (const [key, value] of Object.entries(obj)) { - const fieldPath = prefix ? `${prefix}.${key}` : key; - - if (value.type) { - // If it's a leaf field, add it - fields.push({ - type: value.type as FieldType, - path: fieldPath, - }); - } else if (value.properties) { - // If it's an object, go deeper - fields = fields.concat(extractFields(value.properties, fieldPath)); - } - } - - return fields; - } - - return extractFields(properties); -}; diff --git a/yarn.lock b/yarn.lock index 08eca5626d051..97471afbbd37a 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6574,6 +6574,10 @@ version "0.0.0" uid "" +"@kbn/onechat-genai-utils@link:x-pack/platform/packages/shared/onechat/onechat-genai-utils": + version "0.0.0" + uid "" + "@kbn/onechat-plugin@link:x-pack/platform/plugins/shared/onechat": version "0.0.0" uid "" From 2243e44aa2947713aca2404314ae91fc0c06ddb4 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Mon, 16 Jun 2025 14:22:10 +0200 Subject: [PATCH 08/17] move / rename things --- .../packages/shared/onechat/onechat-common/tools/constants.ts | 1 + .../shared/onechat/server/services/agents/agents_service.ts | 2 +- .../agents/{conversational => chat}/convert_graph_events.ts | 0 .../server/services/agents/{conversational => chat}/graph.ts | 0 .../server/services/agents/{conversational => chat}/handler.ts | 0 .../server/services/agents/{conversational => chat}/index.ts | 0 .../services/agents/{conversational => chat}/provider.ts | 0 .../services/agents/{conversational => chat}/run_chat_agent.ts | 0 .../services/agents/{conversational => chat}/system_prompt.ts | 0 .../utils/from_langchain_messages.test.ts | 0 .../{conversational => chat}/utils/from_langchain_messages.ts | 0 .../services/agents/{conversational => chat}/utils/index.ts | 0 .../utils/to_langchain_messages.test.ts | 0 .../{conversational => chat}/utils/to_langchain_messages.ts | 0 .../utils/tool_provider_to_langchain_tools.ts | 0 .../server/services/agents/{research => researcher}/graph.ts | 2 +- .../server/services/agents/{research => researcher}/index.ts | 0 .../server/services/agents/{research => researcher}/prompts.ts | 0 .../agents/{research => researcher}/researcher_as_tool.ts | 3 +-- .../agents/{research => researcher}/run_researcher_agent.ts | 0 .../server/services/agents/{research => researcher}/utils.ts | 2 +- .../plugins/shared/onechat/server/tools/register_tools.ts | 2 +- 22 files changed, 6 insertions(+), 6 deletions(-) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/convert_graph_events.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/graph.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/handler.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/index.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/provider.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/run_chat_agent.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/system_prompt.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/utils/from_langchain_messages.test.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/utils/from_langchain_messages.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/utils/index.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/utils/to_langchain_messages.test.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/utils/to_langchain_messages.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{conversational => chat}/utils/tool_provider_to_langchain_tools.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{research => researcher}/graph.ts (98%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{research => researcher}/index.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{research => researcher}/prompts.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{research => researcher}/researcher_as_tool.ts (93%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{research => researcher}/run_researcher_agent.ts (100%) rename x-pack/platform/plugins/shared/onechat/server/services/agents/{research => researcher}/utils.ts (89%) diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts index b6ff77b126d61..b1fa476507895 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.ts @@ -17,6 +17,7 @@ export const BuiltinToolIds = { getDocumentById: 'get_document_by_id', generateEsql: 'generate_esql', executeEsql: 'execute_esql', + researcherAgent: 'researcher_agent', }; /** diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/agents_service.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/agents_service.ts index f0456e1a52b17..54e5cd2835ac9 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/agents_service.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/agents_service.ts @@ -9,7 +9,7 @@ import type { Logger } from '@kbn/logging'; import type { Runner } from '@kbn/onechat-server'; import type { AgentsServiceSetup, AgentsServiceStart } from './types'; import { createInternalRegistry } from './utils'; -import { createDefaultAgentProvider } from './conversational'; +import { createDefaultAgentProvider } from './chat'; export interface AgentsServiceSetupDeps { logger: Logger; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/convert_graph_events.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/convert_graph_events.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/convert_graph_events.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/convert_graph_events.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/graph.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/graph.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/graph.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/handler.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/handler.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/handler.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/handler.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/index.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/index.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/index.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/index.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/provider.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/provider.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/provider.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/provider.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/run_chat_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/run_chat_agent.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/system_prompt.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/system_prompt.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/from_langchain_messages.test.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/from_langchain_messages.test.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/from_langchain_messages.test.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/from_langchain_messages.test.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/from_langchain_messages.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/from_langchain_messages.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/from_langchain_messages.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/from_langchain_messages.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/index.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/index.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/index.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/index.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/to_langchain_messages.test.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.test.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/to_langchain_messages.test.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.test.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/to_langchain_messages.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/to_langchain_messages.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/tool_provider_to_langchain_tools.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/tool_provider_to_langchain_tools.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/tool_provider_to_langchain_tools.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/tool_provider_to_langchain_tools.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts similarity index 98% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts index ac3f81d4dd0e3..344ca0b7767fc 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/graph.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -14,7 +14,7 @@ import type { Logger } from '@kbn/core/server'; import { InferenceChatModel } from '@kbn/inference-langchain'; import { getReflectionPrompt, getExecutionPrompt, getAnswerPrompt } from './prompts'; import { extractToolResults } from './utils'; -import { getToolCalls, extractTextContent } from '../conversational/utils/from_langchain_messages'; +import { getToolCalls, extractTextContent } from '../chat/utils/from_langchain_messages'; // tool_choice: toolName diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/index.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/index.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/research/index.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/index.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/prompts.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/research/prompts.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts similarity index 93% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts index be1863f8e7e2e..736af35ae04c2 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/researcher_as_tool.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts @@ -8,7 +8,6 @@ import { z } from '@kbn/zod'; import type { RegisteredTool } from '@kbn/onechat-server'; import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; -import { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { runSearchAgent } from './run_researcher_agent'; const researcherSchema = z.object({ @@ -21,7 +20,7 @@ export interface ResearcherResponse { export const researcherTool = (): RegisteredTool => { return { - id: 'researcher', + id: BuiltinToolIds.researcherAgent, description: 'An agentic researcher agent to perform search tasks', schema: researcherSchema, handler: async ({ instructions }, { toolProvider, request, modelProvider, runner, logger }) => { diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts similarity index 100% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/research/run_researcher_agent.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/utils.ts similarity index 89% rename from x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts rename to x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/utils.ts index 56d69e495d186..9439dffdbb1d7 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/research/utils.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/utils.ts @@ -6,7 +6,7 @@ */ import { BaseMessage, isToolMessage } from '@langchain/core/messages'; -import { extractTextContent } from '../conversational/utils/from_langchain_messages'; +import { extractTextContent } from '../chat/utils/from_langchain_messages'; interface ToolResult { toolCallId: string; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts index 41dded5df5c94..b5fb64539cf65 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts @@ -17,7 +17,7 @@ import { listIndicesTool, indexExplorerTool, } from './retrieval'; -import { researcherTool } from '../services/agents/research'; +import { researcherTool } from '../services/agents/researcher'; export const registerTools = ({ tools: registry }: { tools: ToolsServiceSetup }) => { const tools: Array> = [ From 55b66d8be56bc91376dab9b58e7ce632053c35c2 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Mon, 16 Jun 2025 15:10:14 +0200 Subject: [PATCH 09/17] start moving back to the agent --- .../framework/compose_provider.ts | 50 +++++++++++++++++++ .../onechat-genai-utils/framework/index.ts | 13 +++++ .../onechat-genai-utils/langchain/index.ts | 6 +++ .../tools/index_explorer.ts | 2 +- .../services/agents/chat/run_chat_agent.ts | 4 +- .../services/agents/chat/utils/index.ts | 2 +- .../chat/utils/to_langchain_messages.test.ts | 10 ++-- .../chat/utils/to_langchain_messages.ts | 2 +- .../services/agents/researcher/graph.ts | 4 +- .../agents/researcher/run_researcher_agent.ts | 48 ++++++++++++------ .../services/chat/utils/generate_title.ts | 4 +- .../server/tools/retrieval/generate_esql.ts | 5 +- 12 files changed, 120 insertions(+), 30 deletions(-) create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/compose_provider.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/index.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/compose_provider.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/compose_provider.ts new file mode 100644 index 0000000000000..7ee70b3254380 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/compose_provider.ts @@ -0,0 +1,50 @@ +/* + * 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 type { KibanaRequest } from '@kbn/core-http-server'; +import type { PlainIdToolIdentifier, ToolProviderId, ToolDescriptor } from '@kbn/onechat-common'; +import type { ToolProvider, ExecutableTool } from '@kbn/onechat-server'; + +export interface ByToolIdRule { + type: 'by_tool_id'; + providerId: ToolProviderId; + toolIds: PlainIdToolIdentifier[]; +} + +export interface ByProviderIdRule { + type: 'by_provider_id'; + providerId: ToolProviderId; +} + +export type ToolFilterRule = ByToolIdRule | ByProviderIdRule; + +const matches = (rule: ToolFilterRule, tool: ToolDescriptor): boolean => { + if (rule.type === 'by_tool_id') { + return tool.meta.providerId === rule.providerId && rule.toolIds.includes(tool.id); + } else if (rule.type === 'by_provider_id') { + return tool.meta.providerId === rule.providerId; + } else { + throw new Error('Unknown rule type'); + } +}; + +const anyMatch = (rules: ToolFilterRule[], tool: ToolDescriptor): boolean => { + return rules.some((rule) => matches(rule, tool)); +}; + +export const filterProviderTools = async ({ + provider, + rules, + request, +}: { + provider: ToolProvider; + rules: ToolFilterRule[]; + request: KibanaRequest; +}): Promise => { + const tools = await provider.list({ request }); + return tools.filter((tool) => anyMatch(rules, tool)); +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/index.ts new file mode 100644 index 0000000000000..2fbd383a04100 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/framework/index.ts @@ -0,0 +1,13 @@ +/* + * 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. + */ + +export { + filterProviderTools, + type ByProviderIdRule, + type ToolFilterRule, + type ByToolIdRule, +} from './compose_provider'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts new file mode 100644 index 0000000000000..1fec1c76430eb --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts @@ -0,0 +1,6 @@ +/* + * 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. + */ diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts index a1715315b5308..c7e846f536c9e 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/index_explorer.ts @@ -10,7 +10,7 @@ import { z } from '@kbn/zod'; import type { MappingTypeMapping } from '@elastic/elasticsearch/lib/api/types'; import type { ScopedModel } from '@kbn/onechat-server'; import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; -import { ListIndexInfo, listIndices } from './utils/listings'; +import { ListIndexInfo, listIndices } from './steps/list_indices'; import { getIndexMappings } from './steps/get_mappings'; export interface RelevantIndex { diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts index 4032b6754b979..dd3e40ec92ab9 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts @@ -22,7 +22,7 @@ import type { ExecutableTool, ToolProvider, } from '@kbn/onechat-server'; -import { providerToLangchainTools, toLangchainTool, conversationLangchainMessages } from './utils'; +import { providerToLangchainTools, toLangchainTool, conversationToLangchainMessages } from './utils'; import { createAgentGraph } from './graph'; import { convertGraphEvents, addRoundCompleteEvent } from './convert_graph_events'; @@ -93,7 +93,7 @@ export const runChatAgent: RunChatAgentFn = async ( const langchainTools = Array.isArray(tools) ? tools.map((tool) => toLangchainTool({ tool, logger })) : await providerToLangchainTools({ request, toolProvider: tools, logger }); - const initialMessages = conversationLangchainMessages({ + const initialMessages = conversationToLangchainMessages({ nextInput, previousRounds: conversation, }); diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/index.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/index.ts index 7b3cba794524e..a5ec190b06d6a 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/index.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/index.ts @@ -5,6 +5,6 @@ * 2.0. */ -export { conversationLangchainMessages } from './to_langchain_messages'; +export { conversationToLangchainMessages } from './to_langchain_messages'; export { toLangchainTool, providerToLangchainTools } from './tool_provider_to_langchain_tools'; export { extractTextContent } from './from_langchain_messages'; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.test.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.test.ts index 9dbba5850d3cf..7282c341e6877 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.test.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.test.ts @@ -7,7 +7,7 @@ import { isHumanMessage, isAIMessage, AIMessage, ToolMessage } from '@langchain/core/messages'; import { ToolCallWithResult, ToolCallStep, ConversationRoundStepType } from '@kbn/onechat-common'; -import { conversationLangchainMessages } from './to_langchain_messages'; +import { conversationToLangchainMessages } from './to_langchain_messages'; describe('conversationLangchainMessages', () => { const makeRoundInput = (message: string) => ({ message }); @@ -30,7 +30,7 @@ describe('conversationLangchainMessages', () => { it('returns only the user message if no previous rounds', () => { const nextInput = makeRoundInput('hello'); - const result = conversationLangchainMessages({ previousRounds: [], nextInput }); + const result = conversationToLangchainMessages({ previousRounds: [], nextInput }); expect(result).toHaveLength(1); expect(isHumanMessage(result[0])).toBe(true); expect(result[0].content).toBe('hello'); @@ -45,7 +45,7 @@ describe('conversationLangchainMessages', () => { }, ]; const nextInput = makeRoundInput('how are you?'); - const result = conversationLangchainMessages({ previousRounds, nextInput }); + const result = conversationToLangchainMessages({ previousRounds, nextInput }); expect(result).toHaveLength(3); @@ -68,7 +68,7 @@ describe('conversationLangchainMessages', () => { }, ]; const nextInput = makeRoundInput('next'); - const result = conversationLangchainMessages({ previousRounds, nextInput }); + const result = conversationToLangchainMessages({ previousRounds, nextInput }); // 1 user + 1 tool call (AI + Tool) + 1 assistant + 1 user expect(result).toHaveLength(5); const [ @@ -106,7 +106,7 @@ describe('conversationLangchainMessages', () => { }, ]; const nextInput = makeRoundInput('bye'); - const result = conversationLangchainMessages({ previousRounds, nextInput }); + const result = conversationToLangchainMessages({ previousRounds, nextInput }); // 1 user + 1 assistant + 1 user + 1 tool call (AI + Tool) + 1 assistant + 1 user expect(result).toHaveLength(7); const [ diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts index 7317f1af5e748..ae80b9d372dab 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts @@ -17,7 +17,7 @@ import { toolIdToLangchain } from './tool_provider_to_langchain_tools'; /** * Converts a conversation to langchain format */ -export const conversationLangchainMessages = ({ +export const conversationToLangchainMessages = ({ previousRounds, nextInput, }: { diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts index 344ca0b7767fc..35150b42be6fc 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -35,9 +35,9 @@ export interface ExecutedAction { // tools: // - index explorer -// - fulltext search +// - relevance search +// - nl search // - get_document_by_id -// - ES|QL? interface ReflectionResult { isSufficient: boolean; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts index 9cea9d51d1522..10ccdd1ff3959 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -10,15 +10,21 @@ import { Observable, from, filter, shareReplay, firstValueFrom, map, lastValueFr import type { Logger } from '@kbn/logging'; import { StreamEvent } from '@langchain/core/tracers/log_stream'; import type { KibanaRequest } from '@kbn/core-http-server'; -import { ChatAgentEvent, isRoundCompleteEvent } from '@kbn/onechat-common'; +import { + ChatAgentEvent, + isRoundCompleteEvent, + BuiltinToolIds, + builtinToolProviderId, +} from '@kbn/onechat-common'; import type { ModelProvider, ScopedRunner, ToolProvider } from '@kbn/onechat-server'; +import { filterProviderTools } from '@kbn/onechat-genai-utils/framework'; import { providerToLangchainTools, toLangchainTool, - conversationLangchainMessages, -} from '../conversational/utils'; + conversationToLangchainMessages, +} from '../chat/utils'; import { createAgentGraph } from './graph'; -import { convertGraphEvents, addRoundCompleteEvent } from '../conversational/convert_graph_events'; +import { convertGraphEvents, addRoundCompleteEvent } from '../chat/convert_graph_events'; export interface RunSearchAgentContext { logger: Logger; @@ -63,7 +69,26 @@ export const runSearchAgent: RunChatAgentFn = async ( { logger, request, modelProvider } ) => { const model = await modelProvider.getDefaultModel(); - const langchainTools = await providerToLangchainTools({ request, toolProvider, logger }); + + const researcherTools = await filterProviderTools({ + request, + provider: toolProvider, + rules: [ + { + type: 'by_tool_id', + providerId: builtinToolProviderId, + toolIds: [ + BuiltinToolIds.relevanceSearch, + BuiltinToolIds.naturalLanguageSearch, + BuiltinToolIds.indexExplorer, + BuiltinToolIds.getDocumentById, + ], + }, + ], + }); + + const langchainTools = researcherTools.map((tool) => toLangchainTool({ tool, logger })); + const agentGraph = await createAgentGraph({ logger, chatModel: model.chatModel, @@ -92,7 +117,9 @@ export const runSearchAgent: RunChatAgentFn = async ( shareReplay() ); - events$.subscribe(onEvent); + events$.subscribe((event) => { + // later we should emit reasoning events from there. + }); await lastValueFrom(events$); @@ -102,15 +129,6 @@ export const runSearchAgent: RunChatAgentFn = async ( }; }; -export const extractRound = async (events$: Observable) => { - return await firstValueFrom( - events$.pipe( - filter(isRoundCompleteEvent), - map((event) => event.data.round) - ) - ); -}; - const isStreamEvent = (input: any): input is StreamEvent => { return 'event' in input; }; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/chat/utils/generate_title.ts b/x-pack/platform/plugins/shared/onechat/server/services/chat/utils/generate_title.ts index 5753d6e2a0125..5a407b7331244 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/chat/utils/generate_title.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/chat/utils/generate_title.ts @@ -9,7 +9,7 @@ import { z } from '@kbn/zod'; import { BaseMessageLike } from '@langchain/core/messages'; import type { InferenceChatModel } from '@kbn/inference-langchain'; import type { ConversationRound, RoundInput } from '@kbn/onechat-common'; -import { conversationLangchainMessages } from '../../agents/conversational/utils'; +import { conversationToLangchainMessages } from '../../agents/chat/utils'; export const generateConversationTitle = async ({ previousRounds, @@ -31,7 +31,7 @@ export const generateConversationTitle = async ({ 'system', "'You are a helpful assistant. Assume the following messages is the start of a conversation between you and a user; give this conversation a title based on the content below", ], - ...conversationLangchainMessages({ previousRounds, nextInput }), + ...conversationToLangchainMessages({ previousRounds, nextInput }), ]; const { title } = await structuredModel.invoke(prompt); diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts index 9200b2f25b47b..4fd9d32eeeff3 100644 --- a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.ts @@ -24,7 +24,10 @@ const nlToEsqlToolSchema = z.object({ .describe('(optional) Additional context that could be useful to generate the ES|QL query'), }); -export const generateEsqlTool = (): RegisteredTool => { +export const generateEsqlTool = (): RegisteredTool< + typeof nlToEsqlToolSchema, + GenerateEsqlResponse +> => { return { id: BuiltinToolIds.generateEsql, description: 'Generate an ES|QL query from a natural language query.', From 71c68fcb0dac2b8711a865dadb3690b799fc9dca Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Tue, 17 Jun 2025 09:10:12 +0200 Subject: [PATCH 10/17] better action log representation + improved prompts --- .../services/agents/researcher/backlog.ts | 39 ++++ .../services/agents/researcher/graph.ts | 117 +++++----- .../services/agents/researcher/prompts.ts | 211 +++++++++++++----- .../agents/researcher/researcher_as_tool.ts | 19 +- .../agents/researcher/run_researcher_agent.ts | 15 +- 5 files changed, 285 insertions(+), 116 deletions(-) create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts new file mode 100644 index 0000000000000..a4cf4a13df4d7 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts @@ -0,0 +1,39 @@ +/* + * 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. + */ + +export interface ActionResult { + researchGoal: string; + toolName: string; + arguments: any; + response: any; +} + +export interface ReflectionResult { + isSufficient: boolean; + nextQuestions: string[]; + reasoning?: string; +} + +export type BacklogItem = ActionResult | ReflectionResult; + +export const isReflectionResult = (item: BacklogItem): item is ReflectionResult => { + return 'isSufficient' in item; +}; + +export const isActionResult = (item: BacklogItem): item is ActionResult => { + return 'toolName' in item; +}; + +export const lastReflectionResult = (backlog: BacklogItem[]): ReflectionResult => { + for (let i = backlog.length - 1; i >= 0; i--) { + const current = backlog[i]; + if (isReflectionResult(current)) { + return current; + } + } + throw new Error('No reflection result found'); +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts index 35150b42be6fc..1c9003a8fcae6 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -12,47 +12,21 @@ import { ToolNode } from '@langchain/langgraph/prebuilt'; import type { StructuredTool } from '@langchain/core/tools'; import type { Logger } from '@kbn/core/server'; import { InferenceChatModel } from '@kbn/inference-langchain'; +import { getToolCalls, extractTextContent } from '../chat/utils/from_langchain_messages'; import { getReflectionPrompt, getExecutionPrompt, getAnswerPrompt } from './prompts'; import { extractToolResults } from './utils'; -import { getToolCalls, extractTextContent } from '../chat/utils/from_langchain_messages'; - -// tool_choice: toolName - -export interface PlannedAction { - knowledgeGap: string; -} - -export interface ExecutedAction { - knowledgeGap: string; - toolName: string; - arguments: any; - response: any; -} - -// -// process queue -> create knowledge entries -> reason -// - -// tools: -// - index explorer -// - relevance search -// - nl search -// - get_document_by_id +import { ActionResult, ReflectionResult, BacklogItem, lastReflectionResult } from './backlog'; -interface ReflectionResult { - isSufficient: boolean; - knowledgeGaps: string[]; - reasoning?: string; +export interface ResearchGoal { + question: string; } export const createAgentGraph = async ({ chatModel, tools, - systemPrompt, }: { chatModel: InferenceChatModel; tools: StructuredTool[]; - systemPrompt?: string; logger: Logger; }) => { const StateAnnotation = Annotation.Root({ @@ -60,19 +34,18 @@ export const createAgentGraph = async ({ initialQuery: Annotation(), // the search query cycleBudget: Annotation(), // budget in number of cycles - TODO // internal state - actionsQueue: Annotation({ + actionsQueue: Annotation({ reducer: (state, actions) => { return actions ?? state; }, default: () => [], }), - processedActions: Annotation({ + backlog: Annotation({ reducer: (current, next) => { return [...current, ...next]; }, default: () => [], }), - lastReflectionResult: Annotation(), // outputs generatedAnswer: Annotation(), }); @@ -81,8 +54,8 @@ export const createAgentGraph = async ({ * Initialize the flow by adding a first index explorer call to the action queue. */ const initialize = async (state: typeof StateAnnotation.State) => { - const firstAction: PlannedAction = { - knowledgeGap: state.initialQuery, + const firstAction: ResearchGoal = { + question: state.initialQuery, }; return { actionsQueue: [firstAction], @@ -99,8 +72,8 @@ export const createAgentGraph = async ({ const response = await executionModel.invoke( getExecutionPrompt({ - nextAction: nextItem, - executedActions: state.processedActions, + currentResearchGoal: nextItem, + backlog: state.backlog, }) ); const toolCalls = getToolCalls(response); @@ -110,17 +83,24 @@ export const createAgentGraph = async ({ const toolMessages = await toolNode.invoke([response]); const toolResults = extractToolResults(toolMessages); - const processedActions: ExecutedAction[] = []; - processedActions.push({ - ...nextItem, - toolName: toolCalls[0].toolId.toolId, - arguments: toolCalls[0].args, - response: toolResults[0].result, - }); + const actionResults: ActionResult[] = []; + for (let i = 0; i < toolResults.length; i++) { + const toolCall = toolCalls[i]; + const toolResult = toolResults[i]; + if (toolCall && toolResult) { + const actionResult: ActionResult = { + researchGoal: nextItem.question, + toolName: toolCall.toolId.toolId, + arguments: toolCall.args, + response: toolResult.result, + }; + actionResults.push(actionResult); + } + } return { actionsQueue: queue, - processedActions: [processedActions], + backlog: [actionResults], }; }; @@ -135,17 +115,21 @@ export const createAgentGraph = async ({ const reflection = async (state: typeof StateAnnotation.State) => { const reflectModel = chatModel.withStructuredOutput( z.object({ - isSufficient: z - .boolean() - .describe('Whether the provided info are sufficient to answer the user question'), - knowledgeGaps: z - .array(z.string()) - .describe('A description of what information is missing or needs clarification'), + isSufficient: z.boolean().describe( + `Set to true if the current information fully answers the user question without requiring further research. + Set to false if any knowledge gaps or unresolved sub-problems remain.` + ), + nextQuestions: z.array(z.string()).describe( + `A list of self-contained, actionable research questions or sub-problems that need to be explored + further to fully answer the user question. Leave empty if isSufficient is true.` + ), reasoning: z .string() .optional() .describe( - 'Optional reasoning on why the provided info are sufficient or not. Can be used as scratch pad for thoughts.' + `Brief internal reasoning explaining why the current information is sufficient or not. + You may list what was already answered, what gaps exist, or whether decomposition was necessary. + Use this as your thought process or scratchpad before producing the final output.` ), }) ); @@ -153,36 +137,51 @@ export const createAgentGraph = async ({ const response: ReflectionResult = await reflectModel.invoke( getReflectionPrompt({ userQuery: state.initialQuery, - summaries: state.processedActions, + backlog: state.backlog, }) ); console.log('*** reflection response: ', response); return { - lastReflectionResult: response, + backlog: [...state.backlog, response], actionsQueue: [ ...state.actionsQueue, - response.knowledgeGaps.map((gap) => ({ knowledgeGap: gap })), + response.nextQuestions.map((nextQuestion) => ({ question: nextQuestion })), ], }; }; const evaluateReflection = async (state: typeof StateAnnotation.State) => { - if (state.lastReflectionResult.isSufficient) { + const reflectionResult = lastReflectionResult(state.backlog); + console.log( + '*** evaluateReflection - state: ', + reflectionResult.isSufficient, + reflectionResult.reasoning + ); + + if (reflectionResult.isSufficient) { return 'answer'; } return 'process_queue_item'; }; const answer = async (state: typeof StateAnnotation.State) => { - const response = await chatModel.invoke( + console.log('*** answer - start'); + + const answerModel = chatModel.withConfig({ + tags: ['researcher-answer'], + }); + + const response = await answerModel.invoke( getAnswerPrompt({ userQuery: state.initialQuery, - executedActions: state.processedActions, + backlog: state.backlog, }) ); + console.log('*** answer - raw response: ', response); + const generatedAnswer = extractTextContent(response); console.log('*** answer - response: ', generatedAnswer); @@ -194,10 +193,12 @@ export const createAgentGraph = async ({ // note: the node names are used in the event convertion logic, they should *not* be changed const graph = new StateGraph(StateAnnotation) + // nodes .addNode('initialize', initialize) .addNode('process_queue_item', processQueueItem) .addNode('reflection', reflection) .addNode('answer', answer) + // edges .addEdge('__start__', 'initialize') .addEdge('initialize', 'process_queue_item') .addConditionalEdges('process_queue_item', evaluateQueue, { diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts index 8cddecb20598f..2afe07f585efb 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts @@ -6,43 +6,115 @@ */ import type { BaseMessageLike } from '@langchain/core/messages'; -import type { PlannedAction, ExecutedAction } from './graph'; +import type { ResearchGoal } from './graph'; +import { + isActionResult, + isReflectionResult, + BacklogItem, + ReflectionResult, + ActionResult, +} from './backlog'; + +const renderBacklog = (backlog: BacklogItem[]): string => { + const renderItem = (item: BacklogItem, i: number) => { + if (isActionResult(item)) { + return renderActionResult(item, i); + } + if (isReflectionResult(item)) { + return renderReflectionResult(item, i); + } + return `Unknown item type`; + }; + + return backlog.map((item, i) => renderItem(item, i)).join('\n\n'); +}; + +const renderReflectionResult = ( + { isSufficient, nextQuestions, reasoning }: ReflectionResult, + index: number +): string => { + return `### Cycle ${index + 1} + + At cycle "${index + 1}", you reflected on the data gathered so far: + + - You decided that the current information were ${ + isSufficient ? '*sufficient*' : '*insufficient*' + } to fully answer the question, with the following reasoning: ${reasoning} + + ${ + nextQuestions.length > 0 + ? `- You identified the following questions to follow up on: +${nextQuestions.map((question) => ` - ${question}`).join('\n')}` + : '' + } + `; +}; + +const renderActionResult = (actionResult: ActionResult, index: number): string => { + return `### Cycle ${index + 1} + + At cycle "${index + 1}", you performed the following action: + + - Action type: tool execution + + - Tool name: ${actionResult.toolName} + + - Tool parameters: + \`\`\`json + ${JSON.stringify(actionResult.arguments, undefined, 2)} + \`\`\` + + - Tool response: + \`\`\`json + ${JSON.stringify(actionResult.response, undefined, 2)} + \`\`\` + `; +}; export const getExecutionPrompt = ({ - nextAction, - executedActions, + currentResearchGoal, + backlog, }: { - nextAction: PlannedAction; - executedActions: ExecutedAction[]; + currentResearchGoal: ResearchGoal; + backlog: BacklogItem[]; }): BaseMessageLike[] => { return [ [ 'system', - `You are an expert research assistant from the Elasticsearch company. + `You are a research agent at Elasticsearch with access to external tools. + + Your task: + - Based on a research goal, choose the most appropriate tool to help resolve it. + - You will also be provided with a list of past actions and results. Instructions: - You will be with a goal, and a list of already executed actions. With those information, - please choose which tool to call to reach the goal. + - You must select one tool and invoke it with the most relevant and precise parameters. + - Choose the tool that will best help fulfill the current research goal. + - Some tools (e.g., search) may require contextual information (such as an index name or prior step result). Retrieve it from the action history if needed. + - Do not repeat a tool invocation that has already been attempted with the same or equivalent parameters. + - Think carefully about what the goal requires and which tool best advances it. - Requirements: - - You *must* call a tool - - Be attentive, as some tools may require information from the previously executed action. For - example, search tools usually require to target an index, which can be retrieved using the index explorer tool. - - Be careful to not call the same tool twice with the same parameters. Check the action history. + Constraints: + - Tool use is mandatory. You must respond with a tool call. + - Do not speculate or summarize. Only act by selecting the best next tool and invoking it. - Action history and current goal will be provided in the next user message. + Output format: + Respond using the tool-calling schema provided by the system. + + Additional information: + - The current date is ${new Date().toISOString()}. `, ], [ 'user', ` - ### Current goal: + ### Current Research Goal - Trying to find information about: "${nextAction.knowledgeGap}" + Trying to find information about: "${currentResearchGoal.question}" - ### Action history: + ### Previous Actions - ${executedActions.map((action) => JSON.stringify(action, undefined, 2)).join('\n')} + ${renderBacklog(backlog)} `, ], ]; @@ -50,45 +122,81 @@ export const getExecutionPrompt = ({ export const getReflectionPrompt = ({ userQuery, - summaries, + backlog, }: { userQuery: string; - summaries: any[]; + backlog: BacklogItem[]; }): BaseMessageLike[] => { return [ [ 'system', - `You are an expert research assistant from the Elasticsearch company analyzing summaries about "${userQuery}". + `You are an expert research assistant from the Elasticsearch company analyzing information about the user's question: "${userQuery}". Instructions: - Identify knowledge gaps or areas that need deeper exploration and generate the corresponding follow-up queries. (1 or multiple). - - If provided summaries are sufficient to answer the user's question, don't generate a follow-up query. - - Focus on technical details, implementation specifics, or emerging trends that weren't fully covered. + - Analyze the completeness and depth of the provided summaries. + - Identify any missing, unclear, or shallow information. + - If necessary, break down complex questions into smaller sub-problems. + - Your goal is to generate a precise list of actionable questions that will help drive the research forward. - Requirements: - - Ensure the follow-up query is self-contained and includes necessary context for web search. + Guidelines: + - Only generate questions if the current information is incomplete or insufficient. + - Focus on technical depth, implementation details, trade-offs, edge cases, or emerging trends. + - Each question must be self-contained and ready to be used for search or further investigation. + + Additional information: + - The current date is ${new Date().toISOString()}. Output Format: - Format your response as a JSON object with these exact keys: - - "is_sufficient": true or false - - "knowledge_gap": Describe what information is missing or needs clarification - - "follow_up_queries": Write a specific question to address this gap + - "isSufficient": true or false + - "nextQuestions": list of standalone research questions (empty if isSufficient is true) + - "reasoning": internal reasoning (brief thought process for your analysis) - Example 1: if information are sufficient: + ### Example 1: information is sufficient \`\`\`json - { + { "isSufficient": true, - "knowledgeGaps": [], - } + "nextQuestions": [], + "reasoning": "The provided summaries fully explain how Elasticsearch handles vector search, including indexing, retrieval, and trade-offs." + } + \`\`\` + + ### Example 2: minor gaps or missing details + \`\`\`json + { + "isSufficient": false, + "nextQuestions": [ + "How does Elasticsearch query performance scale with large document sizes?", + "What is the default scoring mechanism used in Elasticsearch for dense vector fields?" + ], + "reasoning": "While the summaries explain vector search basics, they lack detail on scaling performance and scoring behavior." + } + \`\`\` + + ### Example 3: complex decomposition + \`\`\`json + { + "isSufficient": false, + "nextQuestions": [ + "What is the architecture of Elasticsearch when used as a retrieval component in RAG pipelines with LLMs?", + "How does hybrid search compare to dense retrieval in Elasticsearch in terms of accuracy and recall?", + "What are the performance and cost trade-offs between using vector search and keyword-based search in Elasticsearch?" + ], + "reasoning": "The summaries cover general Elasticsearch features but miss details about RAG architectures, hybrid retrieval comparisons, and performance trade-offs." + } \`\`\` `, ], [ 'user', ` - ### Summaries: + ## User question - ${summaries.map((summary) => JSON.stringify(summary, undefined, 2)).join('\n')} + "${userQuery}" + + ## Backlog + + ${renderBacklog(backlog)} `, ], ]; @@ -96,28 +204,33 @@ export const getReflectionPrompt = ({ export const getAnswerPrompt = ({ userQuery, - executedActions, + backlog, }: { userQuery: string; - executedActions: ExecutedAction[]; + backlog: BacklogItem[]; }): BaseMessageLike[] => { return [ [ 'system', - `Generate a high-quality answer to the user's question based on the provided summaries. + `You are a senior technical expert from the Elasticsearch company. + Your role is to provide a clear, well-reasoned answer to the user's question using the information gathered by prior research steps. Instructions: + - Carefully read the user's original question and the gathered information. + - Synthesize an accurate response that directly answers the user's question. + - Do not hedge. If the information is complete, provide a confident and final answer. + - If there are still uncertainties or unresolved issues, acknowledge them clearly and state what is known and what is not. + - Prefer structured, organized output (e.g., use paragraphs, bullet points, or sections if helpful). + + Guidelines: + - Do not mention the research process or that you are an AI or assistant. + - Do not mention that the answer was generated based on previous steps. + - Do not repeat the user's question or summarize the JSON input. + - Do not speculate beyond the gathered information unless logically inferred from it. + + Additional information: - The current date is ${new Date().toISOString()}. - - You are the final step of a multi-step research process, don't mention that you are the final step. - - You have access to all the information gathered from the previous steps. - - You have access to the user's question. - - Generate a high-quality answer to the user's question based on the provided summaries and the user's question. - User Context: - - {research_topic} - - Summaries: - {summaries} `, ], [ @@ -129,9 +242,7 @@ export const getAnswerPrompt = ({ ### Gathered information - \`\`\`json - ${JSON.stringify(executedActions, undefined, 2)} - \`\`\` + ${renderBacklog(backlog.filter(isActionResult))} `, ], ]; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts index 736af35ae04c2..22da4714df90a 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts @@ -11,7 +11,7 @@ import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; import { runSearchAgent } from './run_researcher_agent'; const researcherSchema = z.object({ - instructions: z.string().describe('Instructions for the researcher'), + instructions: z.string().describe('Research instructions for the agent'), }); export interface ResearcherResponse { @@ -21,7 +21,22 @@ export interface ResearcherResponse { export const researcherTool = (): RegisteredTool => { return { id: BuiltinToolIds.researcherAgent, - description: 'An agentic researcher agent to perform search tasks', + description: `An agentic researcher tool to perform search and analysis tasks. + + Can be used to perform "deep search" tasks where a single query or search is not enough + and where we need some kind of more in depth-research with multiple search requests and analysis. + + Example where the agent should be used: + - "Summarize the changes between our previous and current work from home policy" + - "Find the vulnerabilities involved in our latest alerts and gather information about them" + - Any time the user explicitly asks to use this tool + + Example where the agent should not be used (in favor of more simple search tools): + - "Show me the last 5 documents in the index 'foo'" + - "Show me my latest alerts" + + Notes: + - Please include all useful information in the instructions, as the agent has no other context. `, schema: researcherSchema, handler: async ({ instructions }, { toolProvider, request, modelProvider, runner, logger }) => { const searchAgentResult = await runSearchAgent( diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts index 10ccdd1ff3959..8fc539fec35d1 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -97,15 +97,17 @@ export const runSearchAgent: RunChatAgentFn = async ( }); const eventStream = agentGraph.streamEvents( - { initialQuery: instructions }, + { + initialQuery: instructions, + cycleBudget: 10, + }, { version: 'v2', runName: agentGraphName, metadata: { graphName: agentGraphName, - // runId, }, - recursionLimit: 10, + recursionLimit: 30, callbacks: [], } ); @@ -121,11 +123,12 @@ export const runSearchAgent: RunChatAgentFn = async ( // later we should emit reasoning events from there. }); - await lastValueFrom(events$); + // event: 'on_chain_end', name: 'researcher-agent' + const lastEvent = await lastValueFrom(events$); + const generatedAnswer = lastEvent.data.output.generatedAnswer; - // return await extractRound(events$); return { - answer: 'hello', + answer: generatedAnswer, }; }; From 06c2ec83e092420af14b17552eb2ce0fbc96a83b Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Tue, 17 Jun 2025 10:47:51 +0200 Subject: [PATCH 11/17] implement cycle budget --- .../services/agents/researcher/graph.ts | 13 ++++++++--- .../services/agents/researcher/prompts.ts | 21 ++++++++++++++++++ .../agents/researcher/run_researcher_agent.ts | 22 ++++++++++++------- 3 files changed, 45 insertions(+), 11 deletions(-) diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts index 1c9003a8fcae6..a29bbad13af2c 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -34,6 +34,7 @@ export const createAgentGraph = async ({ initialQuery: Annotation(), // the search query cycleBudget: Annotation(), // budget in number of cycles - TODO // internal state + remainingCycles: Annotation(), actionsQueue: Annotation({ reducer: (state, actions) => { return actions ?? state; @@ -59,6 +60,7 @@ export const createAgentGraph = async ({ }; return { actionsQueue: [firstAction], + remainingCycles: state.cycleBudget, }; }; @@ -100,7 +102,7 @@ export const createAgentGraph = async ({ return { actionsQueue: queue, - backlog: [actionResults], + backlog: [...actionResults], }; }; @@ -138,21 +140,26 @@ export const createAgentGraph = async ({ getReflectionPrompt({ userQuery: state.initialQuery, backlog: state.backlog, + maxFollowUpQuestions: 3, + remainingCycles: state.remainingCycles - 1, }) ); console.log('*** reflection response: ', response); + console.log('*** reflection remainingCycles: ', state.remainingCycles); return { + remainingCycles: state.remainingCycles - 1, backlog: [...state.backlog, response], actionsQueue: [ ...state.actionsQueue, - response.nextQuestions.map((nextQuestion) => ({ question: nextQuestion })), + ...response.nextQuestions.map((nextQuestion) => ({ question: nextQuestion })), ], }; }; const evaluateReflection = async (state: typeof StateAnnotation.State) => { + const remainingCycles = state.remainingCycles; const reflectionResult = lastReflectionResult(state.backlog); console.log( '*** evaluateReflection - state: ', @@ -160,7 +167,7 @@ export const createAgentGraph = async ({ reflectionResult.reasoning ); - if (reflectionResult.isSufficient) { + if (reflectionResult.isSufficient || remainingCycles <= 0) { return 'answer'; } return 'process_queue_item'; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts index 2afe07f585efb..5e375d490380c 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts @@ -123,9 +123,15 @@ export const getExecutionPrompt = ({ export const getReflectionPrompt = ({ userQuery, backlog, + maxFollowUpQuestions = 3, + remainingCycles, + cycleBoundaries = { exploration: 3, refinement: 2, finalization: 1 }, }: { userQuery: string; backlog: BacklogItem[]; + remainingCycles: number; + maxFollowUpQuestions?: number; + cycleBoundaries?: { exploration: number; refinement: number; finalization: number }; }): BaseMessageLike[] => { return [ [ @@ -138,8 +144,23 @@ export const getReflectionPrompt = ({ - If necessary, break down complex questions into smaller sub-problems. - Your goal is to generate a precise list of actionable questions that will help drive the research forward. + Cycle Awareness: + - The research process is bounded. There is exactly **${remainingCycles} cycles remaining** before a final answer must be produced. + - Use the following strategy based on that number: + - If ${cycleBoundaries.exploration} or more cycles remain: + - You may explore deeper subtopics or decompositions. + - Pursue emerging trends, architectural alternatives, or implementation-specific nuances. + - If ${cycleBoundaries.refinement} or more cycles remain: + - Focus on clarifying known gaps or weak spots in the current summaries. + - Prefer precision over breadth. + - If ${cycleBoundaries.finalization} or less cycle remains: + - There is no time for further exploration. + - Surface only essential missing information that would block the final answer. + - Avoid speculative or marginal questions. + Guidelines: - Only generate questions if the current information is incomplete or insufficient. + - Do not generate more than ${maxFollowUpQuestions} actionable questions. - Focus on technical depth, implementation details, trade-offs, edge cases, or emerging trends. - Each question must be self-contained and ready to be used for search or further investigation. diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts index 8fc539fec35d1..9a63bf477e249 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -38,6 +38,12 @@ export interface RunSearchAgentParams { * The search instructions */ instructions: string; + /** + * Budget, in search cycles, to allocate to the researcher. + * Defaults to 5. + */ + cycleBudget?: number; + /** * Top level tool provider to use to retrieve internal tools */ @@ -48,24 +54,25 @@ export interface RunSearchAgentParams { onEvent?: (event: ChatAgentEvent) => void; } -export interface RunSearchAgentResponse { +export interface RunResearcherAgentResponse { answer: string; } -export type RunChatAgentFn = ( +export type RunResearcherAgentFn = ( params: RunSearchAgentParams, context: RunSearchAgentContext -) => Promise; +) => Promise; const agentGraphName = 'researcher-agent'; +const defaultCycleBudget = 5; const noopOnEvent = () => {}; /** * Create the handler function for the default onechat agent. */ -export const runSearchAgent: RunChatAgentFn = async ( - { instructions, toolProvider, onEvent = noopOnEvent }, +export const runSearchAgent: RunResearcherAgentFn = async ( + { instructions, cycleBudget = defaultCycleBudget, toolProvider, onEvent = noopOnEvent }, { logger, request, modelProvider } ) => { const model = await modelProvider.getDefaultModel(); @@ -93,13 +100,12 @@ export const runSearchAgent: RunChatAgentFn = async ( logger, chatModel: model.chatModel, tools: langchainTools, - systemPrompt: '', }); const eventStream = agentGraph.streamEvents( { initialQuery: instructions, - cycleBudget: 10, + cycleBudget, }, { version: 'v2', @@ -107,7 +113,7 @@ export const runSearchAgent: RunChatAgentFn = async ( metadata: { graphName: agentGraphName, }, - recursionLimit: 30, + recursionLimit: cycleBudget * 10, callbacks: [], } ); From 28e6b220602da8761d5835129adf15d633ec7d98 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Tue, 17 Jun 2025 13:04:53 +0200 Subject: [PATCH 12/17] remove console.log & unused imports --- .../services/agents/researcher/graph.ts | 102 ++++++++----- .../services/agents/researcher/prompts.ts | 135 ++++++++++-------- .../agents/researcher/run_researcher_agent.ts | 31 ++-- 3 files changed, 145 insertions(+), 123 deletions(-) diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts index a29bbad13af2c..df8f1604f3141 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -6,7 +6,7 @@ */ import { z } from '@kbn/zod'; -import { StateGraph, Annotation } from '@langchain/langgraph'; +import { StateGraph, Annotation, Send } from '@langchain/langgraph'; import { BaseMessage } from '@langchain/core/messages'; import { ToolNode } from '@langchain/langgraph/prebuilt'; import type { StructuredTool } from '@langchain/core/tools'; @@ -21,9 +21,10 @@ export interface ResearchGoal { question: string; } -export const createAgentGraph = async ({ +export const createResearcherAgentGraph = async ({ chatModel, tools, + logger: log, }: { chatModel: InferenceChatModel; tools: StructuredTool[]; @@ -32,7 +33,7 @@ export const createAgentGraph = async ({ const StateAnnotation = Annotation.Root({ // inputs initialQuery: Annotation(), // the search query - cycleBudget: Annotation(), // budget in number of cycles - TODO + cycleBudget: Annotation(), // budget in number of cycles // internal state remainingCycles: Annotation(), actionsQueue: Annotation({ @@ -41,6 +42,12 @@ export const createAgentGraph = async ({ }, default: () => [], }), + pendingActions: Annotation({ + reducer: (state, actions) => { + return actions === 'clear' ? [] : [...state, ...actions]; + }, + default: () => [], + }), backlog: Annotation({ reducer: (current, next) => { return [...current, ...next]; @@ -51,10 +58,18 @@ export const createAgentGraph = async ({ generatedAnswer: Annotation(), }); + type StateType = typeof StateAnnotation.State; + + type ResearchStepState = StateType & { + researchGoal: ResearchGoal; + }; + + const stringify = (obj: unknown) => JSON.stringify(obj, null, 2); + /** * Initialize the flow by adding a first index explorer call to the action queue. */ - const initialize = async (state: typeof StateAnnotation.State) => { + const initialize = async (state: StateType) => { const firstAction: ResearchGoal = { question: state.initialQuery, }; @@ -64,10 +79,19 @@ export const createAgentGraph = async ({ }; }; - const processQueueItem = async (state: typeof StateAnnotation.State) => { - const [nextItem, ...queue] = state.actionsQueue; + const dispatchActions = async (state: StateType) => { + return state.actionsQueue.map((action) => { + return new Send('perform_search', { + ...state, + researchGoal: action, + } satisfies ResearchStepState); + }); + }; + + const performSearch = async (state: ResearchStepState) => { + const nextItem = state.researchGoal; - console.log('*** processQueueItem - nextItem: ', nextItem); + log.trace(() => `performSearch - nextItem: ${stringify(nextItem)}`); const toolNode = new ToolNode(tools); const executionModel = chatModel.bindTools(tools); @@ -80,7 +104,7 @@ export const createAgentGraph = async ({ ); const toolCalls = getToolCalls(response); - console.log('*** processQueueItem - toolCalls: ', toolCalls); + log.trace(() => `performSearch - toolCalls: ${stringify(toolCalls)}`); const toolMessages = await toolNode.invoke([response]); const toolResults = extractToolResults(toolMessages); @@ -101,20 +125,26 @@ export const createAgentGraph = async ({ } return { - actionsQueue: queue, - backlog: [...actionResults], + pendingActions: [...actionResults], }; }; - const evaluateQueue = async (state: typeof StateAnnotation.State) => { - const { actionsQueue } = state; - if (actionsQueue.length) { - return 'process_queue_item'; - } - return 'reflection'; + const collectResults = async (state: StateType) => { + log.trace( + () => + `collectResults - pending actions: ${stringify( + state.pendingActions.map((action) => action.researchGoal) + )}` + ); + + return { + pendingActions: 'clear', + actionsQueue: [], + backlog: [...state.pendingActions], + }; }; - const reflection = async (state: typeof StateAnnotation.State) => { + const reflection = async (state: StateType) => { const reflectModel = chatModel.withStructuredOutput( z.object({ isSufficient: z.boolean().describe( @@ -145,12 +175,14 @@ export const createAgentGraph = async ({ }) ); - console.log('*** reflection response: ', response); - console.log('*** reflection remainingCycles: ', state.remainingCycles); + log.trace( + () => + `reflection - remaining cycles: ${state.remainingCycles} - response: ${stringify(response)}` + ); return { remainingCycles: state.remainingCycles - 1, - backlog: [...state.backlog, response], + backlog: [response], actionsQueue: [ ...state.actionsQueue, ...response.nextQuestions.map((nextQuestion) => ({ question: nextQuestion })), @@ -158,24 +190,17 @@ export const createAgentGraph = async ({ }; }; - const evaluateReflection = async (state: typeof StateAnnotation.State) => { + const evaluateReflection = async (state: StateType) => { const remainingCycles = state.remainingCycles; const reflectionResult = lastReflectionResult(state.backlog); - console.log( - '*** evaluateReflection - state: ', - reflectionResult.isSufficient, - reflectionResult.reasoning - ); if (reflectionResult.isSufficient || remainingCycles <= 0) { return 'answer'; } - return 'process_queue_item'; + return dispatchActions(state); }; - const answer = async (state: typeof StateAnnotation.State) => { - console.log('*** answer - start'); - + const answer = async (state: StateType) => { const answerModel = chatModel.withConfig({ tags: ['researcher-answer'], }); @@ -187,11 +212,9 @@ export const createAgentGraph = async ({ }) ); - console.log('*** answer - raw response: ', response); - const generatedAnswer = extractTextContent(response); - console.log('*** answer - response: ', generatedAnswer); + log.trace(() => `answer - response ${stringify(generatedAnswer)}`); return { generatedAnswer, @@ -202,18 +225,19 @@ export const createAgentGraph = async ({ const graph = new StateGraph(StateAnnotation) // nodes .addNode('initialize', initialize) - .addNode('process_queue_item', processQueueItem) + .addNode('perform_search', performSearch) + .addNode('collect_results', collectResults) .addNode('reflection', reflection) .addNode('answer', answer) // edges .addEdge('__start__', 'initialize') - .addEdge('initialize', 'process_queue_item') - .addConditionalEdges('process_queue_item', evaluateQueue, { - process_queue_item: 'process_queue_item', - reflection: 'reflection', + .addConditionalEdges('initialize', dispatchActions, { + perform_search: 'perform_search', }) + .addEdge('perform_search', 'collect_results') + .addEdge('collect_results', 'reflection') .addConditionalEdges('reflection', evaluateReflection, { - process_queue_item: 'process_queue_item', + perform_search: 'perform_search', answer: 'answer', }) .addEdge('answer', '__end__') diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts index 5e375d490380c..f6397917464ce 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts @@ -6,6 +6,7 @@ */ import type { BaseMessageLike } from '@langchain/core/messages'; +import { BuiltinToolIds as Tools } from '@kbn/onechat-common'; import type { ResearchGoal } from './graph'; import { isActionResult, @@ -15,62 +16,6 @@ import { ActionResult, } from './backlog'; -const renderBacklog = (backlog: BacklogItem[]): string => { - const renderItem = (item: BacklogItem, i: number) => { - if (isActionResult(item)) { - return renderActionResult(item, i); - } - if (isReflectionResult(item)) { - return renderReflectionResult(item, i); - } - return `Unknown item type`; - }; - - return backlog.map((item, i) => renderItem(item, i)).join('\n\n'); -}; - -const renderReflectionResult = ( - { isSufficient, nextQuestions, reasoning }: ReflectionResult, - index: number -): string => { - return `### Cycle ${index + 1} - - At cycle "${index + 1}", you reflected on the data gathered so far: - - - You decided that the current information were ${ - isSufficient ? '*sufficient*' : '*insufficient*' - } to fully answer the question, with the following reasoning: ${reasoning} - - ${ - nextQuestions.length > 0 - ? `- You identified the following questions to follow up on: -${nextQuestions.map((question) => ` - ${question}`).join('\n')}` - : '' - } - `; -}; - -const renderActionResult = (actionResult: ActionResult, index: number): string => { - return `### Cycle ${index + 1} - - At cycle "${index + 1}", you performed the following action: - - - Action type: tool execution - - - Tool name: ${actionResult.toolName} - - - Tool parameters: - \`\`\`json - ${JSON.stringify(actionResult.arguments, undefined, 2)} - \`\`\` - - - Tool response: - \`\`\`json - ${JSON.stringify(actionResult.response, undefined, 2)} - \`\`\` - `; -}; - export const getExecutionPrompt = ({ currentResearchGoal, backlog, @@ -83,25 +28,35 @@ export const getExecutionPrompt = ({ 'system', `You are a research agent at Elasticsearch with access to external tools. - Your task: + ### Your task - Based on a research goal, choose the most appropriate tool to help resolve it. - You will also be provided with a list of past actions and results. - Instructions: + ### Instructions - You must select one tool and invoke it with the most relevant and precise parameters. - Choose the tool that will best help fulfill the current research goal. - Some tools (e.g., search) may require contextual information (such as an index name or prior step result). Retrieve it from the action history if needed. - Do not repeat a tool invocation that has already been attempted with the same or equivalent parameters. - Think carefully about what the goal requires and which tool best advances it. - Constraints: + ### Constraints - Tool use is mandatory. You must respond with a tool call. - Do not speculate or summarize. Only act by selecting the best next tool and invoking it. - Output format: + ### Tools description + Your two main search tools are "${Tools.relevanceSearch}" and "${Tools.naturalLanguageSearch}" + - When doing fulltext search, prefer the "${ + Tools.relevanceSearch + }" tool as it performs better for plain fulltext searches. + - For more advanced queries (filtering, aggregation, buckets), use the "${ + Tools.naturalLanguageSearch + }" tool. + + + ### Output format Respond using the tool-calling schema provided by the system. - Additional information: + ### Additional information - The current date is ${new Date().toISOString()}. `, ], @@ -139,7 +94,7 @@ export const getReflectionPrompt = ({ `You are an expert research assistant from the Elasticsearch company analyzing information about the user's question: "${userQuery}". Instructions: - - Analyze the completeness and depth of the provided summaries. + - Analyze the completeness and depth of data available in your backlog history. - Identify any missing, unclear, or shallow information. - If necessary, break down complex questions into smaller sub-problems. - Your goal is to generate a precise list of actionable questions that will help drive the research forward. @@ -268,3 +223,59 @@ export const getAnswerPrompt = ({ ], ]; }; + +const renderBacklog = (backlog: BacklogItem[]): string => { + const renderItem = (item: BacklogItem, i: number) => { + if (isActionResult(item)) { + return renderActionResult(item, i); + } + if (isReflectionResult(item)) { + return renderReflectionResult(item, i); + } + return `Unknown item type`; + }; + + return backlog.map((item, i) => renderItem(item, i)).join('\n\n'); +}; + +const renderReflectionResult = ( + { isSufficient, nextQuestions, reasoning }: ReflectionResult, + index: number +): string => { + return `### Cycle ${index + 1} + + At cycle "${index + 1}", you reflected on the data gathered so far: + + - You decided that the current information were ${ + isSufficient ? '*sufficient*' : '*insufficient*' + } to fully answer the question, with the following reasoning: ${reasoning} + + ${ + nextQuestions.length > 0 + ? `- You identified the following questions to follow up on: +${nextQuestions.map((question) => ` - ${question}`).join('\n')}` + : '' + } + `; +}; + +const renderActionResult = (actionResult: ActionResult, index: number): string => { + return `### Cycle ${index + 1} + + At cycle "${index + 1}", you performed the following action: + + - Action type: tool execution + + - Tool name: ${actionResult.toolName} + + - Tool parameters: + \`\`\`json + ${JSON.stringify(actionResult.arguments, undefined, 2)} + \`\`\` + + - Tool response: + \`\`\`json + ${JSON.stringify(actionResult.response, undefined, 2)} + \`\`\` + `; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts index 9a63bf477e249..fefdca6fbcfc1 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -5,26 +5,15 @@ * 2.0. */ -import { v4 as uuidv4 } from 'uuid'; -import { Observable, from, filter, shareReplay, firstValueFrom, map, lastValueFrom } from 'rxjs'; +import { from, filter, shareReplay, lastValueFrom } from 'rxjs'; import type { Logger } from '@kbn/logging'; import { StreamEvent } from '@langchain/core/tracers/log_stream'; import type { KibanaRequest } from '@kbn/core-http-server'; -import { - ChatAgentEvent, - isRoundCompleteEvent, - BuiltinToolIds, - builtinToolProviderId, -} from '@kbn/onechat-common'; +import { ChatAgentEvent, BuiltinToolIds, builtinToolProviderId } from '@kbn/onechat-common'; import type { ModelProvider, ScopedRunner, ToolProvider } from '@kbn/onechat-server'; import { filterProviderTools } from '@kbn/onechat-genai-utils/framework'; -import { - providerToLangchainTools, - toLangchainTool, - conversationToLangchainMessages, -} from '../chat/utils'; -import { createAgentGraph } from './graph'; -import { convertGraphEvents, addRoundCompleteEvent } from '../chat/convert_graph_events'; +import { toLangchainTool } from '../chat/utils'; +import { createResearcherAgentGraph } from './graph'; export interface RunSearchAgentContext { logger: Logger; @@ -96,7 +85,7 @@ export const runSearchAgent: RunResearcherAgentFn = async ( const langchainTools = researcherTools.map((tool) => toLangchainTool({ tool, logger })); - const agentGraph = await createAgentGraph({ + const agentGraph = await createResearcherAgentGraph({ logger, chatModel: model.chatModel, tools: langchainTools, @@ -118,12 +107,7 @@ export const runSearchAgent: RunResearcherAgentFn = async ( } ); - const events$ = from(eventStream).pipe( - filter(isStreamEvent), - // convertGraphEvents({ graphName: agentGraphName, runName: agentGraphName }), - // addRoundCompleteEvent({ userInput: instructions }), - shareReplay() - ); + const events$ = from(eventStream).pipe(filter(isStreamEvent), shareReplay()); events$.subscribe((event) => { // later we should emit reasoning events from there. @@ -133,6 +117,9 @@ export const runSearchAgent: RunResearcherAgentFn = async ( const lastEvent = await lastValueFrom(events$); const generatedAnswer = lastEvent.data.output.generatedAnswer; + console.log('**** last output'); + console.log(JSON.stringify(lastEvent.data.output, null, 2)); + return { answer: generatedAnswer, }; From 688c98a0d629399d340113d3c3c8b68c99d29023 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 17 Jun 2025 11:18:42 +0000 Subject: [PATCH 13/17] [CI] Auto-commit changed files from 'node scripts/generate codeowners' --- .github/CODEOWNERS | 1 + .../packages/shared/onechat/onechat-genai-utils/package.json | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 5b93f7141f4d9..247a8dab217ea 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -867,6 +867,7 @@ x-pack/platform/packages/shared/ml/runtime_field_utils @elastic/ml-ui x-pack/platform/packages/shared/ml/trained_models_utils @elastic/ml-ui x-pack/platform/packages/shared/onechat/onechat-browser @elastic/workchat-eng x-pack/platform/packages/shared/onechat/onechat-common @elastic/workchat-eng +x-pack/platform/packages/shared/onechat/onechat-genai-utils @elastic/workchat-eng x-pack/platform/packages/shared/onechat/onechat-server @elastic/workchat-eng x-pack/platform/packages/shared/security/api_key_management @elastic/kibana-security x-pack/platform/packages/shared/security/form_components @elastic/kibana-security diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json index 220e3faded6f5..bfb6f78dc0292 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/package.json @@ -2,5 +2,5 @@ "name": "@kbn/onechat-genai-utils", "private": true, "version": "1.0.0", - "license": "Elastic License 2.0 OR AGPL-3.0-only OR SSPL-1.0" + "license": "Elastic License 2.0" } \ No newline at end of file From 9e3c077bb3a0b860aa06ac03865ade5cee24b476 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 17 Jun 2025 11:46:32 +0000 Subject: [PATCH 14/17] [CI] Auto-commit changed files from 'node scripts/yarn_deduplicate' --- .../shared/onechat/onechat-genai-utils/tsconfig.json | 10 +++++++++- .../shared/onechat/onechat-server/tsconfig.json | 1 + x-pack/platform/plugins/shared/onechat/tsconfig.json | 3 ++- 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json index 63f0b5ff33faa..f2e629ec5d882 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json @@ -13,5 +13,13 @@ "exclude": [ "target/**/*" ], - "kbn_references": [] + "kbn_references": [ + "@kbn/core-http-server", + "@kbn/onechat-common", + "@kbn/onechat-server", + "@kbn/core-elasticsearch-server", + "@kbn/inference-common", + "@kbn/inference-plugin", + "@kbn/zod", + ] } diff --git a/x-pack/platform/packages/shared/onechat/onechat-server/tsconfig.json b/x-pack/platform/packages/shared/onechat/onechat-server/tsconfig.json index f48664e99769e..8a9d27680715a 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-server/tsconfig.json +++ b/x-pack/platform/packages/shared/onechat/onechat-server/tsconfig.json @@ -21,5 +21,6 @@ "@kbn/core-http-server", "@kbn/onechat-common", "@kbn/inference-common", + "@kbn/logging", ] } diff --git a/x-pack/platform/plugins/shared/onechat/tsconfig.json b/x-pack/platform/plugins/shared/onechat/tsconfig.json index e6707a8b01452..873cacd519da8 100644 --- a/x-pack/platform/plugins/shared/onechat/tsconfig.json +++ b/x-pack/platform/plugins/shared/onechat/tsconfig.json @@ -36,6 +36,7 @@ "@kbn/core-http-browser", "@kbn/features-plugin", "@kbn/core-logging-server-mocks", - "@kbn/i18n" + "@kbn/i18n", + "@kbn/onechat-genai-utils" ] } From 46839ad28e48e07592596aada70cd03447384fc3 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 17 Jun 2025 12:00:38 +0000 Subject: [PATCH 15/17] [CI] Auto-commit changed files from 'node scripts/eslint --no-cache --fix' --- .../onechat/server/services/agents/chat/run_chat_agent.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts index dd3e40ec92ab9..688afc9471fd1 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/run_chat_agent.ts @@ -22,7 +22,11 @@ import type { ExecutableTool, ToolProvider, } from '@kbn/onechat-server'; -import { providerToLangchainTools, toLangchainTool, conversationToLangchainMessages } from './utils'; +import { + providerToLangchainTools, + toLangchainTool, + conversationToLangchainMessages, +} from './utils'; import { createAgentGraph } from './graph'; import { convertGraphEvents, addRoundCompleteEvent } from './convert_graph_events'; From 8138de3875c027ba9707cddf951fbce595c4d073 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Tue, 17 Jun 2025 14:53:57 +0200 Subject: [PATCH 16/17] remove console.log --- .../server/services/agents/researcher/run_researcher_agent.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts index fefdca6fbcfc1..fc1691e39cbd6 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -117,9 +117,6 @@ export const runSearchAgent: RunResearcherAgentFn = async ( const lastEvent = await lastValueFrom(events$); const generatedAnswer = lastEvent.data.output.generatedAnswer; - console.log('**** last output'); - console.log(JSON.stringify(lastEvent.data.output, null, 2)); - return { answer: generatedAnswer, }; From 44dd5a1fdd0e175bd89953b0331e97ec8cf0aae8 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Fri, 20 Jun 2025 10:12:03 +0200 Subject: [PATCH 17/17] start working on reflection events --- .../onechat/onechat-common/agents/events.ts | 14 +++ .../onechat/onechat-common/agents/index.ts | 3 + .../langchain/graph_events.ts | 71 ++++++++++++++ .../onechat-genai-utils/langchain/index.ts | 12 +++ .../onechat-genai-utils/langchain/messages.ts | 25 +++++ .../agents/chat/convert_graph_events.ts | 29 +++--- .../services/agents/researcher/backlog.ts | 2 +- .../agents/researcher/convert_graph_events.ts | 73 ++++++++++++++ .../services/agents/researcher/graph.ts | 97 ++++++++++--------- .../agents/researcher/researcher_as_tool.ts | 4 +- .../agents/researcher/run_researcher_agent.ts | 31 +++--- 11 files changed, 282 insertions(+), 79 deletions(-) create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/graph_events.ts create mode 100644 x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.ts create mode 100644 x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/convert_graph_events.ts diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/agents/events.ts b/x-pack/platform/packages/shared/onechat/onechat-common/agents/events.ts index e49d2f8a59d6e..767d7cd428324 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/agents/events.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/agents/events.ts @@ -12,6 +12,7 @@ import type { StructuredToolIdentifier } from '../tools/tools'; export enum ChatAgentEventType { toolCall = 'toolCall', toolResult = 'toolResult', + reasoning = 'reasoning', messageChunk = 'messageChunk', messageComplete = 'messageComplete', roundComplete = 'roundComplete', @@ -53,6 +54,18 @@ export const isToolResultEvent = (event: OnechatEvent): event is To return event.type === ChatAgentEventType.toolResult; }; +// reasoning + +export interface ReasoningEventData { + reasoning: string; +} + +export type ReasoningEvent = ChatAgentEventBase; + +export const isReasoningEvent = (event: OnechatEvent): event is ReasoningEvent => { + return event.type === ChatAgentEventType.reasoning; +}; + // Message chunk export interface MessageChunkEventData { @@ -116,6 +129,7 @@ export const isRoundCompleteEvent = ( export type ChatAgentEvent = | ToolCallEvent | ToolResultEvent + | ReasoningEvent | MessageChunkEvent | MessageCompleteEvent | RoundCompleteEvent; diff --git a/x-pack/platform/packages/shared/onechat/onechat-common/agents/index.ts b/x-pack/platform/packages/shared/onechat/onechat-common/agents/index.ts index 23df0aca44ed6..e44a3389dd9ce 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-common/agents/index.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-common/agents/index.ts @@ -19,6 +19,8 @@ export { type ToolResultEventData, type ToolCallEvent, type ToolCallEventData, + type ReasoningEvent, + type ReasoningEventData, type MessageChunkEventData, type MessageChunkEvent, type MessageCompleteEventData, @@ -27,6 +29,7 @@ export { type RoundCompleteEvent, isToolCallEvent, isToolResultEvent, + isReasoningEvent, isMessageChunkEvent, isMessageCompleteEvent, isRoundCompleteEvent, diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/graph_events.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/graph_events.ts new file mode 100644 index 0000000000000..7ec567d4d2a9f --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/graph_events.ts @@ -0,0 +1,71 @@ +/* + * 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 type { AIMessageChunk } from '@langchain/core/messages'; +import { StreamEvent as LangchainStreamEvent } from '@langchain/core/tracers/log_stream'; +import { + ChatAgentEventType, + MessageChunkEvent, + ReasoningEvent, + MessageCompleteEvent, +} from '@kbn/onechat-common/agents'; +import { extractTextContent } from './messages'; + +export const matchGraphName = (event: LangchainStreamEvent, graphName: string): boolean => { + return event.metadata.graphName === graphName; +}; + +export const matchGraphNode = (event: LangchainStreamEvent, nodeName: string): boolean => { + return event.metadata.langgraph_node === nodeName; +}; + +export const matchEvent = (event: LangchainStreamEvent, eventName: string): boolean => { + return event.event === eventName; +}; + +export const matchName = (event: LangchainStreamEvent, name: string): boolean => { + return event.name === name; +}; + +export const hasTag = (event: LangchainStreamEvent, tag: string): boolean => { + return (event.tags ?? []).includes(tag); +}; + +export const createTextChunkEvent = ( + chunk: AIMessageChunk, + { defaultMessageId = 'unknown' }: { defaultMessageId?: string } = {} +): MessageChunkEvent => { + return { + type: ChatAgentEventType.messageChunk, + data: { + messageId: chunk.id ?? defaultMessageId, + textChunk: extractTextContent(chunk), + }, + }; +}; + +export const createMessageEvent = ( + content: string, + { messageId = 'unknown' }: { messageId?: string } = {} +): MessageCompleteEvent => { + return { + type: ChatAgentEventType.messageComplete, + data: { + messageId, + messageContent: content, + }, + }; +}; + +export const createReasoningEvent = (reasoning: string): ReasoningEvent => { + return { + type: ChatAgentEventType.reasoning, + data: { + reasoning, + }, + }; +}; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts index 1fec1c76430eb..5e63301383577 100644 --- a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts @@ -4,3 +4,15 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + +export { + matchGraphName, + matchGraphNode, + matchName, + matchEvent, + hasTag, + createTextChunkEvent, + createMessageEvent, + createReasoningEvent, +} from './graph_events'; +export { extractTextContent } from './messages'; diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.ts new file mode 100644 index 0000000000000..2559ae65b186f --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.ts @@ -0,0 +1,25 @@ +/* + * 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 { BaseMessage, MessageContentComplex } from '@langchain/core/messages'; + +/** + * Extract the text content from a langchain message or chunk. + */ +export const extractTextContent = (message: BaseMessage): string => { + if (typeof message.content === 'string') { + return message.content; + } else { + let content = ''; + for (const item of message.content as MessageContentComplex[]) { + if (item.type === 'text') { + content += item.text; + } + } + return content; + } +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/convert_graph_events.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/convert_graph_events.ts index 823130fbcb789..ad3e69fcfa0d6 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/convert_graph_events.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/convert_graph_events.ts @@ -21,7 +21,14 @@ import { } from '@kbn/onechat-common/agents'; import { RoundInput, ConversationRoundStepType } from '@kbn/onechat-common/chat'; import { StructuredToolIdentifier, toStructuredToolIdentifier } from '@kbn/onechat-common/tools'; -import { extractTextContent, getToolCalls } from './utils/from_langchain_messages'; +import { + matchGraphName, + matchEvent, + matchName, + createTextChunkEvent, + extractTextContent, +} from '@kbn/onechat-genai-utils/langchain'; +import { getToolCalls } from './utils/from_langchain_messages'; export type ConvertedEvents = | MessageChunkEvent @@ -90,21 +97,13 @@ export const convertGraphEvents = ({ } // stream text chunks for the UI - if (event.event === 'on_chat_model_stream') { + if (matchEvent(event, 'on_chat_model_stream')) { const chunk: AIMessageChunk = event.data.chunk; - const chunkEvent: MessageChunkEvent = { - type: ChatAgentEventType.messageChunk, - data: { - messageId: chunk.id ?? 'todo', - textChunk: extractTextContent(chunk), - }, - }; - - return of(chunkEvent); + return of(createTextChunkEvent(chunk)); } // emit tool calls or full message on each agent step - if (event.event === 'on_chain_end' && event.name === 'agent') { + if (matchEvent(event, 'on_chain_end') && matchName(event, 'agent')) { const addedMessages: BaseMessage[] = event.data.output.addedMessages ?? []; const lastMessage = addedMessages[addedMessages.length - 1]; @@ -139,7 +138,7 @@ export const convertGraphEvents = ({ } // emit tool result events - if (event.event === 'on_chain_end' && event.name === 'tools') { + if (matchEvent(event, 'on_chain_end') && matchName(event, 'tools')) { const toolMessages: ToolMessage[] = event.data.output.addedMessages ?? []; const toolResultEvents: ToolResultEvent[] = []; @@ -166,7 +165,3 @@ export const convertGraphEvents = ({ ); }; }; - -const matchGraphName = (event: LangchainStreamEvent, graphName: string): boolean => { - return event.metadata.graphName === graphName; -}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts index a4cf4a13df4d7..4420b42078d67 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/backlog.ts @@ -15,7 +15,7 @@ export interface ActionResult { export interface ReflectionResult { isSufficient: boolean; nextQuestions: string[]; - reasoning?: string; + reasoning: string; } export type BacklogItem = ActionResult | ReflectionResult; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/convert_graph_events.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/convert_graph_events.ts new file mode 100644 index 0000000000000..62cb41b2c9721 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/convert_graph_events.ts @@ -0,0 +1,73 @@ +/* + * 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 { v4 as uuidv4 } from 'uuid'; +import { StreamEvent as LangchainStreamEvent } from '@langchain/core/tracers/log_stream'; +import type { AIMessageChunk } from '@langchain/core/messages'; +import { EMPTY, mergeMap, of, OperatorFunction } from 'rxjs'; +import { + MessageChunkEvent, + MessageCompleteEvent, + ReasoningEvent, +} from '@kbn/onechat-common/agents'; +import { + matchGraphName, + matchEvent, + matchName, + hasTag, + createTextChunkEvent, + createMessageEvent, + createReasoningEvent, +} from '@kbn/onechat-genai-utils/langchain'; +import type { StateType } from './graph'; +import { lastReflectionResult } from './backlog'; + +export type ResearcherAgentEvents = MessageChunkEvent | MessageCompleteEvent | ReasoningEvent; + +export const convertGraphEvents = ({ + graphName, +}: { + graphName: string; +}): OperatorFunction => { + return (streamEvents$) => { + const messageId = uuidv4(); + return streamEvents$.pipe( + mergeMap((event) => { + if (!matchGraphName(event, graphName)) { + return EMPTY; + } + + // response text chunks + if (matchEvent(event, 'on_chat_model_stream') && hasTag(event, 'researcher-answer')) { + const chunk: AIMessageChunk = event.data.chunk; + + const messageChunkEvent = createTextChunkEvent(chunk, { defaultMessageId: messageId }); + return of(messageChunkEvent); + } + + // response message + if (matchEvent(event, 'on_chain_end') && matchName(event, 'answer')) { + const { generatedAnswer } = event.data.output as StateType; + + const messageEvent = createMessageEvent(generatedAnswer); + return of(messageEvent); + } + + // emit reasoning events for "reflection" step + if (matchEvent(event, 'on_chain_end') && matchName(event, 'reflection')) { + const { backlog } = event.data.output as StateType; + const reflectionResult = lastReflectionResult(backlog); + + const reasoningEvent = createReasoningEvent(reflectionResult.reasoning); + return of(reasoningEvent); + } + + return EMPTY; + }) + ); + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts index df8f1604f3141..6b08d459f54fb 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -17,6 +17,40 @@ import { getReflectionPrompt, getExecutionPrompt, getAnswerPrompt } from './prom import { extractToolResults } from './utils'; import { ActionResult, ReflectionResult, BacklogItem, lastReflectionResult } from './backlog'; +const StateAnnotation = Annotation.Root({ + // inputs + initialQuery: Annotation(), // the search query + cycleBudget: Annotation(), // budget in number of cycles + // internal state + remainingCycles: Annotation(), + actionsQueue: Annotation({ + reducer: (state, actions) => { + return actions ?? state; + }, + default: () => [], + }), + pendingActions: Annotation({ + reducer: (state, actions) => { + return actions === 'clear' ? [] : [...state, ...actions]; + }, + default: () => [], + }), + backlog: Annotation({ + reducer: (current, next) => { + return [...current, ...next]; + }, + default: () => [], + }), + // outputs + generatedAnswer: Annotation(), +}); + +export type StateType = typeof StateAnnotation.State; + +type ResearchStepState = StateType & { + researchGoal: ResearchGoal; +}; + export interface ResearchGoal { question: string; } @@ -30,40 +64,6 @@ export const createResearcherAgentGraph = async ({ tools: StructuredTool[]; logger: Logger; }) => { - const StateAnnotation = Annotation.Root({ - // inputs - initialQuery: Annotation(), // the search query - cycleBudget: Annotation(), // budget in number of cycles - // internal state - remainingCycles: Annotation(), - actionsQueue: Annotation({ - reducer: (state, actions) => { - return actions ?? state; - }, - default: () => [], - }), - pendingActions: Annotation({ - reducer: (state, actions) => { - return actions === 'clear' ? [] : [...state, ...actions]; - }, - default: () => [], - }), - backlog: Annotation({ - reducer: (current, next) => { - return [...current, ...next]; - }, - default: () => [], - }), - // outputs - generatedAnswer: Annotation(), - }); - - type StateType = typeof StateAnnotation.State; - - type ResearchStepState = StateType & { - researchGoal: ResearchGoal; - }; - const stringify = (obj: unknown) => JSON.stringify(obj, null, 2); /** @@ -145,26 +145,27 @@ export const createResearcherAgentGraph = async ({ }; const reflection = async (state: StateType) => { - const reflectModel = chatModel.withStructuredOutput( - z.object({ - isSufficient: z.boolean().describe( - `Set to true if the current information fully answers the user question without requiring further research. + const reflectModel = chatModel + .withStructuredOutput( + z.object({ + isSufficient: z.boolean().describe( + `Set to true if the current information fully answers the user question without requiring further research. Set to false if any knowledge gaps or unresolved sub-problems remain.` - ), - nextQuestions: z.array(z.string()).describe( - `A list of self-contained, actionable research questions or sub-problems that need to be explored + ), + nextQuestions: z.array(z.string()).describe( + `A list of self-contained, actionable research questions or sub-problems that need to be explored further to fully answer the user question. Leave empty if isSufficient is true.` - ), - reasoning: z - .string() - .optional() - .describe( + ), + reasoning: z.string().describe( `Brief internal reasoning explaining why the current information is sufficient or not. You may list what was already answered, what gaps exist, or whether decomposition was necessary. Use this as your thought process or scratchpad before producing the final output.` ), - }) - ); + }) + ) + .withConfig({ + tags: ['researcher-reflection'], + }); const response: ReflectionResult = await reflectModel.invoke( getReflectionPrompt({ diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts index 22da4714df90a..38f0ec36f0ac3 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts @@ -8,7 +8,7 @@ import { z } from '@kbn/zod'; import type { RegisteredTool } from '@kbn/onechat-server'; import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; -import { runSearchAgent } from './run_researcher_agent'; +import { runResearcherAgent } from './run_researcher_agent'; const researcherSchema = z.object({ instructions: z.string().describe('Research instructions for the agent'), @@ -39,7 +39,7 @@ export const researcherTool = (): RegisteredTool { - const searchAgentResult = await runSearchAgent( + const searchAgentResult = await runResearcherAgent( { instructions, toolProvider, diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts index fc1691e39cbd6..f23dd1a739a72 100644 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -9,20 +9,26 @@ import { from, filter, shareReplay, lastValueFrom } from 'rxjs'; import type { Logger } from '@kbn/logging'; import { StreamEvent } from '@langchain/core/tracers/log_stream'; import type { KibanaRequest } from '@kbn/core-http-server'; -import { ChatAgentEvent, BuiltinToolIds, builtinToolProviderId } from '@kbn/onechat-common'; +import { + ChatAgentEvent, + BuiltinToolIds, + builtinToolProviderId, + isMessageCompleteEvent, +} from '@kbn/onechat-common'; import type { ModelProvider, ScopedRunner, ToolProvider } from '@kbn/onechat-server'; import { filterProviderTools } from '@kbn/onechat-genai-utils/framework'; import { toLangchainTool } from '../chat/utils'; import { createResearcherAgentGraph } from './graph'; +import { convertGraphEvents } from './convert_graph_events'; -export interface RunSearchAgentContext { +export interface RunResearcherAgentContext { logger: Logger; request: KibanaRequest; modelProvider: ModelProvider; runner: ScopedRunner; } -export interface RunSearchAgentParams { +export interface RunResearcherAgentParams { /** * The search instructions */ @@ -48,8 +54,8 @@ export interface RunResearcherAgentResponse { } export type RunResearcherAgentFn = ( - params: RunSearchAgentParams, - context: RunSearchAgentContext + params: RunResearcherAgentParams, + context: RunResearcherAgentContext ) => Promise; const agentGraphName = 'researcher-agent'; @@ -60,7 +66,7 @@ const noopOnEvent = () => {}; /** * Create the handler function for the default onechat agent. */ -export const runSearchAgent: RunResearcherAgentFn = async ( +export const runResearcherAgent: RunResearcherAgentFn = async ( { instructions, cycleBudget = defaultCycleBudget, toolProvider, onEvent = noopOnEvent }, { logger, request, modelProvider } ) => { @@ -107,15 +113,18 @@ export const runSearchAgent: RunResearcherAgentFn = async ( } ); - const events$ = from(eventStream).pipe(filter(isStreamEvent), shareReplay()); + const events$ = from(eventStream).pipe( + filter(isStreamEvent), + convertGraphEvents({ graphName: agentGraphName }), + shareReplay() + ); - events$.subscribe((event) => { + events$.pipe().subscribe((event) => { // later we should emit reasoning events from there. }); - // event: 'on_chain_end', name: 'researcher-agent' - const lastEvent = await lastValueFrom(events$); - const generatedAnswer = lastEvent.data.output.generatedAnswer; + const lastEvent = await lastValueFrom(events$.pipe(filter(isMessageCompleteEvent))); + const generatedAnswer = lastEvent.data.messageContent; return { answer: generatedAnswer,