[onechat] factorize chat agent execution to runChatAgent function (elastic#222921)

## Summary

Introduce a static `runChatAgent` function to allow re-using the agent's
workflow from anywhere


### Example

```ts
    const completedRound = await runChatAgent(
      {
        nextInput,
        conversation,
        agentGraphName: defaultAgentGraphName,
        runId,
        onEvent: (event) => {
          events.emit(event);
        },
        tools: toolProvider,
      },
      {
        logger,
        runner,
        request,
        modelProvider,
      }
    );
```

Author: pgayvallet

x-pack/platform/packages/shared/onechat/onechat-server/agents/provider.ts

@@ -54,7 +54,7 @@ export interface AgentHandlerContext {
    */
   modelProvider: ModelProvider;
   /**
-   * Tool provider that should be used to list of execute tools.
+   * Tool provider that can be used to list or execute tools.
    */
   toolProvider: ToolProvider;
   /**

x-pack/platform/packages/shared/onechat/onechat-server/src/tools.ts

@@ -100,6 +100,10 @@ export interface ToolHandlerContext {
    * Can be used to access the inference APIs or chatModel.
    */
   modelProvider: ModelProvider;
+  /**
+   * Tool provider that can be used to list or execute tools.
+   */
+  toolProvider: ToolProvider;
   /**
    * Onechat runner scoped to the current execution.
    * Can be used to run other workchat primitive as part of the tool execution.

x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/graph.ts

@@ -12,14 +12,16 @@ 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 { withSystemPrompt } from './system_prompt';
+import { withSystemPrompt, defaultSystemPrompt } from './system_prompt';
 
 export const createAgentGraph = async ({
   chatModel,
   tools,
+  systemPrompt = defaultSystemPrompt,
 }: {
   chatModel: InferenceChatModel;
   tools: StructuredTool[];
+  systemPrompt?: string;
   logger: Logger;
 }) => {
   const StateAnnotation = Annotation.Root({
@@ -44,6 +46,7 @@ export const createAgentGraph = async ({
   const callModel = async (state: typeof StateAnnotation.State) => {
     const response = await model.invoke(
       await withSystemPrompt({
+        systemPrompt,
         messages: [...state.initialMessages, ...state.addedMessages],
       })
     );
@@ -68,7 +71,7 @@ export const createAgentGraph = async ({
     };
   };
 
-  // note: the node names are used in the event convertion logic, they should not be changed
+  // note: the node names are used in the event convertion logic, they should *not* be changed
   const graph = new StateGraph(StateAnnotation)
     .addNode('agent', callModel)
     .addNode('tools', toolHandler)

x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/handler.ts

@@ -5,14 +5,9 @@
  * 2.0.
  */
 
-import { from, filter, shareReplay, firstValueFrom, map } from 'rxjs';
 import type { Logger } from '@kbn/logging';
-import { StreamEvent } from '@langchain/core/tracers/log_stream';
-import { isRoundCompleteEvent } from '@kbn/onechat-common';
 import type { ConversationalAgentHandlerFn } from '@kbn/onechat-server';
-import { providerToLangchainTools, conversationLangchainMessages } from './utils';
-import { createAgentGraph } from './graph';
-import { convertGraphEvents, addRoundCompleteEvent } from './convert_graph_events';
+import { runChatAgent } from './run_chat_agent';
 
 export interface CreateConversationalAgentHandlerParams {
   logger: Logger;
@@ -30,54 +25,29 @@ export const createHandler = ({
     { agentParams: { nextInput, conversation = [] }, runId },
     { request, modelProvider, toolProvider, events, runner }
   ) => {
-    const model = await modelProvider.getDefaultModel();
-    const tools = await providerToLangchainTools({ request, toolProvider, runner, logger });
-    const initialMessages = conversationLangchainMessages({
-      nextInput,
-      previousRounds: conversation,
-    });
-    const agentGraph = await createAgentGraph({ logger, chatModel: model.chatModel, tools });
-
-    const eventStream = agentGraph.streamEvents(
-      { initialMessages },
+    const completedRound = await runChatAgent(
       {
-        version: 'v2',
-        runName: defaultAgentGraphName,
-        metadata: {
-          graphName: defaultAgentGraphName,
-          runId,
+        nextInput,
+        conversation,
+        agentGraphName: defaultAgentGraphName,
+        runId,
+        onEvent: (event) => {
+          events.emit(event);
         },
-        recursionLimit: 10,
-        callbacks: [],
+        tools: toolProvider,
+      },
+      {
+        logger,
+        runner,
+        request,
+        modelProvider,
       }
     );
 
-    const events$ = from(eventStream).pipe(
-      filter(isStreamEvent),
-      convertGraphEvents({ graphName: defaultAgentGraphName, runName: defaultAgentGraphName }),
-      addRoundCompleteEvent({ userInput: nextInput }),
-      shareReplay()
-    );
-
-    events$.subscribe((event) => {
-      events.emit(event);
-    });
-
-    const round = await firstValueFrom(
-      events$.pipe(
-        filter(isRoundCompleteEvent),
-        map((event) => event.data.round)
-      )
-    );
-
     return {
       result: {
-        round,
+        round: completedRound,
       },
     };
   };
 };
-
-const isStreamEvent = (input: any): input is StreamEvent => {
-  return 'event' in input;
-};

x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/index.ts

@@ -6,3 +6,4 @@
  */
 
 export { createDefaultAgentProvider } from './provider';
+export { runChatAgent } from './run_chat_agent';

x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/run_chat_agent.ts

@@ -0,0 +1,144 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { v4 as uuidv4 } from 'uuid';
+import { Observable, from, filter, shareReplay, firstValueFrom, map } 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 {
+  RoundInput,
+  ConversationRound,
+  ChatAgentEvent,
+  isRoundCompleteEvent,
+} from '@kbn/onechat-common';
+import type {
+  ModelProvider,
+  ScopedRunner,
+  ExecutableTool,
+  ToolProvider,
+} from '@kbn/onechat-server';
+import { providerToLangchainTools, toLangchainTool, conversationLangchainMessages } from './utils';
+import { createAgentGraph } from './graph';
+import { convertGraphEvents, addRoundCompleteEvent } from './convert_graph_events';
+
+export interface RunChatAgentContext {
+  logger: Logger;
+  request: KibanaRequest;
+  modelProvider: ModelProvider;
+  runner: ScopedRunner;
+}
+
+export interface RunChatAgentParams {
+  /**
+   * The next message in this conversation that the agent should respond to.
+   */
+  nextInput: RoundInput;
+  /**
+   * Previous rounds of conversation.
+   */
+  conversation?: ConversationRound[];
+  /**
+   * Optional system prompt to override the default one.
+   */
+  systemPrompt?: string;
+  /**
+   * List of tools that will be exposed to the agent.
+   * Either a list of tools or a tool provider.
+   */
+  tools?: ToolProvider | ExecutableTool[];
+  /**
+   * In case of nested calls (e.g calling from a tool), allows to define the runId.
+   */
+  runId?: string;
+  /**
+   * Handler to react to the agent's events.
+   */
+  onEvent?: (event: ChatAgentEvent) => void;
+  /**
+   * Can be used to override the graph's name. Used for tracing.
+   */
+  agentGraphName?: string;
+}
+
+export type RunChatAgentFn = (
+  params: RunChatAgentParams,
+  context: RunChatAgentContext
+) => Promise<ConversationRound>;
+
+const defaultAgentGraphName = 'default-onechat-agent';
+
+const noopOnEvent = () => {};
+
+/**
+ * Create the handler function for the default onechat agent.
+ */
+export const runChatAgent: RunChatAgentFn = async (
+  {
+    nextInput,
+    conversation = [],
+    tools = [],
+    onEvent = noopOnEvent,
+    runId = uuidv4(),
+    systemPrompt,
+    agentGraphName = defaultAgentGraphName,
+  },
+  { logger, request, modelProvider }
+) => {
+  const model = await modelProvider.getDefaultModel();
+  const langchainTools = Array.isArray(tools)
+    ? tools.map((tool) => toLangchainTool({ tool, logger }))
+    : await providerToLangchainTools({ request, toolProvider: tools, logger });
+  const initialMessages = conversationLangchainMessages({
+    nextInput,
+    previousRounds: conversation,
+  });
+  const agentGraph = await createAgentGraph({
+    logger,
+    chatModel: model.chatModel,
+    tools: langchainTools,
+    systemPrompt,
+  });
+
+  const eventStream = agentGraph.streamEvents(
+    { initialMessages },
+    {
+      version: 'v2',
+      runName: agentGraphName,
+      metadata: {
+        graphName: agentGraphName,
+        runId,
+      },
+      recursionLimit: 10,
+      callbacks: [],
+    }
+  );
+
+  const events$ = from(eventStream).pipe(
+    filter(isStreamEvent),
+    convertGraphEvents({ graphName: agentGraphName, runName: agentGraphName }),
+    addRoundCompleteEvent({ userInput: nextInput }),
+    shareReplay()
+  );
+
+  events$.subscribe(onEvent);
+
+  return await extractRound(events$);
+};
+
+export const extractRound = async (events$: Observable<ChatAgentEvent>) => {
+  return await firstValueFrom(
+    events$.pipe(
+      filter(isRoundCompleteEvent),
+      map((event) => event.data.round)
+    )
+  );
+};
+
+const isStreamEvent = (input: any): input is StreamEvent => {
+  return 'event' in input;
+};

x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/system_prompt.ts

@@ -7,17 +7,25 @@
 
 import { BaseMessage, BaseMessageLike } from '@langchain/core/messages';
 
-const getSystemPrompt = () => {
-  return `You are a helpful chat assistant from the Elasticsearch company.
+export const defaultSystemPrompt =
+  'You are a helpful chat assistant from the Elasticsearch company.';
 
-  You have tools at your disposal that you can use to answer the user's question.
+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 = ({ messages }: { messages: BaseMessage[] }): BaseMessageLike[] => {
-  return [['system', getSystemPrompt()], ...messages];
+export const withSystemPrompt = ({
+  systemPrompt,
+  messages,
+}: {
+  systemPrompt: string;
+  messages: BaseMessage[];
+}): BaseMessageLike[] => {
+  return [['system', getFullSystemPrompt(systemPrompt)], ...messages];
 };

x-pack/platform/plugins/shared/onechat/server/services/agents/conversational/utils/tool_provider_to_langchain_tools.ts

@@ -9,44 +9,37 @@ import { StructuredTool, tool as toTool } from '@langchain/core/tools';
 import { Logger } from '@kbn/logging';
 import type { KibanaRequest } from '@kbn/core-http-server';
 import { toolDescriptorToIdentifier, toSerializedToolIdentifier } from '@kbn/onechat-common';
-import type { ToolProvider, ExecutableTool, ScopedRunner } from '@kbn/onechat-server';
+import type { ToolProvider, ExecutableTool } from '@kbn/onechat-server';
 
 export const providerToLangchainTools = async ({
   request,
   toolProvider,
   logger,
-  runner,
 }: {
   request: KibanaRequest;
   toolProvider: ToolProvider;
   logger: Logger;
-  runner: ScopedRunner;
 }): Promise<StructuredTool[]> => {
   const allTools = await toolProvider.list({ request });
   return Promise.all(
     allTools.map((tool) => {
-      return toLangchainTool({ tool, logger, runner });
+      return toLangchainTool({ tool, logger });
     })
   );
 };
 
 export const toLangchainTool = ({
   tool,
   logger,
-  runner,
 }: {
   tool: ExecutableTool;
-  runner: ScopedRunner;
   logger: Logger;
 }): StructuredTool => {
   const toolId = toolDescriptorToIdentifier(tool);
   return toTool(
     async (input) => {
       try {
-        const toolReturn = await runner.runTool({
-          toolId,
-          toolParams: input,
-        });
+        const toolReturn = await tool.execute({ toolParams: input });
         return JSON.stringify(toolReturn.result);
       } catch (e) {
         logger.warn(`error calling tool ${tool.name}: ${e.message}`);

x-pack/platform/plugins/shared/onechat/server/services/runner/run_agent.ts

@@ -11,6 +11,7 @@ import type {
   ConversationalAgentParams,
   RunAgentReturn,
 } from '@kbn/onechat-server';
+import { internalProviderToPublic } from '../tools/utils';
 import { createAgentEventEmitter, forkContextForAgentRun } from './utils';
 import type { RunnerManager } from './runner';
 
@@ -29,7 +30,10 @@ export const createAgentHandlerContext = <TParams = Record<string, unknown>>({
     esClient: elasticsearch.client.asScoped(request),
     modelProvider: modelProviderFactory({ request, defaultConnectorId }),
     runner: manager.getRunner(),
-    toolProvider: toolsService.registry.asPublicRegistry(),
+    toolProvider: internalProviderToPublic({
+      provider: toolsService.registry,
+      getRunner: manager.getRunner,
+    }),
     events: createAgentEventEmitter({ eventHandler: onEvent, context: manager.context }),
   };
 };