@rollyjoely/mlflow-langchain 0.1.3

package/README.md

@@ -0,0 +1,42 @@
+# MLflow LangChain Integration
+
+Auto-instrumentation for [LangChain](https://js.langchain.com/) chat models with MLflow Tracing.
+
+## Installation
+
+```bash
+npm install mlflow-langchain mlflow-tracing @langchain/core
+```
+
+## Usage
+
+```typescript
+import * as mlflow from 'mlflow-tracing';
+import { tracedModel } from 'mlflow-langchain';
+import { ChatAnthropic } from '@langchain/anthropic';
+
+mlflow.init({
+  trackingUri: 'http://localhost:5000',
+  experimentId: '<experiment-id>',
+});
+
+const model = tracedModel(new ChatAnthropic({ model: 'claude-sonnet-4-5-20250514' }));
+
+// Both invoke() and stream() are automatically traced
+const result = await model.invoke([{ role: 'user', content: 'Hello!' }]);
+```
+
+## Supported Models
+
+- `ChatAnthropic` (`@langchain/anthropic`)
+- `ChatOpenAI` (`@langchain/openai`)
+- `ChatXAI` (`@langchain/xai`)
+- Any `BaseChatModel` subclass with `invoke()` / `stream()` methods
+
+## Features
+
+- Traces `invoke()` and `stream()` calls as LLM spans
+- Captures input messages and output content
+- Extracts token usage from `usage_metadata`
+- Auto-detects message format from model class name
+- Preserves tracing through `bindTools()` calls

package/dist/index.d.ts

@@ -0,0 +1,23 @@
+/**
+ * MLflow Tracing wrapper for LangChain BaseChatModel.
+ *
+ * Wraps LangChain chat models to produce well-formatted MLflow spans
+ * for invoke() and stream() calls. Works with all LangChain providers:
+ * ChatAnthropic, ChatOpenAI, ChatXAI, and any BaseChatModel subclass.
+ */
+/**
+ * Create a traced version of a LangChain BaseChatModel with MLflow tracing.
+ *
+ * Wraps `invoke()` and `stream()` to produce LLM spans with:
+ * - Input messages
+ * - Output content
+ * - Token usage from usage_metadata
+ * - Message format (auto-detected from model class name)
+ *
+ * Also wraps `bindTools()` so that models with bound tools remain traced.
+ *
+ * @param model - The LangChain BaseChatModel instance to trace
+ * @returns Traced model with the same interface
+ */
+export declare function tracedModel<T = any>(model: T): T;
+//# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA;;;;;;GAMG;AAsCH;;;;;;;;;;;;;GAaG;AACH,wBAAgB,WAAW,CAAC,CAAC,GAAG,GAAG,EAAE,KAAK,EAAE,CAAC,GAAG,CAAC,CAoChD"}

package/dist/index.js

@@ -0,0 +1,313 @@
+"use strict";
+/**
+ * MLflow Tracing wrapper for LangChain BaseChatModel.
+ *
+ * Wraps LangChain chat models to produce well-formatted MLflow spans
+ * for invoke() and stream() calls. Works with all LangChain providers:
+ * ChatAnthropic, ChatOpenAI, ChatXAI, and any BaseChatModel subclass.
+ */
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.tracedModel = tracedModel;
+const mlflow_tracing_1 = require("@rollyjoely/mlflow-tracing");
+const TRACED_METHODS = ['invoke', 'stream'];
+/**
+ * Known LangChain model class names mapped to MLflow message format strings.
+ * Used to set the MESSAGE_FORMAT span attribute for proper UI rendering.
+ */
+const MODEL_CLASS_TO_FORMAT = {
+    ChatAnthropic: 'langchain-anthropic',
+    ChatOpenAI: 'langchain-openai',
+    ChatXAI: 'langchain-openai',
+    ChatGoogleGenerativeAI: 'langchain-gemini',
+};
+/**
+ * Detect the message format from the model's class name.
+ * Falls back to 'langchain' for unknown model types.
+ */
+function detectMessageFormat(model) {
+    const className = model?.constructor?.name;
+    if (className && className in MODEL_CLASS_TO_FORMAT) {
+        return MODEL_CLASS_TO_FORMAT[className];
+    }
+    return 'langchain';
+}
+/**
+ * Create a traced version of a LangChain BaseChatModel with MLflow tracing.
+ *
+ * Wraps `invoke()` and `stream()` to produce LLM spans with:
+ * - Input messages
+ * - Output content
+ * - Token usage from usage_metadata
+ * - Message format (auto-detected from model class name)
+ *
+ * Also wraps `bindTools()` so that models with bound tools remain traced.
+ *
+ * @param model - The LangChain BaseChatModel instance to trace
+ * @returns Traced model with the same interface
+ */
+function tracedModel(model) {
+    if (!model || typeof model !== 'object') {
+        return model;
+    }
+    const messageFormat = detectMessageFormat(model);
+    return new Proxy(model, {
+        get(target, prop, receiver) {
+            const original = Reflect.get(target, prop, receiver);
+            if (typeof original === 'function') {
+                if (prop === 'invoke') {
+                    // eslint-disable-next-line @typescript-eslint/ban-types
+                    return wrapInvoke(original, target, messageFormat);
+                }
+                if (prop === 'stream') {
+                    // eslint-disable-next-line @typescript-eslint/ban-types
+                    return wrapStream(original, target, messageFormat);
+                }
+                if (prop === 'bindTools') {
+                    // Ensure bindTools returns a traced model
+                    return function (...args) {
+                        // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+                        const bound = original.apply(target, args);
+                        return tracedModel(bound);
+                    };
+                }
+                // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
+                return original.bind(target);
+            }
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+            return original;
+        },
+    });
+}
+/**
+ * Wrap the invoke() method with MLflow tracing.
+ * Creates an LLM span that captures inputs, outputs, and token usage.
+ */
+// eslint-disable-next-line @typescript-eslint/ban-types
+function wrapInvoke(fn, target, messageFormat) {
+    return function (...args) {
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+        return (0, mlflow_tracing_1.withSpan)(async (span) => {
+            span.setInputs(serializeInput(args[0]));
+            const result = await fn.apply(target, args);
+            span.setOutputs(serializeOutput(result));
+            try {
+                const usage = extractTokenUsage(result);
+                if (usage) {
+                    span.setAttribute(mlflow_tracing_1.SpanAttributeKey.TOKEN_USAGE, usage);
+                }
+            }
+            catch (error) {
+                console.debug('Error extracting token usage from LangChain response', error);
+            }
+            span.setAttribute(mlflow_tracing_1.SpanAttributeKey.MESSAGE_FORMAT, messageFormat);
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+            return result;
+        }, { name: 'ChatModel', spanType: mlflow_tracing_1.SpanType.LLM });
+    };
+}
+/**
+ * Wrap the stream() method with MLflow tracing.
+ * Creates an LLM span that wraps the async iterator, collecting chunks
+ * and recording the aggregated result when iteration completes.
+ */
+// eslint-disable-next-line @typescript-eslint/ban-types
+function wrapStream(fn, target, messageFormat) {
+    return function (...args) {
+        const inputs = serializeInput(args[0]);
+        // stream() returns an IterableReadableStream (async iterable).
+        // We need to wrap the async iterator to trace the full lifecycle.
+        const streamPromise = fn.apply(target, args);
+        // LangChain's stream() can be sync or async depending on version.
+        // Handle both cases by always wrapping in a promise.
+        if (streamPromise && typeof streamPromise.then === 'function') {
+            // Async: stream() returns a Promise<IterableReadableStream>
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+            return streamPromise.then((stream) => {
+                return wrapAsyncIterable(stream, inputs, messageFormat);
+            });
+        }
+        // Sync: stream() returns an IterableReadableStream directly
+        return wrapAsyncIterable(streamPromise, inputs, messageFormat);
+    };
+}
+/**
+ * Wrap an async iterable (stream result) with MLflow tracing.
+ * Returns a proxy that intercepts Symbol.asyncIterator to add span tracking.
+ */
+function wrapAsyncIterable(stream, inputs, messageFormat) {
+    let tracingClaimed = false;
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+    return new Proxy(stream, {
+        get(target, prop, receiver) {
+            const original = Reflect.get(target, prop, receiver);
+            if (prop === Symbol.asyncIterator) {
+                return function () {
+                    if (tracingClaimed) {
+                        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-call
+                        return target[Symbol.asyncIterator]();
+                    }
+                    tracingClaimed = true;
+                    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-argument
+                    return wrapStreamIterator(target[Symbol.asyncIterator](), inputs, messageFormat);
+                };
+            }
+            if (typeof original === 'function') {
+                // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
+                return original.bind(target);
+            }
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+            return original;
+        },
+    });
+}
+/**
+ * Wrap an async iterator with MLflow span tracking.
+ * Collects all chunks during iteration and records aggregated outputs on completion.
+ */
+async function* wrapStreamIterator(iterator, inputs, messageFormat) {
+    const parentSpan = (0, mlflow_tracing_1.getCurrentActiveSpan)();
+    const span = (0, mlflow_tracing_1.startSpan)({ name: 'ChatModel', spanType: mlflow_tracing_1.SpanType.LLM, parent: parentSpan ?? undefined });
+    span.setInputs(inputs);
+    const chunks = [];
+    let iterationError;
+    try {
+        while (true) {
+            const { value, done } = await iterator.next();
+            if (done) {
+                break;
+            }
+            chunks.push(value);
+            yield value;
+        }
+    }
+    catch (error) {
+        iterationError = error;
+        throw error;
+    }
+    finally {
+        if (iterationError) {
+            span.setAttribute(mlflow_tracing_1.SpanAttributeKey.MESSAGE_FORMAT, messageFormat);
+            span.setStatus(mlflow_tracing_1.SpanStatusCode.ERROR, iterationError.message);
+            span.end();
+        }
+        else {
+            try {
+                // Aggregate chunks using LangChain's concat pattern
+                const aggregated = aggregateChunks(chunks);
+                if (aggregated) {
+                    span.setOutputs(serializeOutput(aggregated));
+                    const usage = extractTokenUsage(aggregated);
+                    if (usage) {
+                        span.setAttribute(mlflow_tracing_1.SpanAttributeKey.TOKEN_USAGE, usage);
+                    }
+                }
+            }
+            catch (e) {
+                console.debug('Could not aggregate stream chunks', e);
+            }
+            span.setAttribute(mlflow_tracing_1.SpanAttributeKey.MESSAGE_FORMAT, messageFormat);
+            span.end();
+        }
+    }
+}
+/**
+ * Aggregate LangChain AIMessageChunks using their concat() method.
+ * This is the same pattern used by LangChain internally for combining streamed chunks.
+ */
+function aggregateChunks(chunks) {
+    if (chunks.length === 0)
+        return undefined;
+    if (chunks.length === 1)
+        return chunks[0];
+    try {
+        // LangChain AIMessageChunk implements concat()
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-return, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call
+        return chunks.slice(1).reduce((acc, chunk) => acc.concat(chunk), chunks[0]);
+    }
+    catch {
+        // Fallback: return last chunk if concat not available
+        return chunks[chunks.length - 1];
+    }
+}
+/**
+ * Serialize LangChain message input for span recording.
+ * Handles both array of BaseMessage objects and raw input formats.
+ */
+function serializeInput(input) {
+    if (!input)
+        return input;
+    // If input is an array of LangChain messages, serialize them
+    if (Array.isArray(input)) {
+        return input.map((msg) => {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+            if (msg && typeof msg === 'object' && msg.content !== undefined) {
+                return {
+                    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-assignment
+                    role: msg._getType?.() ?? msg.constructor?.name ?? 'unknown',
+                    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-assignment
+                    content: msg.content,
+                    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+                    ...(msg.name ? { name: msg.name } : {}),
+                    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+                    ...(msg.tool_calls ? { tool_calls: msg.tool_calls } : {}),
+                };
+            }
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+            return msg;
+        });
+    }
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+    return input;
+}
+/**
+ * Serialize LangChain message output for span recording.
+ * Extracts the key fields from AIMessage/AIMessageChunk.
+ */
+function serializeOutput(output) {
+    if (!output || typeof output !== 'object')
+        return output;
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+    if (output.content !== undefined) {
+        const serialized = {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-assignment
+            content: output.content,
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-assignment
+            response_metadata: output.response_metadata,
+        };
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+        if (output.tool_calls && output.tool_calls.length > 0) {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access
+            serialized.tool_calls = output.tool_calls;
+        }
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+        if (output.usage_metadata) {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access
+            serialized.usage_metadata = output.usage_metadata;
+        }
+        return serialized;
+    }
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+    return output;
+}
+/**
+ * Extract token usage from a LangChain response.
+ * LangChain provides usage_metadata with input_tokens/output_tokens fields.
+ */
+function extractTokenUsage(response) {
+    // LangChain standardized usage_metadata
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+    const usage = response?.usage_metadata;
+    if (!usage) {
+        return undefined;
+    }
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+    const inputTokens = usage.input_tokens ?? usage.inputTokens ?? 0;
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+    const outputTokens = usage.output_tokens ?? usage.outputTokens ?? 0;
+    const totalTokens = inputTokens + outputTokens;
+    return {
+        input_tokens: inputTokens,
+        output_tokens: outputTokens,
+        total_tokens: totalTokens,
+    };
+}

package/package.json

@@ -0,0 +1,52 @@
+{
+  "name": "@rollyjoely/mlflow-langchain",
+  "version": "0.1.3",
+  "description": "LangChain integration package for MLflow Tracing",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/mlflow/mlflow.git"
+  },
+  "homepage": "https://mlflow.org/",
+  "author": {
+    "name": "MLflow",
+    "url": "https://mlflow.org/"
+  },
+  "bugs": {
+    "url": "https://github.com/mlflow/mlflow/issues"
+  },
+  "license": "Apache-2.0",
+  "keywords": [
+    "mlflow",
+    "tracing",
+    "observability",
+    "opentelemetry",
+    "llm",
+    "langchain",
+    "javascript",
+    "typescript"
+  ],
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "scripts": {
+    "build": "tsc",
+    "test": "jest",
+    "lint": "eslint . --ext .ts --max-warnings 0",
+    "lint:fix": "eslint . --ext .ts --fix",
+    "format": "prettier --write .",
+    "format:check": "prettier --check ."
+  },
+  "peerDependencies": {
+    "@langchain/core": ">=0.3.0",
+    "@rollyjoely/mlflow-tracing": "^0.1.3"
+  },
+  "devDependencies": {
+    "jest": "^29.6.2",
+    "typescript": "^5.8.3"
+  },
+  "engines": {
+    "node": ">=18"
+  },
+  "files": [
+    "dist/"
+  ]
+}