Add Detailed Streaming to the Tool Agent (#4155)

* Add Detailed Streaming to the Tool Agent * lint fix --------- Co-authored-by: Henry <hzj94@hotmail.com>
2026-06-28 19:00:59 +03:00 · 2025-04-02 13:14:35 -04:00
parent cb06df4584
commit 9957184680
2 changed files with 118 additions and 3 deletions
@@ -30,6 +30,7 @@ import { LangWatch, LangWatchSpan, LangWatchTrace, autoconvertTypedValues } from
 import { DataSource } from 'typeorm'
 import { ChatGenerationChunk } from '@langchain/core/outputs'
 import { AIMessageChunk } from '@langchain/core/messages'
+import { Serialized } from '@langchain/core/load/serializable'

 interface AgentRun extends Run {
    actions: AgentAction[]
@@ -1499,3 +1500,86 @@ export class AnalyticHandler {
        }
    }
 }
+
+/**
+ * Custom callback handler for streaming detailed intermediate information
+ * during agent execution, specifically tool invocation inputs and outputs.
+ */
+export class CustomStreamingHandler extends BaseCallbackHandler {
+    name = 'custom_streaming_handler'
+
+    private sseStreamer: IServerSideEventStreamer
+    private chatId: string
+
+    constructor(sseStreamer: IServerSideEventStreamer, chatId: string) {
+        super()
+        this.sseStreamer = sseStreamer
+        this.chatId = chatId
+    }
+
+    /**
+     * Handle the start of a tool invocation
+     */
+    async handleToolStart(tool: Serialized, input: string, runId: string, parentRunId?: string): Promise<void> {
+        if (!this.sseStreamer) return
+
+        const toolName = typeof tool === 'object' && tool.name ? tool.name : 'unknown-tool'
+        const toolInput = typeof input === 'string' ? input : JSON.stringify(input, null, 2)
+
+        // Stream the tool invocation details using the agent_trace event type for consistency
+        this.sseStreamer.streamCustomEvent(this.chatId, 'agent_trace', {
+            step: 'tool_start',
+            name: toolName,
+            input: toolInput,
+            runId,
+            parentRunId: parentRunId || null
+        })
+    }
+
+    /**
+     * Handle the end of a tool invocation
+     */
+    async handleToolEnd(output: string | object, runId: string, parentRunId?: string): Promise<void> {
+        if (!this.sseStreamer) return
+
+        const toolOutput = typeof output === 'string' ? output : JSON.stringify(output, null, 2)
+
+        // Stream the tool output details using the agent_trace event type for consistency
+        this.sseStreamer.streamCustomEvent(this.chatId, 'agent_trace', {
+            step: 'tool_end',
+            output: toolOutput,
+            runId,
+            parentRunId: parentRunId || null
+        })
+    }
+
+    /**
+     * Handle tool errors
+     */
+    async handleToolError(error: Error, runId: string, parentRunId?: string): Promise<void> {
+        if (!this.sseStreamer) return
+
+        // Stream the tool error details using the agent_trace event type for consistency
+        this.sseStreamer.streamCustomEvent(this.chatId, 'agent_trace', {
+            step: 'tool_error',
+            error: error.message,
+            runId,
+            parentRunId: parentRunId || null
+        })
+    }
+
+    /**
+     * Handle agent actions
+     */
+    async handleAgentAction(action: AgentAction, runId: string, parentRunId?: string): Promise<void> {
+        if (!this.sseStreamer) return
+
+        // Stream the agent action details using the agent_trace event type for consistency
+        this.sseStreamer.streamCustomEvent(this.chatId, 'agent_trace', {
+            step: 'agent_action',
+            action: JSON.stringify(action),
+            runId,
+            parentRunId: parentRunId || null
+        })
+    }
+}