add logs to component chains/agents

2026-06-25 17:01:04 +03:00 · 2023-07-11 01:53:22 +01:00
parent 2bcc2f90b8
commit eb19c206cf
23 changed files with 414 additions and 199 deletions
@@ -0,0 +1,180 @@
+import { BaseTracer, Run, BaseCallbackHandler } from 'langchain/callbacks'
+import { AgentAction, ChainValues } from 'langchain/schema'
+import { Logger } from 'winston'
+import { Server } from 'socket.io'
+
+interface AgentRun extends Run {
+    actions: AgentAction[]
+}
+
+function tryJsonStringify(obj: unknown, fallback: string) {
+    try {
+        return JSON.stringify(obj, null, 2)
+    } catch (err) {
+        return fallback
+    }
+}
+
+function elapsed(run: Run): string {
+    if (!run.end_time) return ''
+    const elapsed = run.end_time - run.start_time
+    if (elapsed < 1000) {
+        return `${elapsed}ms`
+    }
+    return `${(elapsed / 1000).toFixed(2)}s`
+}
+
+export class ConsoleCallbackHandler extends BaseTracer {
+    name = 'console_callback_handler' as const
+    logger: Logger
+
+    protected persistRun(_run: Run) {
+        return Promise.resolve()
+    }
+
+    constructor(logger: Logger) {
+        super()
+        this.logger = logger
+    }
+
+    // utility methods
+
+    getParents(run: Run) {
+        const parents: Run[] = []
+        let currentRun = run
+        while (currentRun.parent_run_id) {
+            const parent = this.runMap.get(currentRun.parent_run_id)
+            if (parent) {
+                parents.push(parent)
+                currentRun = parent
+            } else {
+                break
+            }
+        }
+        return parents
+    }
+
+    getBreadcrumbs(run: Run) {
+        const parents = this.getParents(run).reverse()
+        const string = [...parents, run]
+            .map((parent) => {
+                const name = `${parent.execution_order}:${parent.run_type}:${parent.name}`
+                return name
+            })
+            .join(' > ')
+        return string
+    }
+
+    // logging methods
+
+    onChainStart(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(`[chain/start] [${crumbs}] Entering Chain run with input: ${tryJsonStringify(run.inputs, '[inputs]')}`)
+    }
+
+    onChainEnd(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[chain/end] [${crumbs}] [${elapsed(run)}] Exiting Chain run with output: ${tryJsonStringify(run.outputs, '[outputs]')}`
+        )
+    }
+
+    onChainError(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[chain/error] [${crumbs}] [${elapsed(run)}] Chain run errored with error: ${tryJsonStringify(run.error, '[error]')}`
+        )
+    }
+
+    onLLMStart(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        const inputs = 'prompts' in run.inputs ? { prompts: (run.inputs.prompts as string[]).map((p) => p.trim()) } : run.inputs
+        this.logger.verbose(`[llm/start] [${crumbs}] Entering LLM run with input: ${tryJsonStringify(inputs, '[inputs]')}`)
+    }
+
+    onLLMEnd(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[llm/end] [${crumbs}] [${elapsed(run)}] Exiting LLM run with output: ${tryJsonStringify(run.outputs, '[response]')}`
+        )
+    }
+
+    onLLMError(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[llm/error] [${crumbs}] [${elapsed(run)}] LLM run errored with error: ${tryJsonStringify(run.error, '[error]')}`
+        )
+    }
+
+    onToolStart(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(`[tool/start] [${crumbs}] Entering Tool run with input: "${run.inputs.input?.trim()}"`)
+    }
+
+    onToolEnd(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(`[tool/end] [${crumbs}] [${elapsed(run)}] Exiting Tool run with output: "${run.outputs?.output?.trim()}"`)
+    }
+
+    onToolError(run: Run) {
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[tool/error] [${crumbs}] [${elapsed(run)}] Tool run errored with error: ${tryJsonStringify(run.error, '[error]')}`
+        )
+    }
+
+    onAgentAction(run: Run) {
+        const agentRun = run as AgentRun
+        const crumbs = this.getBreadcrumbs(run)
+        this.logger.verbose(
+            `[agent/action] [${crumbs}] Agent selected action: ${tryJsonStringify(
+                agentRun.actions[agentRun.actions.length - 1],
+                '[action]'
+            )}`
+        )
+    }
+}
+
+/**
+ * Custom chain handler class
+ */
+export class CustomChainHandler extends BaseCallbackHandler {
+    name = 'custom_chain_handler'
+    isLLMStarted = false
+    socketIO: Server
+    socketIOClientId = ''
+    skipK = 0 // Skip streaming for first K numbers of handleLLMStart
+    returnSourceDocuments = false
+
+    constructor(socketIO: Server, socketIOClientId: string, skipK?: number, returnSourceDocuments?: boolean) {
+        super()
+        this.socketIO = socketIO
+        this.socketIOClientId = socketIOClientId
+        this.skipK = skipK ?? this.skipK
+        this.returnSourceDocuments = returnSourceDocuments ?? this.returnSourceDocuments
+    }
+
+    handleLLMStart() {
+        if (this.skipK > 0) this.skipK -= 1
+    }
+
+    handleLLMNewToken(token: string) {
+        if (this.skipK === 0) {
+            if (!this.isLLMStarted) {
+                this.isLLMStarted = true
+                this.socketIO.to(this.socketIOClientId).emit('start', token)
+            }
+            this.socketIO.to(this.socketIOClientId).emit('token', token)
+        }
+    }
+
+    handleLLMEnd() {
+        this.socketIO.to(this.socketIOClientId).emit('end')
+    }
+
+    handleChainEnd(outputs: ChainValues): void | Promise<void> {
+        if (this.returnSourceDocuments) {
+            this.socketIO.to(this.socketIOClientId).emit('sourceDocuments', outputs?.sourceDocuments)
+        }
+    }
+}