@visibe.ai/node 0.1.4

package/dist/esm/integrations/langchain.js

@@ -0,0 +1,269 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { randomUUID } from 'node:crypto';
+import { calculateCost } from '../utils';
+// Re-export the storage so openai.ts and langgraph.ts can share the same instance.
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export const activeLangChainStorage = new AsyncLocalStorage();
+// ---------------------------------------------------------------------------
+// LangChain token extraction
+// Different providers nest token usage in different locations.
+// Check in the order specified by the spec.
+// ---------------------------------------------------------------------------
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function extractTokenUsage(output) {
+    const usage = output?.llmOutput?.tokenUsage
+        ?? output?.llmOutput?.usage
+        ?? output?.generations?.[0]?.[0]?.generationInfo?.usage;
+    // Use ?? not || so token counts of 0 are preserved correctly.
+    return {
+        inputTokens: usage?.promptTokens ?? usage?.input_tokens ?? 0,
+        outputTokens: usage?.completionTokens ?? usage?.output_tokens ?? 0,
+    };
+}
+// ---------------------------------------------------------------------------
+// LangChainCallback
+// ---------------------------------------------------------------------------
+export class LangChainCallback {
+    nextSpanId() {
+        return `step_${++this.stepCounter}`;
+    }
+    constructor(options) {
+        // Required by @langchain/core v1+ for proper callback registration.
+        // Without `name`, ensureHandler() wraps via fromMethods() which drops prototype methods.
+        // Without `awaitHandlers`, callbacks run in a background queue (p-queue) and fire
+        // after model.invoke() returns — causing spans to be missed on flush/completeTrace.
+        this.name = 'visibe-langchain-callback';
+        this.awaitHandlers = true;
+        this.raiseError = false;
+        // Maps LangChain runId → our spanId so we can set parent_span_id.
+        this.runIdToSpanId = new Map();
+        // Tracks start times so we can compute durationMs.
+        this.pendingLLMCalls = new Map(); // runId → startMs
+        this.pendingToolCalls = new Map();
+        this.stepCounter = 0;
+        // Agents we have already emitted agent_start spans for.
+        this.seenAgents = new Set();
+        // Token / call accumulators — updated by handleLLMEnd, read by patchCompiledStateGraph
+        // and patchRunnableSequence to populate completeTrace totals.
+        this.totalInputTokens = 0;
+        this.totalOutputTokens = 0;
+        this.llmCallCount = 0;
+        this.visibe = options.visibe;
+        this.traceId = options.traceId;
+        this.agentName = options.agentName;
+    }
+    // ---------------------------------------------------------------------------
+    // LLM events
+    // ---------------------------------------------------------------------------
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleLLMStart(_llm, _messages, runId) {
+        this.pendingLLMCalls.set(runId, Date.now());
+    }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleLLMEnd(output, runId, parentRunId) {
+        const startMs = this.pendingLLMCalls.get(runId) ?? Date.now();
+        this.pendingLLMCalls.delete(runId);
+        const { inputTokens, outputTokens } = extractTokenUsage(output);
+        const gen = output?.generations?.[0]?.[0];
+        const model = gen?.generationInfo?.model ?? this.agentName;
+        const cost = calculateCost(model, inputTokens, outputTokens);
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const rawText = gen?.text ?? gen?.message?.content ?? '';
+        const outputText = typeof rawText === 'string' ? rawText : JSON.stringify(rawText);
+        const spanId = this.nextSpanId();
+        this.runIdToSpanId.set(runId, spanId);
+        const parentSpanId = parentRunId ? this.runIdToSpanId.get(parentRunId) : undefined;
+        const span = this.visibe.buildLLMSpan({
+            spanId,
+            parentSpanId,
+            agentName: this.agentName,
+            model,
+            status: 'success',
+            inputTokens,
+            outputTokens,
+            inputText: '', // LangChain doesn't surface the raw prompt here
+            outputText,
+            durationMs: Date.now() - startMs,
+        });
+        this.visibe.batcher.add(this.traceId, span);
+        // Update local accumulators (used by patchCompiledStateGraph / patchRunnableSequence).
+        this.totalInputTokens += inputTokens;
+        this.totalOutputTokens += outputTokens;
+        this.llmCallCount++;
+        // Notify track() accumulator if running inside a group tracker.
+        this._onLLMSpan?.(inputTokens, outputTokens, cost);
+    }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleLLMError(err, runId) {
+        this.pendingLLMCalls.delete(runId);
+        this.visibe.batcher.add(this.traceId, this.visibe.buildErrorSpan({
+            spanId: this.nextSpanId(),
+            errorType: err?.constructor?.name ?? 'Error',
+            errorMessage: err?.message ?? String(err),
+        }));
+    }
+    // ---------------------------------------------------------------------------
+    // Tool events
+    // ---------------------------------------------------------------------------
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleToolStart(_tool, input, runId) {
+        this.pendingToolCalls.set(runId, { startMs: Date.now(), inputText: input });
+    }
+    async handleToolEnd(output, runId, parentRunId) {
+        const pending = this.pendingToolCalls.get(runId);
+        this.pendingToolCalls.delete(runId);
+        const spanId = this.nextSpanId();
+        this.runIdToSpanId.set(runId, spanId);
+        const parentSpanId = parentRunId ? this.runIdToSpanId.get(parentRunId) : undefined;
+        const span = this.visibe.buildToolSpan({
+            spanId,
+            parentSpanId,
+            toolName: 'tool',
+            agentName: this.agentName,
+            status: 'success',
+            durationMs: pending ? Date.now() - pending.startMs : 0,
+            inputText: pending?.inputText ?? '',
+            outputText: output,
+        });
+        this.visibe.batcher.add(this.traceId, span);
+        this._onToolSpan?.();
+    }
+    async handleToolError(err, runId) {
+        this.pendingToolCalls.delete(runId);
+        this.visibe.batcher.add(this.traceId, this.visibe.buildErrorSpan({
+            spanId: this.nextSpanId(),
+            errorType: err?.constructor?.name ?? 'Error',
+            errorMessage: err?.message ?? String(err),
+        }));
+    }
+    // ---------------------------------------------------------------------------
+    // Chain events
+    // ---------------------------------------------------------------------------
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleChainStart(chain, _inputs, runId, parentRunId) {
+        // Emit an agent_start span the first time we see a named chain.
+        const chainName = chain?.id?.at(-1) ?? '';
+        if (chainName && !this.seenAgents.has(chainName)) {
+            this.seenAgents.add(chainName);
+            const spanId = this.nextSpanId();
+            this.runIdToSpanId.set(runId, spanId);
+            void parentRunId; // suppress unused warning
+        }
+    }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleChainEnd(_outputs, _runId) { }
+    async handleChainError(err, _runId) {
+        this.visibe.batcher.add(this.traceId, this.visibe.buildErrorSpan({
+            spanId: this.nextSpanId(),
+            errorType: err?.constructor?.name ?? 'Error',
+            errorMessage: err?.message ?? String(err),
+        }));
+    }
+}
+// ---------------------------------------------------------------------------
+// patchRunnableSequence — patches RunnableSequence so pipe() chains are
+// automatically instrumented.  Called from index.ts patchFramework().
+// ---------------------------------------------------------------------------
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function patchRunnableSequence(lcModule, visibe) {
+    const RunnableSequence = lcModule?.RunnableSequence;
+    if (!RunnableSequence)
+        return () => { };
+    const originalInvoke = RunnableSequence.prototype.invoke;
+    const originalStream = RunnableSequence.prototype.stream;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunnableSequence.prototype.invoke = async function (input, config) {
+        // If already inside a LangChain trace, pass through.
+        if (activeLangChainStorage.getStore() !== undefined) {
+            return originalInvoke.call(this, input, config);
+        }
+        const traceId = randomUUID();
+        const startedAt = new Date().toISOString();
+        const startMs = Date.now();
+        await visibe.apiClient.createTrace({
+            trace_id: traceId,
+            name: 'langchain',
+            framework: 'langchain',
+            started_at: startedAt,
+            ...(visibe.sessionId ? { session_id: visibe.sessionId } : {}),
+        });
+        const cb = new LangChainCallback({ visibe, traceId, agentName: 'langchain' });
+        let result;
+        let status = 'completed';
+        try {
+            result = await activeLangChainStorage.run(cb, () => originalInvoke.call(this, input, _mergeCallbacks(config, cb)));
+        }
+        catch (err) {
+            status = 'failed';
+            throw err;
+        }
+        finally {
+            visibe.batcher.flush();
+            await visibe.apiClient.completeTrace(traceId, {
+                status,
+                ended_at: new Date().toISOString(),
+                duration_ms: Date.now() - startMs,
+                total_tokens: cb.totalInputTokens + cb.totalOutputTokens,
+                total_input_tokens: cb.totalInputTokens,
+                total_output_tokens: cb.totalOutputTokens,
+            });
+        }
+        return result;
+    };
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    RunnableSequence.prototype.stream = async function* (input, config) {
+        if (activeLangChainStorage.getStore() !== undefined) {
+            yield* originalStream.call(this, input, config);
+            return;
+        }
+        const traceId = randomUUID();
+        const startedAt = new Date().toISOString();
+        const startMs = Date.now();
+        await visibe.apiClient.createTrace({
+            trace_id: traceId,
+            name: 'langchain',
+            framework: 'langchain',
+            started_at: startedAt,
+            ...(visibe.sessionId ? { session_id: visibe.sessionId } : {}),
+        });
+        const cb = new LangChainCallback({ visibe, traceId, agentName: 'langchain' });
+        let status = 'completed';
+        try {
+            // RunnableSequence.stream() is an async function (not async generator) returning
+            // a Promise<AsyncIterable>. activeLangChainStorage.run returns that Promise,
+            // so we must await before yield*.
+            const gen = await activeLangChainStorage.run(cb, () => originalStream.call(this, input, _mergeCallbacks(config, cb)));
+            yield* gen;
+        }
+        catch (err) {
+            status = 'failed';
+            throw err;
+        }
+        finally {
+            visibe.batcher.flush();
+            await visibe.apiClient.completeTrace(traceId, {
+                status,
+                ended_at: new Date().toISOString(),
+                duration_ms: Date.now() - startMs,
+                total_tokens: cb.totalInputTokens + cb.totalOutputTokens,
+                total_input_tokens: cb.totalInputTokens,
+                total_output_tokens: cb.totalOutputTokens,
+            });
+        }
+    };
+    return () => {
+        RunnableSequence.prototype.invoke = originalInvoke;
+        RunnableSequence.prototype.stream = originalStream;
+    };
+}
+// ---------------------------------------------------------------------------
+// Private helpers
+// ---------------------------------------------------------------------------
+// Merge our callback into an existing LangChain config object.
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function _mergeCallbacks(config, cb) {
+    if (!config)
+        return { callbacks: [cb] };
+    const existing = Array.isArray(config.callbacks) ? config.callbacks : [];
+    return { ...config, callbacks: [...existing, cb] };
+}

