UNPKG

@langchain/core

Version:
529 lines (528 loc) 18.4 kB
import { BaseTracer } from "./base.js"; import { IterableReadableStream } from "../utils/stream.js"; import { AIMessageChunk } from "../messages/ai.js"; import { GenerationChunk } from "../outputs.js"; function assignName({ name, serialized, }) { if (name !== undefined) { return name; } if (serialized?.name !== undefined) { return serialized.name; } else if (serialized?.id !== undefined && Array.isArray(serialized?.id)) { return serialized.id[serialized.id.length - 1]; } return "Unnamed"; } export const isStreamEventsHandler = (handler) => handler.name === "event_stream_tracer"; /** * Class that extends the `BaseTracer` class from the * `langchain.callbacks.tracers.base` module. It represents a callback * handler that logs the execution of runs and emits `RunLog` instances to a * `RunLogStream`. */ export class EventStreamCallbackHandler extends BaseTracer { constructor(fields) { super({ _awaitHandler: true, ...fields }); Object.defineProperty(this, "autoClose", { enumerable: true, configurable: true, writable: true, value: true }); Object.defineProperty(this, "includeNames", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "includeTypes", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "includeTags", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "excludeNames", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "excludeTypes", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "excludeTags", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "runInfoMap", { enumerable: true, configurable: true, writable: true, value: new Map() }); Object.defineProperty(this, "tappedPromises", { enumerable: true, configurable: true, writable: true, value: new Map() }); Object.defineProperty(this, "transformStream", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "writer", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "receiveStream", { enumerable: true, configurable: true, writable: true, value: void 0 }); Object.defineProperty(this, "name", { enumerable: true, configurable: true, writable: true, value: "event_stream_tracer" }); Object.defineProperty(this, "lc_prefer_streaming", { enumerable: true, configurable: true, writable: true, value: true }); this.autoClose = fields?.autoClose ?? true; this.includeNames = fields?.includeNames; this.includeTypes = fields?.includeTypes; this.includeTags = fields?.includeTags; this.excludeNames = fields?.excludeNames; this.excludeTypes = fields?.excludeTypes; this.excludeTags = fields?.excludeTags; this.transformStream = new TransformStream(); this.writer = this.transformStream.writable.getWriter(); this.receiveStream = IterableReadableStream.fromReadableStream(this.transformStream.readable); } [Symbol.asyncIterator]() { return this.receiveStream; } async persistRun(_run) { // This is a legacy method only called once for an entire run tree // and is therefore not useful here } _includeRun(run) { const runTags = run.tags ?? []; let include = this.includeNames === undefined && this.includeTags === undefined && this.includeTypes === undefined; if (this.includeNames !== undefined) { include = include || this.includeNames.includes(run.name); } if (this.includeTypes !== undefined) { include = include || this.includeTypes.includes(run.runType); } if (this.includeTags !== undefined) { include = include || runTags.find((tag) => this.includeTags?.includes(tag)) !== undefined; } if (this.excludeNames !== undefined) { include = include && !this.excludeNames.includes(run.name); } if (this.excludeTypes !== undefined) { include = include && !this.excludeTypes.includes(run.runType); } if (this.excludeTags !== undefined) { include = include && runTags.every((tag) => !this.excludeTags?.includes(tag)); } return include; } async *tapOutputIterable(runId, outputStream) { const firstChunk = await outputStream.next(); if (firstChunk.done) { return; } const runInfo = this.runInfoMap.get(runId); // Run has finished, don't issue any stream events. // An example of this is for runnables that use the default // implementation of .stream(), which delegates to .invoke() // and calls .onChainEnd() before passing it to the iterator. if (runInfo === undefined) { yield firstChunk.value; return; } // Match format from handlers below function _formatOutputChunk(eventType, data) { if (eventType === "llm" && typeof data === "string") { return new GenerationChunk({ text: data }); } return data; } let tappedPromise = this.tappedPromises.get(runId); // if we are the first to tap, issue stream events if (tappedPromise === undefined) { let tappedPromiseResolver; tappedPromise = new Promise((resolve) => { tappedPromiseResolver = resolve; }); this.tappedPromises.set(runId, tappedPromise); try { const event = { event: `on_${runInfo.runType}_stream`, run_id: runId, name: runInfo.name, tags: runInfo.tags, metadata: runInfo.metadata, data: {}, }; await this.send({ ...event, data: { chunk: _formatOutputChunk(runInfo.runType, firstChunk.value), }, }, runInfo); yield firstChunk.value; for await (const chunk of outputStream) { // Don't yield tool and retriever stream events if (runInfo.runType !== "tool" && runInfo.runType !== "retriever") { await this.send({ ...event, data: { chunk: _formatOutputChunk(runInfo.runType, chunk), }, }, runInfo); } yield chunk; } } finally { // eslint-disable-next-line @typescript-eslint/no-non-null-assertion tappedPromiseResolver(); // Don't delete from the promises map to keep track of which runs have been tapped. } } else { // otherwise just pass through yield firstChunk.value; for await (const chunk of outputStream) { yield chunk; } } } async send(payload, run) { if (this._includeRun(run)) { await this.writer.write(payload); } } async sendEndEvent(payload, run) { const tappedPromise = this.tappedPromises.get(payload.run_id); if (tappedPromise !== undefined) { void tappedPromise.then(() => { void this.send(payload, run); }); } else { await this.send(payload, run); } } async onLLMStart(run) { const runName = assignName(run); const runType = run.inputs.messages !== undefined ? "chat_model" : "llm"; const runInfo = { tags: run.tags ?? [], metadata: run.extra?.metadata ?? {}, name: runName, runType, inputs: run.inputs, }; this.runInfoMap.set(run.id, runInfo); const eventName = `on_${runType}_start`; await this.send({ event: eventName, data: { input: run.inputs, }, name: runName, tags: run.tags ?? [], run_id: run.id, metadata: run.extra?.metadata ?? {}, }, runInfo); } async onLLMNewToken(run, token, // eslint-disable-next-line @typescript-eslint/no-explicit-any kwargs) { const runInfo = this.runInfoMap.get(run.id); let chunk; let eventName; if (runInfo === undefined) { throw new Error(`onLLMNewToken: Run ID ${run.id} not found in run map.`); } // Top-level streaming events are covered by tapOutputIterable if (this.runInfoMap.size === 1) { return; } if (runInfo.runType === "chat_model") { eventName = "on_chat_model_stream"; if (kwargs?.chunk === undefined) { chunk = new AIMessageChunk({ content: token, id: `run-${run.id}` }); } else { chunk = kwargs.chunk.message; } } else if (runInfo.runType === "llm") { eventName = "on_llm_stream"; if (kwargs?.chunk === undefined) { chunk = new GenerationChunk({ text: token }); } else { chunk = kwargs.chunk; } } else { throw new Error(`Unexpected run type ${runInfo.runType}`); } await this.send({ event: eventName, data: { chunk, }, run_id: run.id, name: runInfo.name, tags: runInfo.tags, metadata: runInfo.metadata, }, runInfo); } async onLLMEnd(run) { const runInfo = this.runInfoMap.get(run.id); this.runInfoMap.delete(run.id); let eventName; if (runInfo === undefined) { throw new Error(`onLLMEnd: Run ID ${run.id} not found in run map.`); } const generations = run.outputs?.generations; // eslint-disable-next-line @typescript-eslint/no-explicit-any let output; if (runInfo.runType === "chat_model") { for (const generation of generations ?? []) { if (output !== undefined) { break; } output = generation[0]?.message; } eventName = "on_chat_model_end"; } else if (runInfo.runType === "llm") { output = { generations: generations?.map((generation) => { return generation.map((chunk) => { return { text: chunk.text, generationInfo: chunk.generationInfo, }; }); }), llmOutput: run.outputs?.llmOutput ?? {}, }; eventName = "on_llm_end"; } else { throw new Error(`onLLMEnd: Unexpected run type: ${runInfo.runType}`); } await this.sendEndEvent({ event: eventName, data: { output, input: runInfo.inputs, }, run_id: run.id, name: runInfo.name, tags: runInfo.tags, metadata: runInfo.metadata, }, runInfo); } async onChainStart(run) { const runName = assignName(run); const runType = run.run_type ?? "chain"; const runInfo = { tags: run.tags ?? [], metadata: run.extra?.metadata ?? {}, name: runName, runType: run.run_type, }; let eventData = {}; // Workaround Runnable core code not sending input when transform streaming. if (run.inputs.input === "" && Object.keys(run.inputs).length === 1) { eventData = {}; runInfo.inputs = {}; } else if (run.inputs.input !== undefined) { eventData.input = run.inputs.input; runInfo.inputs = run.inputs.input; } else { eventData.input = run.inputs; runInfo.inputs = run.inputs; } this.runInfoMap.set(run.id, runInfo); await this.send({ event: `on_${runType}_start`, data: eventData, name: runName, tags: run.tags ?? [], run_id: run.id, metadata: run.extra?.metadata ?? {}, }, runInfo); } async onChainEnd(run) { const runInfo = this.runInfoMap.get(run.id); this.runInfoMap.delete(run.id); if (runInfo === undefined) { throw new Error(`onChainEnd: Run ID ${run.id} not found in run map.`); } const eventName = `on_${run.run_type}_end`; const inputs = run.inputs ?? runInfo.inputs ?? {}; const outputs = run.outputs?.output ?? run.outputs; const data = { output: outputs, input: inputs, }; if (inputs.input && Object.keys(inputs).length === 1) { data.input = inputs.input; runInfo.inputs = inputs.input; } await this.sendEndEvent({ event: eventName, data, run_id: run.id, name: runInfo.name, tags: runInfo.tags, metadata: runInfo.metadata ?? {}, }, runInfo); } async onToolStart(run) { const runName = assignName(run); const runInfo = { tags: run.tags ?? [], metadata: run.extra?.metadata ?? {}, name: runName, runType: "tool", inputs: run.inputs ?? {}, }; this.runInfoMap.set(run.id, runInfo); await this.send({ event: "on_tool_start", data: { input: run.inputs ?? {}, }, name: runName, run_id: run.id, tags: run.tags ?? [], metadata: run.extra?.metadata ?? {}, }, runInfo); } async onToolEnd(run) { const runInfo = this.runInfoMap.get(run.id); this.runInfoMap.delete(run.id); if (runInfo === undefined) { throw new Error(`onToolEnd: Run ID ${run.id} not found in run map.`); } if (runInfo.inputs === undefined) { throw new Error(`onToolEnd: Run ID ${run.id} is a tool call, and is expected to have traced inputs.`); } const output = run.outputs?.output === undefined ? run.outputs : run.outputs.output; await this.sendEndEvent({ event: "on_tool_end", data: { output, input: runInfo.inputs, }, run_id: run.id, name: runInfo.name, tags: runInfo.tags, metadata: runInfo.metadata, }, runInfo); } async onRetrieverStart(run) { const runName = assignName(run); const runType = "retriever"; const runInfo = { tags: run.tags ?? [], metadata: run.extra?.metadata ?? {}, name: runName, runType, inputs: { query: run.inputs.query, }, }; this.runInfoMap.set(run.id, runInfo); await this.send({ event: "on_retriever_start", data: { input: { query: run.inputs.query, }, }, name: runName, tags: run.tags ?? [], run_id: run.id, metadata: run.extra?.metadata ?? {}, }, runInfo); } async onRetrieverEnd(run) { const runInfo = this.runInfoMap.get(run.id); this.runInfoMap.delete(run.id); if (runInfo === undefined) { throw new Error(`onRetrieverEnd: Run ID ${run.id} not found in run map.`); } await this.sendEndEvent({ event: "on_retriever_end", data: { output: run.outputs?.documents ?? run.outputs, input: runInfo.inputs, }, run_id: run.id, name: runInfo.name, tags: runInfo.tags, metadata: runInfo.metadata, }, runInfo); } // eslint-disable-next-line @typescript-eslint/no-explicit-any async handleCustomEvent(eventName, data, runId) { const runInfo = this.runInfoMap.get(runId); if (runInfo === undefined) { throw new Error(`handleCustomEvent: Run ID ${runId} not found in run map.`); } await this.send({ event: "on_custom_event", run_id: runId, name: eventName, tags: runInfo.tags, metadata: runInfo.metadata, data, }, runInfo); } async finish() { const pendingPromises = [...this.tappedPromises.values()]; void Promise.all(pendingPromises).finally(() => { void this.writer.close(); }); } }