[onechat] add tool progression events

x-pack/.i18nrc.json

@@ -109,7 +109,7 @@
     "xpack.observabilityLogsOverview": [
       "platform/packages/shared/logs-overview/src/components"
     ],
-    "xpack.onechat": ["platform/plugins/shared/onechat"],
+    "xpack.onechat": ["platform/plugins/shared/onechat", "platform/packages/shared/onechat"],
     "xpack.osquery": [
       "platform/plugins/shared/osquery"
     ],

x-pack/platform/packages/shared/onechat/onechat-common/chat/conversation.ts

@@ -39,10 +39,20 @@ export type ConversationRoundStepMixin<TType extends ConversationRoundStepType,
   type: TType;
 };
 
+/**
+ * Tool call progress which were emitted during the tool execution
+ */
+export interface ToolCallProgress {
+  /**
+   * The full text message
+   */
+  message: string;
+}
+
 /**
  * Represents a tool call with the corresponding result.
  */
-export interface ToolCallWithResult<T = ToolResult[]> {
+export interface ToolCallWithResult {
   /**
    * Id of the tool call, as returned by the LLM
    */
@@ -55,15 +65,19 @@ export interface ToolCallWithResult<T = ToolResult[]> {
    * Arguments the tool was called with.
    */
   params: Record<string, any>;
+  /**
+   * List of progress message which were send during that tool call
+   */
+  progression?: ToolCallProgress[];
   /**
    * Result of the tool
    */
-  results: T;
+  results: ToolResult[];
 }
 
-export type ToolCallStep<T = ToolResult[]> = ConversationRoundStepMixin<
+export type ToolCallStep = ConversationRoundStepMixin<
   ConversationRoundStepType.toolCall,
-  ToolCallWithResult<T>
+  ToolCallWithResult
 >;
 
 export const createToolCallStep = (toolCallWithResult: ToolCallWithResult): ToolCallStep => {
@@ -96,17 +110,17 @@ export const isReasoningStep = (step: ConversationRoundStep): step is ReasoningS
 /**
  * Defines all possible types for round steps.
  */
-export type ConversationRoundStep<T = ToolResult[]> = ToolCallStep<T> | ReasoningStep;
+export type ConversationRoundStep = ToolCallStep | ReasoningStep;
 
 /**
  * Represents a round in a conversation, containing all the information
  * related to this particular round.
  */
-export interface ConversationRound<T = ToolResult[]> {
+export interface ConversationRound {
   /** The user input that initiated the round */
   input: RoundInput;
   /** List of intermediate steps before the end result, such as tool calls */
-  steps: Array<ConversationRoundStep<T>>;
+  steps: ConversationRoundStep[];
   /** The final response from the assistant */
   response: AssistantResponse;
   /** when tracing is enabled, contains the traceId associated with this round */

x-pack/platform/packages/shared/onechat/onechat-common/chat/events.ts

@@ -11,6 +11,7 @@ import type { ConversationRound } from './conversation';
 
 export enum ChatEventType {
   toolCall = 'tool_call',
+  toolProgress = 'tool_progress',
   toolResult = 'tool_result',
   reasoning = 'reasoning',
   messageChunk = 'message_chunk',
@@ -39,6 +40,21 @@ export const isToolCallEvent = (event: OnechatEvent<string, any>): event is Tool
   return event.type === ChatEventType.toolCall;
 };
 
+// Tool progress
+
+export interface ToolProgressEventData {
+  tool_call_id: string;
+  message: string;
+}
+
+export type ToolProgressEvent = ChatEventBase<ChatEventType.toolProgress, ToolProgressEventData>;
+
+export const isToolProgressEvent = (
+  event: OnechatEvent<string, any>
+): event is ToolProgressEvent => {
+  return event.type === ChatEventType.toolProgress;
+};
+
 // Tool result
 
 export interface ToolResultEventData {
@@ -158,6 +174,7 @@ export const isConversationUpdatedEvent = (
  */
 export type ChatAgentEvent =
   | ToolCallEvent
+  | ToolProgressEvent
   | ToolResultEvent
   | ReasoningEvent
   | MessageChunkEvent

x-pack/platform/packages/shared/onechat/onechat-common/chat/index.ts

@@ -17,6 +17,7 @@ export {
   type ConversationRoundStep,
   type ReasoningStepData,
   type ReasoningStep,
+  type ToolCallProgress,
   ConversationRoundStepType,
   isToolCallStep,
   isReasoningStep,
@@ -32,6 +33,8 @@ export {
   type ChatAgentEvent,
   type ToolResultEvent,
   type ToolResultEventData,
+  type ToolProgressEvent,
+  type ToolProgressEventData,
   type ToolCallEvent,
   type ToolCallEventData,
   type ReasoningEvent,
@@ -44,6 +47,7 @@ export {
   type RoundCompleteEvent,
   isToolCallEvent,
   isToolResultEvent,
+  isToolProgressEvent,
   isReasoningEvent,
   isMessageChunkEvent,
   isMessageCompleteEvent,

x-pack/platform/packages/shared/onechat/onechat-common/index.ts

@@ -98,6 +98,8 @@ export {
   type ConversationUpdatedEvent,
   type ConversationUpdatedEventData,
   type ChatAgentEvent,
+  type ToolProgressEvent,
+  type ToolProgressEventData,
   type ToolResultEvent,
   type ToolResultEventData,
   type ToolCallEvent,
@@ -110,6 +112,7 @@ export {
   type MessageCompleteEvent,
   type RoundCompleteEventData,
   type RoundCompleteEvent,
+  type ToolCallProgress,
   isToolCallEvent,
   isToolResultEvent,
   isReasoningEvent,
@@ -118,4 +121,5 @@ export {
   isRoundCompleteEvent,
   isConversationCreatedEvent,
   isConversationUpdatedEvent,
+  isToolProgressEvent,
 } from './chat';

x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.test.ts

@@ -46,7 +46,6 @@ describe('extractToolReturn', () => {
           },
         },
       ],
-      runId: 'unknown',
     });
   });
 

x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/messages.ts

@@ -77,7 +77,6 @@ export const extractToolReturn = (message: ToolMessage): RunToolReturn => {
     if (content.startsWith('Error:')) {
       return {
         results: [{ type: ToolResultType.error, data: { message: content } }],
-        runId: 'unknown',
       };
     } else {
       throw new Error(`No artifact attached to tool message: ${JSON.stringify(message)}`);

x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/tools.ts

@@ -9,7 +9,16 @@ import type { StructuredTool } from '@langchain/core/tools';
 import { tool as toTool } from '@langchain/core/tools';
 import type { Logger } from '@kbn/logging';
 import type { KibanaRequest } from '@kbn/core-http-server';
-import type { ToolProvider, ExecutableTool, RunToolReturn } from '@kbn/onechat-server';
+import type { ChatAgentEvent } from '@kbn/onechat-common';
+import { ChatEventType } from '@kbn/onechat-common';
+import type {
+  AgentEventEmitterFn,
+  ExecutableTool,
+  OnechatToolEvent,
+  RunToolReturn,
+  ToolProvider,
+  ToolEventHandlerFn,
+} from '@kbn/onechat-server';
 import { ToolResultType } from '@kbn/onechat-common/tools/tool_result';
 import type { ToolCall } from './messages';
 
@@ -30,18 +39,20 @@ export const toolsToLangchain = async ({
   request,
   tools,
   logger,
+  sendEvent,
 }: {
   request: KibanaRequest;
   tools: ToolProvider | ExecutableTool[];
   logger: Logger;
+  sendEvent?: AgentEventEmitterFn;
 }): Promise<ToolsAndMappings> => {
   const allTools = Array.isArray(tools) ? tools : await tools.list({ request });
   const onechatToLangchainIdMap = createToolIdMappings(allTools);
 
   const convertedTools = await Promise.all(
     allTools.map((tool) => {
       const toolId = onechatToLangchainIdMap.get(tool.id);
-      return toolToLangchain({ tool, logger, toolId });
+      return toolToLangchain({ tool, logger, toolId, sendEvent });
     })
   );
 
@@ -83,25 +94,35 @@ export const toolToLangchain = ({
   tool,
   toolId,
   logger,
+  sendEvent,
 }: {
   tool: ExecutableTool;
   toolId?: string;
   logger: Logger;
+  sendEvent?: AgentEventEmitterFn;
 }): StructuredTool => {
   return toTool(
-    async (input): Promise<[string, RunToolReturn]> => {
+    async (input, config): Promise<[string, RunToolReturn]> => {
+      let onEvent: ToolEventHandlerFn | undefined;
+      if (sendEvent) {
+        const toolCallId = config.configurable?.tool_call_id ?? config.toolCall?.id ?? 'unknown';
+        const convertEvent = getToolEventConverter({ toolCallId });
+        onEvent = (event) => {
+          sendEvent(convertEvent(event));
+        };
+      }
+
       try {
         logger.debug(`Calling tool ${tool.id} with params: ${JSON.stringify(input, null, 2)}`);
-        const toolReturn = await tool.execute({ toolParams: input });
-        const content = JSON.stringify({ results: toolReturn.results }); // wrap in a results object to conform to bedrock format
+        const toolReturn = await tool.execute({ toolParams: input, onEvent });
+        const content = JSON.stringify({ results: toolReturn.results });
         logger.debug(`Tool ${tool.id} returned reply of length ${content.length}`);
         return [content, toolReturn];
       } catch (e) {
         logger.warn(`error calling tool ${tool.id}: ${e}`);
         logger.debug(e.stack);
 
         const errorToolReturn: RunToolReturn = {
-          runId: tool.id,
           results: [
             {
               type: ToolResultType.error,
@@ -141,3 +162,18 @@ function reverseMap<K, V>(map: Map<K, V>): Map<V, K> {
   }
   return reversed;
 }
+
+const getToolEventConverter = ({ toolCallId }: { toolCallId: string }) => {
+  return (toolEvent: OnechatToolEvent): ChatAgentEvent => {
+    if (toolEvent.type === ChatEventType.toolProgress) {
+      return {
+        type: ChatEventType.toolProgress,
+        data: {
+          ...toolEvent.data,
+          tool_call_id: toolCallId,
+        },
+      };
+    }
+    throw new Error(`Invalid tool call type ${toolEvent.type}`);
+  };
+};

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/search/graph.ts

@@ -10,7 +10,7 @@ import type { BaseMessage } from '@langchain/core/messages';
 import { isToolMessage } from '@langchain/core/messages';
 import { messagesStateReducer } from '@langchain/langgraph';
 import { ToolNode } from '@langchain/langgraph/prebuilt';
-import type { ScopedModel } from '@kbn/onechat-server';
+import type { ScopedModel, ToolEventEmitter } from '@kbn/onechat-server';
 import type { ElasticsearchClient, Logger } from '@kbn/core/server';
 import type { ToolResult } from '@kbn/onechat-common/tools';
 import { ToolResultType } from '@kbn/onechat-common/tools';
@@ -19,6 +19,7 @@ import { indexExplorer } from '../index_explorer';
 import { createNaturalLanguageSearchTool, createRelevanceSearchTool } from './inner_tools';
 import { getSearchPrompt } from './prompts';
 import type { SearchTarget } from './types';
+import { progressMessages } from './i18n';
 
 const StateAnnotation = Annotation.Root({
   // inputs
@@ -45,19 +46,23 @@ export const createSearchToolGraph = ({
   model,
   esClient,
   logger,
+  events,
 }: {
   model: ScopedModel;
   esClient: ElasticsearchClient;
   logger: Logger;
+  events?: ToolEventEmitter;
 }) => {
   const tools = [
-    createRelevanceSearchTool({ model, esClient }),
-    createNaturalLanguageSearchTool({ model, esClient }),
+    createRelevanceSearchTool({ model, esClient, events }),
+    createNaturalLanguageSearchTool({ model, esClient, events }),
   ];
 
   const toolNode = new ToolNode<typeof StateAnnotation.State.messages>(tools);
 
   const selectAndValidateIndex = async (state: StateType) => {
+    events?.reportProgress(progressMessages.selectingTarget());
+
     const explorerRes = await indexExplorer({
       nlQuery: state.nlQuery,
       indexPattern: state.targetPattern ?? '*',
@@ -69,6 +74,8 @@ export const createSearchToolGraph = ({
 
     if (explorerRes.resources.length > 0) {
       const selectedResource = explorerRes.resources[0];
+      events?.reportProgress(progressMessages.selectedTarget(selectedResource.name));
+
       return {
         indexIsValid: true,
         searchTarget: { type: selectedResource.type, name: selectedResource.name },
@@ -90,6 +97,7 @@ export const createSearchToolGraph = ({
   });
 
   const callSearchAgent = async (state: StateType) => {
+    events?.reportProgress(progressMessages.resolvingSearchStrategy());
     const response = await searchModel.invoke(
       getSearchPrompt({ nlQuery: state.nlQuery, searchTarget: state.searchTarget })
     );

x-pack/platform/packages/shared/onechat/onechat-genai-utils/tools/search/i18n.ts

@@ -0,0 +1,45 @@
+/*
+ * 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 { i18n } from '@kbn/i18n';
+
+export const progressMessages = {
+  selectingTarget: () => {
+    return i18n.translate('xpack.onechat.tools.search.progress.selectingTarget', {
+      defaultMessage: 'Selecting the best target for this query',
+    });
+  },
+  selectedTarget: (target: string) => {
+    return i18n.translate('xpack.onechat.tools.search.progress.selectedTarget', {
+      defaultMessage: 'Selected "{target}" as the next search target',
+      values: {
+        target,
+      },
+    });
+  },
+  resolvingSearchStrategy: () => {
+    return i18n.translate('xpack.onechat.tools.search.progress.searchStrategy', {
+      defaultMessage: 'Thinking about the search strategy to use',
+    });
+  },
+  performingRelevanceSearch: ({ term }: { term: string }) => {
+    return i18n.translate('xpack.onechat.tools.search.progress.performingRelevanceSearch', {
+      defaultMessage: 'Searching documents for "{term}"',
+      values: {
+        term,
+      },
+    });
+  },
+  performingNlSearch: ({ query }: { query: string }) => {
+    return i18n.translate('xpack.onechat.tools.search.progress.performingTextSearch', {
+      defaultMessage: 'Generating an ES|QL for "{query}"',
+      values: {
+        query,
+      },
+    });
+  },
+};