package/dist/esm/integrations/langgraph.js

@@ -0,0 +1,168 @@
+import { randomUUID } from 'node:crypto';
+import { LangChainCallback, activeLangChainStorage } from './langchain';
+// ---------------------------------------------------------------------------
+// LangGraphCallback
+// Extends LangChainCallback and adds node-level agent_start spans.
+// ---------------------------------------------------------------------------
+export class LangGraphCallback extends LangChainCallback {
+    constructor(options) {
+        super(options);
+        this.nodeNames = new Set(options.nodeNames ?? []);
+    }
+    // Override handleChainStart to emit agent_start for known graph nodes.
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async handleChainStart(chain, inputs, runId, parentRunId) {
+        // LangGraph may surface the node key in chain.name (run name) or chain.id[-1] (class name).
+        const chainName = chain?.name ?? chain?.id?.at(-1) ?? '';
+        if (chainName && this.nodeNames.has(chainName)) {
+            const spanId = this.nextSpanId();
+            this.runIdToSpanId.set(runId, spanId);
+            // Emit an agent_start span for this node.
+            // type MUST be exactly "agent_start" — the backend validates this string.
+            this.visibe.batcher.add(this.traceId, this.visibe.buildAgentStartSpan({
+                spanId,
+                agentName: chainName,
+            }));
+            // Don't call super — we've already set the runId mapping.
+            return;
+        }
+        await super.handleChainStart(chain, inputs, runId, parentRunId);
+    }
+}
+// ---------------------------------------------------------------------------
+// patchCompiledStateGraph
+// Patches the CompiledStateGraph class so every graph created after init()
+// is automatically instrumented.  Class-level patching is required — patching
+// a single instance only instruments that one instance.
+// ---------------------------------------------------------------------------
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function patchCompiledStateGraph(lgModule, visibe) {
+    const CompiledStateGraph = lgModule?.CompiledStateGraph;
+    if (!CompiledStateGraph)
+        return () => { };
+    const originalInvoke = CompiledStateGraph.prototype.invoke;
+    const originalStream = CompiledStateGraph.prototype.stream;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    CompiledStateGraph.prototype.invoke = async function (input, config) {
+        if (activeLangChainStorage.getStore() !== undefined) {
+            return originalInvoke.call(this, input, config);
+        }
+        const traceId = randomUUID();
+        const startedAt = new Date().toISOString();
+        const startMs = Date.now();
+        const graphName = this.name ?? 'langgraph';
+        await visibe.apiClient.createTrace({
+            trace_id: traceId,
+            name: graphName,
+            framework: 'langgraph',
+            started_at: startedAt,
+            ...(visibe.sessionId ? { session_id: visibe.sessionId } : {}),
+        });
+        // Collect node names from the graph's node registry.
+        const nodeNames = _extractNodeNames(this);
+        const cb = new LangGraphCallback({
+            visibe,
+            traceId,
+            agentName: graphName,
+            nodeNames,
+        });
+        let result;
+        let status = 'completed';
+        try {
+            result = await activeLangChainStorage.run(cb, () => originalInvoke.call(this, input, _mergeCallbacks(config, cb)));
+        }
+        catch (err) {
+            status = 'failed';
+            throw err;
+        }
+        finally {
+            visibe.batcher.flush();
+            await visibe.apiClient.completeTrace(traceId, {
+                status,
+                ended_at: new Date().toISOString(),
+                duration_ms: Date.now() - startMs,
+                total_tokens: cb.totalInputTokens + cb.totalOutputTokens,
+                total_input_tokens: cb.totalInputTokens,
+                total_output_tokens: cb.totalOutputTokens,
+            });
+        }
+        return result;
+    };
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    CompiledStateGraph.prototype.stream = async function* (input, config) {
+        if (activeLangChainStorage.getStore() !== undefined) {
+            // Pregel.stream() is a regular async function returning a Promise<AsyncIterable>.
+            // We must await it before yield*.
+            yield* (await originalStream.call(this, input, config));
+            return;
+        }
+        const traceId = randomUUID();
+        const startedAt = new Date().toISOString();
+        const startMs = Date.now();
+        const graphName = this.name ?? 'langgraph';
+        await visibe.apiClient.createTrace({
+            trace_id: traceId,
+            name: graphName,
+            framework: 'langgraph',
+            started_at: startedAt,
+            ...(visibe.sessionId ? { session_id: visibe.sessionId } : {}),
+        });
+        const nodeNames = _extractNodeNames(this);
+        const cb = new LangGraphCallback({
+            visibe,
+            traceId,
+            agentName: graphName,
+            nodeNames,
+        });
+        let status = 'completed';
+        try {
+            // activeLangChainStorage.run returns the callback's return value synchronously,
+            // which is a Promise<AsyncIterable> from Pregel.stream(). Await before yield*.
+            const gen = await activeLangChainStorage.run(cb, () => originalStream.call(this, input, _mergeCallbacks(config, cb)));
+            yield* gen;
+        }
+        catch (err) {
+            status = 'failed';
+            throw err;
+        }
+        finally {
+            visibe.batcher.flush();
+            await visibe.apiClient.completeTrace(traceId, {
+                status,
+                ended_at: new Date().toISOString(),
+                duration_ms: Date.now() - startMs,
+                total_tokens: cb.totalInputTokens + cb.totalOutputTokens,
+                total_input_tokens: cb.totalInputTokens,
+                total_output_tokens: cb.totalOutputTokens,
+            });
+        }
+    };
+    return () => {
+        CompiledStateGraph.prototype.invoke = originalInvoke;
+        CompiledStateGraph.prototype.stream = originalStream;
+    };
+}
+// ---------------------------------------------------------------------------
+// Private helpers
+// ---------------------------------------------------------------------------
+// Extract the set of node names from a compiled graph's internal node registry.
+// LangGraph stores nodes in `this.nodes` (a Map or plain object keyed by name).
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function _extractNodeNames(graph) {
+    try {
+        const nodes = graph.nodes;
+        if (nodes instanceof Map)
+            return Array.from(nodes.keys());
+        if (nodes && typeof nodes === 'object')
+            return Object.keys(nodes);
+    }
+    catch { /* ignore */ }
+    return [];
+}
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function _mergeCallbacks(config, cb) {
+    if (!config)
+        return { callbacks: [cb] };
+    const existing = Array.isArray(config.callbacks) ? config.callbacks : [];
+    return { ...config, callbacks: [...existing, cb] };
+}