diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e14c7a3919d6a..827ad0095acea 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -868,6 +868,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/package.json b/package.json index ee963e0d5c98a..4b7070e341eab 100644 --- a/package.json +++ b/package.json @@ -739,6 +739,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 c71459a2a45fa..f7049bc019fc2 100644 --- a/tsconfig.base.json +++ b/tsconfig.base.json @@ -1406,6 +1406,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/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-common/index.ts b/x-pack/platform/packages/shared/onechat/onechat-common/index.ts index 1873dfe5a5a78..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,6 +23,8 @@ export { createBuiltinToolId, builtinToolProviderId, unknownToolProviderId, + 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 new file mode 100644 index 0000000000000..b1fa476507895 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-common/tools/constants.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. + */ + +/** + * Ids of built-in onechat tools + */ +export const BuiltinToolIds = { + indexExplorer: 'index_explorer', + relevanceSearch: 'relevance_search', + naturalLanguageSearch: 'nl_search', + listIndices: 'list_indices', + getIndexMapping: 'get_index_mapping', + getDocumentById: 'get_document_by_id', + generateEsql: 'generate_esql', + executeEsql: 'execute_esql', + researcherAgent: 'researcher_agent', +}; + +/** + * Common set of tags used for platform tools. + */ +export const BuiltinTags = { + /** + * 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..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,3 +23,4 @@ export { builtinToolProviderId, unknownToolProviderId, } from './tools'; +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/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/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/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 new file mode 100644 index 0000000000000..5e63301383577 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts @@ -0,0 +1,18 @@ +/* + * 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 { + 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/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..bfb6f78dc0292 --- /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" +} \ No newline at end of file diff --git a/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts new file mode 100644 index 0000000000000..2b95805c81911 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/generate_esql.ts @@ -0,0 +1,90 @@ +/* + * 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 { filter, toArray, firstValueFrom } from 'rxjs'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import { isChatCompletionMessageEvent, isChatCompletionEvent } from '@kbn/inference-common'; +import { naturalLanguageToEsql } from '@kbn/inference-plugin/server'; +import type { ScopedModel } from '@kbn/onechat-server'; +import { indexExplorer } from './index_explorer'; +import { getIndexMappings } from './steps/get_mappings'; +import { extractEsqlQueries } from './utils/esql'; + +export interface GenerateEsqlResponse { + answer: string; + queries: string[]; +} + +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, + }); + 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: ` + Your task is to generate an ES|QL query. + + - User query: "${query}", + - Additional context: "${context ?? 'N/A'} + - Index to use: "${selectedIndex}" + - Mapping of this index: + \`\`\`json + ${JSON.stringify(mappings, undefined, 2)} + \`\`\` + + Given those info, please generate an ES|QL query to address the user request + `, + }); + + 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, + }; +}; 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..c7e846f536c9e --- /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 './steps/list_indices'; +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/packages/shared/onechat/onechat-genai-utils/tools/utils/esql.ts b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/esql.ts new file mode 100644 index 0000000000000..88ca462d51a7f --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/utils/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 { 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> = []; + + 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; +}; + +export 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/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..f2e629ec5d882 --- /dev/null +++ b/x-pack/platform/packages/shared/onechat/onechat-genai-utils/tsconfig.json @@ -0,0 +1,25 @@ +{ + "extends": "../../../../../../tsconfig.base.json", + "compilerOptions": { + "outDir": "target/types", + "types": [ + "jest", + "node" + ] + }, + "include": [ + "**/*.ts", + ], + "exclude": [ + "target/**/*" + ], + "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/src/tools.ts b/x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts index d78c2d8035008..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 { @@ -22,7 +23,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. @@ -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/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/server/plugin.ts b/x-pack/platform/plugins/shared/onechat/server/plugin.ts index f4e7587aea81d..321a18cbb6823 100644 --- a/x-pack/platform/plugins/shared/onechat/server/plugin.ts +++ b/x-pack/platform/plugins/shared/onechat/server/plugin.ts @@ -19,6 +19,7 @@ import type { import { registerRoutes } from './routes'; import { ServiceManager } from './services'; import { registerFeatures } from './features'; +import { registerTools } from './tools'; import { ONECHAT_MCP_SERVER_UI_SETTING_ID } from '../common/constants'; export class OnechatPlugin @@ -50,6 +51,8 @@ export class OnechatPlugin registerFeatures({ features: pluginsSetup.features }); + registerTools({ tools: serviceSetups.tools }); + coreSetup.uiSettings.register({ [ONECHAT_MCP_SERVER_UI_SETTING_ID]: { description: i18n.translate('onechat.uiSettings.mcpServer.description', { 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 88% 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 index 823130fbcb789..ad3e69fcfa0d6 100644 --- 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 @@ -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/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 95% 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 index 4032b6754b979..688afc9471fd1 100644 --- 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 @@ -22,7 +22,11 @@ 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 +97,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/system_prompt.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/system_prompt.ts new file mode 100644 index 0000000000000..35a569c628278 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/system_prompt.ts @@ -0,0 +1,47 @@ +/* + * 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 { BaseMessage, BaseMessageLike } from '@langchain/core/messages'; +import { BuiltinToolIds } from '@kbn/onechat-common'; + +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. + + - 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. + + - 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. + + - 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) => { + return `${systemPrompt} + + ### Additional info + - The current date is: ${new Date().toISOString()} + - You can use markdown format to structure your response + `; +}; + +export const withSystemPrompt = ({ + systemPrompt, + messages, +}: { + systemPrompt: string; + messages: BaseMessage[]; +}): BaseMessageLike[] => { + return [['system', getFullSystemPrompt(systemPrompt)], ...messages]; +}; 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 84% 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 index 7b3cba794524e..a5ec190b06d6a 100644 --- 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 @@ -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/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 92% 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 index 9dbba5850d3cf..7282c341e6877 100644 --- 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 @@ -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/conversational/utils/to_langchain_messages.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/chat/utils/to_langchain_messages.ts similarity index 97% 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 index 7317f1af5e748..ae80b9d372dab 100644 --- 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 @@ -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/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/conversational/system_prompt.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts deleted file mode 100644 index b5b27c66605e4..0000000000000 --- a/x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts +++ /dev/null @@ -1,31 +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 { BaseMessage, BaseMessageLike } from '@langchain/core/messages'; - -export const defaultSystemPrompt = - 'You are a helpful chat assistant from the Elasticsearch company.'; - -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 - `; -}; - -export const withSystemPrompt = ({ - systemPrompt, - messages, -}: { - systemPrompt: string; - messages: BaseMessage[]; -}): BaseMessageLike[] => { - return [['system', getFullSystemPrompt(systemPrompt)], ...messages]; -}; 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..4420b42078d67 --- /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/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 new file mode 100644 index 0000000000000..6b08d459f54fb --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/graph.ts @@ -0,0 +1,248 @@ +/* + * 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, Send } 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 { getToolCalls, extractTextContent } from '../chat/utils/from_langchain_messages'; +import { getReflectionPrompt, getExecutionPrompt, getAnswerPrompt } from './prompts'; +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; +} + +export const createResearcherAgentGraph = async ({ + chatModel, + tools, + logger: log, +}: { + chatModel: InferenceChatModel; + tools: StructuredTool[]; + logger: Logger; +}) => { + 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: StateType) => { + const firstAction: ResearchGoal = { + question: state.initialQuery, + }; + return { + actionsQueue: [firstAction], + remainingCycles: state.cycleBudget, + }; + }; + + 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; + + log.trace(() => `performSearch - nextItem: ${stringify(nextItem)}`); + + const toolNode = new ToolNode(tools); + const executionModel = chatModel.bindTools(tools); + + const response = await executionModel.invoke( + getExecutionPrompt({ + currentResearchGoal: nextItem, + backlog: state.backlog, + }) + ); + const toolCalls = getToolCalls(response); + + log.trace(() => `performSearch - toolCalls: ${stringify(toolCalls)}`); + + const toolMessages = await toolNode.invoke([response]); + const toolResults = extractToolResults(toolMessages); + + 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 { + pendingActions: [...actionResults], + }; + }; + + 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: 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. + 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().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({ + userQuery: state.initialQuery, + backlog: state.backlog, + maxFollowUpQuestions: 3, + remainingCycles: state.remainingCycles - 1, + }) + ); + + log.trace( + () => + `reflection - remaining cycles: ${state.remainingCycles} - response: ${stringify(response)}` + ); + + return { + remainingCycles: state.remainingCycles - 1, + backlog: [response], + actionsQueue: [ + ...state.actionsQueue, + ...response.nextQuestions.map((nextQuestion) => ({ question: nextQuestion })), + ], + }; + }; + + const evaluateReflection = async (state: StateType) => { + const remainingCycles = state.remainingCycles; + const reflectionResult = lastReflectionResult(state.backlog); + + if (reflectionResult.isSufficient || remainingCycles <= 0) { + return 'answer'; + } + return dispatchActions(state); + }; + + const answer = async (state: StateType) => { + const answerModel = chatModel.withConfig({ + tags: ['researcher-answer'], + }); + + const response = await answerModel.invoke( + getAnswerPrompt({ + userQuery: state.initialQuery, + backlog: state.backlog, + }) + ); + + const generatedAnswer = extractTextContent(response); + + log.trace(() => `answer - response ${stringify(generatedAnswer)}`); + + return { + generatedAnswer, + }; + }; + + // 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('perform_search', performSearch) + .addNode('collect_results', collectResults) + .addNode('reflection', reflection) + .addNode('answer', answer) + // edges + .addEdge('__start__', 'initialize') + .addConditionalEdges('initialize', dispatchActions, { + perform_search: 'perform_search', + }) + .addEdge('perform_search', 'collect_results') + .addEdge('collect_results', 'reflection') + .addConditionalEdges('reflection', evaluateReflection, { + perform_search: 'perform_search', + answer: 'answer', + }) + .addEdge('answer', '__end__') + .compile(); + + return graph; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/index.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/index.ts new file mode 100644 index 0000000000000..45be5d76abcb8 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/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/researcher/prompts.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts new file mode 100644 index 0000000000000..f6397917464ce --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/prompts.ts @@ -0,0 +1,281 @@ +/* + * 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 { BuiltinToolIds as Tools } from '@kbn/onechat-common'; +import type { ResearchGoal } from './graph'; +import { + isActionResult, + isReflectionResult, + BacklogItem, + ReflectionResult, + ActionResult, +} from './backlog'; + +export const getExecutionPrompt = ({ + currentResearchGoal, + backlog, +}: { + currentResearchGoal: ResearchGoal; + backlog: BacklogItem[]; +}): BaseMessageLike[] => { + return [ + [ + 'system', + `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 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 + - 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. + + ### 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 + - The current date is ${new Date().toISOString()}. + `, + ], + [ + 'user', + ` + ### Current Research Goal + + Trying to find information about: "${currentResearchGoal.question}" + + ### Previous Actions + + ${renderBacklog(backlog)} + `, + ], + ]; +}; + +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 [ + [ + 'system', + `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 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. + + 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. + + Additional information: + - The current date is ${new Date().toISOString()}. + + Output Format: + - Format your response as a JSON object with these exact keys: + - "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: information is sufficient + \`\`\`json + { + "isSufficient": true, + "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', + ` + ## User question + + "${userQuery}" + + ## Backlog + + ${renderBacklog(backlog)} + `, + ], + ]; +}; + +export const getAnswerPrompt = ({ + userQuery, + backlog, +}: { + userQuery: string; + backlog: BacklogItem[]; +}): BaseMessageLike[] => { + return [ + [ + 'system', + `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()}. + + `, + ], + [ + 'user', + ` + ### User question + + "${userQuery}" + + ### Gathered information + + ${renderBacklog(backlog.filter(isActionResult))} + `, + ], + ]; +}; + +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/researcher_as_tool.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts new file mode 100644 index 0000000000000..38f0ec36f0ac3 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/researcher_as_tool.ts @@ -0,0 +1,58 @@ +/* + * 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 { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; +import { runResearcherAgent } from './run_researcher_agent'; + +const researcherSchema = z.object({ + instructions: z.string().describe('Research instructions for the agent'), +}); + +export interface ResearcherResponse { + answer: string; +} + +export const researcherTool = (): RegisteredTool => { + return { + id: BuiltinToolIds.researcherAgent, + 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 runResearcherAgent( + { + instructions, + toolProvider, + }, + { request, modelProvider, runner, logger } + ); + + return { + answer: searchAgentResult.answer, + }; + }, + meta: { + tags: [BuiltinTags.retrieval], + }, + }; +}; 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 new file mode 100644 index 0000000000000..f23dd1a739a72 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/run_researcher_agent.ts @@ -0,0 +1,136 @@ +/* + * 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 { 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, + 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 RunResearcherAgentContext { + logger: Logger; + request: KibanaRequest; + modelProvider: ModelProvider; + runner: ScopedRunner; +} + +export interface RunResearcherAgentParams { + /** + * 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 + */ + toolProvider: ToolProvider; + /** + * Handler to react to the agent's events. + */ + onEvent?: (event: ChatAgentEvent) => void; +} + +export interface RunResearcherAgentResponse { + answer: string; +} + +export type RunResearcherAgentFn = ( + params: RunResearcherAgentParams, + context: RunResearcherAgentContext +) => Promise; + +const agentGraphName = 'researcher-agent'; +const defaultCycleBudget = 5; + +const noopOnEvent = () => {}; + +/** + * Create the handler function for the default onechat agent. + */ +export const runResearcherAgent: RunResearcherAgentFn = async ( + { instructions, cycleBudget = defaultCycleBudget, toolProvider, onEvent = noopOnEvent }, + { logger, request, modelProvider } +) => { + const model = await modelProvider.getDefaultModel(); + + 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 createResearcherAgentGraph({ + logger, + chatModel: model.chatModel, + tools: langchainTools, + }); + + const eventStream = agentGraph.streamEvents( + { + initialQuery: instructions, + cycleBudget, + }, + { + version: 'v2', + runName: agentGraphName, + metadata: { + graphName: agentGraphName, + }, + recursionLimit: cycleBudget * 10, + callbacks: [], + } + ); + + const events$ = from(eventStream).pipe( + filter(isStreamEvent), + convertGraphEvents({ graphName: agentGraphName }), + shareReplay() + ); + + events$.pipe().subscribe((event) => { + // later we should emit reasoning events from there. + }); + + const lastEvent = await lastValueFrom(events$.pipe(filter(isMessageCompleteEvent))); + const generatedAnswer = lastEvent.data.messageContent; + + return { + answer: generatedAnswer, + }; +}; + +const isStreamEvent = (input: any): input is StreamEvent => { + return 'event' in input; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/utils.ts b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/utils.ts new file mode 100644 index 0000000000000..9439dffdbb1d7 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/services/agents/researcher/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 '../chat/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/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/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/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/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..b5fb64539cf65 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/register_tools.ts @@ -0,0 +1,38 @@ +/* + * 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 { + getDocumentByIdTool, + executeEsqlTool, + naturalLanguageSearchTool, + generateEsqlTool, + relevanceSearchTool, + getIndexMappingsTool, + listIndicesTool, + indexExplorerTool, +} from './retrieval'; +import { researcherTool } from '../services/agents/researcher'; + +export const registerTools = ({ tools: registry }: { tools: ToolsServiceSetup }) => { + const tools: Array> = [ + getDocumentByIdTool(), + executeEsqlTool(), + naturalLanguageSearchTool(), + generateEsqlTool(), + relevanceSearchTool(), + getIndexMappingsTool(), + listIndicesTool(), + indexExplorerTool(), + researcherTool(), + ]; + + tools.forEach((tool) => { + registry.register(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..48630bd154728 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/execute_esql.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. + */ + +import { z } from '@kbn/zod'; +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 const executeEsqlTool = (): RegisteredTool => { + return { + 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: [BuiltinTags.retrieval], + }, + }; +}; 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..4fd9d32eeeff3 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/generate_esql.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 { 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< + typeof nlToEsqlToolSchema, + GenerateEsqlResponse +> => { + 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 new file mode 100644 index 0000000000000..c44b1e6c7296e --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_document_by_id.ts @@ -0,0 +1,33 @@ +/* + * 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 { 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 const getDocumentByIdTool = (): RegisteredTool< + typeof getDocumentByIdSchema, + GetDocumentByIdResult +> => { + return { + 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: [BuiltinTags.retrieval], + }, + }; +}; 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..9beb2b69f5bee --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/get_index_mapping.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 { z } from '@kbn/zod'; +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 const getIndexMappingsTool = (): RegisteredTool< + typeof getIndexMappingsSchema, + GetIndexMappingsResult +> => { + return { + 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: [BuiltinTags.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 new file mode 100644 index 0000000000000..859da7c162dcf --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index.ts @@ -0,0 +1,15 @@ +/* + * 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 { 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 new file mode 100644 index 0000000000000..34d290ecd2b86 --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/index_explorer.ts @@ -0,0 +1,51 @@ +/* + * 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 { 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.'), + 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 const indexExplorerTool = (): RegisteredTool< + typeof indexExplorerSchema, + IndexExplorerResponse +> => { + return { + 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*'. + 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: [BuiltinTags.retrieval], + }, + }; +}; 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..279409670328d --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/list_indices.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. + */ + +import { z } from '@kbn/zod'; +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 + .string() + .optional() + .describe( + '(optional) pattern to filter indices by. Defaults to *. Leave empty to list all indices (recommended)' + ), +}); + +export const listIndicesTool = (): RegisteredTool => { + return { + 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: [BuiltinTags.retrieval], + }, + }; +}; diff --git a/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_search.ts b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_search.ts new file mode 100644 index 0000000000000..a875b42c1808d --- /dev/null +++ b/x-pack/platform/plugins/shared/onechat/server/tools/retrieval/nl_search.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 { z } from '@kbn/zod'; +import { BuiltinToolIds, BuiltinTags } from '@kbn/onechat-common'; +import type { RegisteredTool } from '@kbn/onechat-server'; +import { naturalLanguageSearch, NaturalLanguageSearchResponse } from '@kbn/onechat-genai-utils'; + +const searchDslSchema = z.object({ + 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 const naturalLanguageSearchTool = (): RegisteredTool< + typeof searchDslSchema, + NaturalLanguageSearchResponse +> => { + return { + 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(); + return naturalLanguageSearch({ + query, + context, + index, + model, + esClient: esClient.asCurrentUser, + }); + }, + meta: { + 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/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" ] } diff --git a/yarn.lock b/yarn.lock index 1d38de93a598a..f257d580c62e2 100644 --- a/yarn.lock +++ b/yarn.lock @@ -6583,6 +6583,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 ""