@insightfactory.ai/insightfactory-databricks-langgraph-tracer 1.0.0-dev.0

package/dist/databricks-tracer.d.ts

@@ -0,0 +1,235 @@
+/**
+ * Custom LangChain `BaseTracer` that emits the shared schema to Databricks
+ * MLflow via `@mlflow/core`.
+ *
+ * This is the TypeScript counterpart of the Python `DatabricksLangGraphTracer`
+ * and emits the **identical** shared trace/span schema
+ * (`schema/trace-schema.json`):
+ *
+ * - **Root run create** → trace **tags** (`source`, `langgraph.run_id` /
+ *   `graph_id` / `env` / `api_revision`, conditional `langgraph.thread_id`) and
+ *   trace **metadata** (`mlflow.trace.session` / `mlflow.trace.user` when a
+ *   thread / user is present).
+ * - **Child run create** → `langgraph.node` / `langgraph.step` on the live node
+ *   span (from LangGraph's run metadata).
+ * - **LLM run update (end)** → resolved `mlflow.llm.model` / `mlflow.llm.provider`,
+ *   the reserved `mlflow.chat.tokenUsage` (incl. the cache-read / cache-creation
+ *   slots), the non-reserved token-count attrs, the reserved 3-key
+ *   `mlflow.llm.cost` plus the non-reserved cache cost line-items, and the
+ *   OpenInference-style `llm.*` dashboard keys (`llm.model_name` /
+ *   `llm.model_provider` / `llm.usage.prompt_tokens_cost` /
+ *   `llm.usage.completion_tokens_cost` — the keys the Databricks experiment
+ *   Overview cost charts aggregate) — accumulating a per-root cost aggregate
+ *   (totals **and** a per-model cost/token rollup).
+ * - **Root run update (end), before `span.end()`** → the reserved 3-key
+ *   `mlflow.trace.cost` metadata rollup, the per-model `cost.by_model` tag
+ *   (JSON), and the `cost.unknown_model` tag.
+ *
+ * Uses `startSpan` with an explicit `parent` ({@link LiveSpan}) so nested runs
+ * keep hierarchy without relying on the OTEL active context (callbacks are not
+ * executed inside `withSpan`). {@link BaseTracer} invokes `persistRun` only for
+ * root runs; every run (root included) is closed from `onRunUpdate`.
+ */
+import { SpanType } from "@mlflow/core";
+import { BaseTracer, type Run } from "@langchain/core/tracers/base";
+import { type TokenUsage } from "./cost.js";
+/** Maps LangChain `run_type` strings to MLflow {@link SpanType} (default CHAIN). */
+export declare function mapLangChainRunTypeToMlflowSpanType(runType: string): SpanType;
+/** Best-effort JSON-serializable copy for span inputs/outputs (handles cycles). */
+export declare function serializeTracePayload(value: unknown): unknown;
+/**
+ * Multimodal content-reference externalization (issues #16, #21). Chat-model span
+ * inputs preserve multimodal image / PDF / file content parts, but the inline
+ * base64 payload of each is removed before recording: the part is replaced with a
+ * lightweight reference so the Unity Catalog trace tables never carry the multi-MB
+ * data URIs that pushed large invoice traces past the SQL inline read limit.
+ *
+ * By default a part becomes a `{ type, _omitted: true, bytes: N }` placeholder. A
+ * consumer can instead supply a {@link ContentRefResolver} mapping a part to its
+ * own reference object — for example, the Unity Catalog volume path the image
+ * was loaded from, so the original can be re-fetched at runtime. The resolver runs
+ * only on the stored copy of the inputs; the live message sent to the model is
+ * untouched (so prompt caching is unaffected).
+ *
+ * When `maxStringChars` is set, any over-long *plain-text* leaf — in span **inputs or
+ * outputs** — is replaced with a compact `{ _truncated: true, chars: N, bytes: M,
+ * preview: "…" }` placeholder before recording: large text, not just multimodal
+ * bytes, can also push a trace past the SQL inline read limit. The
+ * multimodal payload externalization above runs on inputs only (LLM messages);
+ * outputs are arbitrary chain / tool JSON, so only the text cap is applied to them
+ * (see {@link externalizeContentRefs}). Both run on a copy; the live values are never
+ * mutated.
+ */
+/** Context passed to a {@link ContentRefResolver} for one multimodal part. */
+export interface ContentPartContext {
+    /** Position within the message `content` array (the page number for a pure-image message). */
+    index: number;
+    /** Encoded length of the inline payload being externalized. */
+    bytes: number;
+    /**
+     * The run's metadata (`run.extra.metadata` — LangGraph / RunnableConfig metadata).
+     * This is a **shallow copy** of the run metadata, so adding or replacing top-level
+     * keys is harmless (it cannot corrupt the live run state shared with sibling/child
+     * spans); its nested values are still shared, so do not mutate those.
+     */
+    metadata: Record<string, unknown>;
+}
+/**
+ * Consumer hook: maps a multimodal content part (+ its context) to the reference
+ * object stored in the trace, or `undefined` / `null` to fall back to the default
+ * `{ type, _omitted: true, bytes: N }` placeholder.
+ */
+export type ContentRefResolver = (part: Record<string, unknown>, ctx: ContentPartContext) => Record<string, unknown> | undefined | null;
+/**
+ * Resolve the text cap: explicit option wins, else `DATABRICKS_TRACING_MAX_STRING_CHARS`.
+ * Returns `undefined` (cap disabled) when neither is set, or when the value is not a
+ * positive integer. A non-positive / unparseable env value is ignored (warned), not
+ * thrown, so a stray override never breaks tracing startup.
+ */
+export declare function resolveMaxStringChars(explicit?: number): number | undefined;
+/**
+ * Externalize multimodal content parts' inline payloads and cap text.
+ * Public entry point. Takes a **shallow copy** of `metadata` once (when payloads are
+ * externalized with a resolver) so a resolver mutating `ctx.metadata` cannot corrupt
+ * the live run metadata, which the caller passes by reference and which is shared with
+ * sibling / child spans (resilient by design — a faulty resolver must not affect
+ * tracing). Then walks via {@link externalizeWalk}.
+ *
+ * `externalizePayloads` controls the multimodal-part stripping:
+ * - **inputs** (`true`) — the #21 behaviour: inline image / PDF / file payloads in
+ *   message `content` arrays are replaced with references (inputs are LLM messages,
+ *   where the `content` / `data` / `base64` shapes are genuinely multimodal).
+ * - **outputs** (`false`) — only the text cap is applied. Span outputs are arbitrary
+ *   chain / tool return values; the multimodal probe matches a bare `data` / `base64`
+ *   string in any `content` list, so running it would strip legitimate non-binary
+ *   output to an `_omitted` placeholder. A genuinely huge inline payload in an output
+ *   is still bounded by the text cap when it is set.
+ */
+export declare function externalizeContentRefs(obj: unknown, resolver: ContentRefResolver | undefined, metadata: Record<string, unknown>, maxStringChars?: number, externalizePayloads?: boolean): unknown;
+/**
+ * Extract normalized token usage from an LLM run end payload, extended for the
+ * Anthropic `cache_creation` category to match the Python `_extract_usage`.
+ *
+ * Prefers LangChain normalized `usage_metadata`
+ * (`input_token_details.cache_read` / `.cache_creation`,
+ * `output_token_details.reasoning`); falls back to raw `response_metadata.usage`
+ * (`input_tokens_details.cached_tokens` / `cache_read_input_tokens`,
+ * `cache_creation_input_tokens`, `output_tokens_details.reasoning_tokens`).
+ * Missing cache/reasoning sub-fields default to `0`. Returns `undefined` when
+ * neither source is present.
+ */
+export declare function extractUsage(run: Run): TokenUsage | undefined;
+/**
+ * Derive the canonical LiteLLM catalog key for an LLM run, aligned with the
+ * Python `_extract_model`.
+ *
+ * Sources, in order: invocation params `model` (probed both top-level on
+ * `run.extra.invocation_params` and nested under `run.extra.metadata`), then
+ * `response_metadata.model_name` — each qualified via `qualify` (default
+ * {@link qualifyDeploymentName}; consumers with non-canonical deployment
+ * names supply a resolver via the tracer's `pricingKeyResolver` option).
+ */
+export declare function extractModel(run: Run, qualify?: (name: string) => string): string | undefined;
+export declare class DatabricksLangGraphTracer extends BaseTracer {
+    name: string;
+    private readonly source;
+    private readonly qualifyModel;
+    private readonly contentRefResolver;
+    private readonly maxStringChars;
+    private readonly spanMap;
+    private readonly traceAggregates;
+    private readonly rootByRunId;
+    constructor(options?: {
+        source?: string;
+        /**
+         * Maps a deployment/model name to its LiteLLM catalog key, for
+         * deployments whose names diverge from the catalog (e.g. a tenant
+         * alias, or third-party models like `kimi-*` priced under another
+         * provider). Return `undefined` to fall through to the default
+         * `azure/<name>` → `azure_ai/<name>` probe.
+         */
+        pricingKeyResolver?: (name: string) => string | undefined;
+        /**
+         * Maps each multimodal content part to the reference stored in the trace in
+         * place of its inline base64 (issues #16/#21). When omitted, parts are
+         * externalized to a `{ type, _omitted: true, bytes: N }` placeholder. Runs on
+         * a copy of the inputs only — the live message sent to the model is untouched.
+         */
+        contentRefResolver?: ContentRefResolver;
+        /**
+         * Opt into capping over-long *plain-text* leaves in span inputs / outputs to a
+         * `{ _truncated: true, … }` placeholder (large text, not just
+         * multimodal bytes, can push a trace past the SQL inline read limit). Falls
+         * back to `DATABRICKS_TRACING_MAX_STRING_CHARS`, else off. A distinct knob from
+         * the multimodal reference handling above.
+         */
+        maxStringChars?: number;
+    });
+    onRunCreate(run: Run): Promise<void>;
+    onRunUpdate(run: Run): Promise<void>;
+    /**
+     * Write the root trace tags + session/user metadata (Python
+     * `_write_root_trace_fields`). Always emits `source` + the `langgraph.run_id` /
+     * `graph_id` / `env` / `api_revision` tags. A thread id additionally produces
+     * the `langgraph.thread_id` tag and the reserved `mlflow.trace.session`
+     * metadata; a user identity produces the reserved `mlflow.trace.user` metadata.
+     */
+    private writeRootTraceFields;
+    /**
+     * Write `langgraph.node` / `langgraph.step` onto a child span. LangGraph stamps
+     * these onto every per-node run's metadata (including the node's nested LLM
+     * run), mirroring the Python `_write_child_langgraph_context`.
+     */
+    private writeChildLangGraphContext;
+    /**
+     * Write model / provider / token-usage / cost attrs on a live LLM span (Python
+     * `_enrich_llm_span`). Token usage is written whenever usage is extractable;
+     * model / provider whenever the model resolves. Cost needs both: the reserved
+     * 3-key `mlflow.llm.cost` plus the non-reserved cache-read / cache-creation
+     * cost line-items (Option B) are written and folded into the root aggregate; a
+     * resolved-but-unpriceable model is recorded for the `cost.unknown_model` tag.
+     * Whenever the model resolves with usage, the span is also folded into the
+     * per-model `cost.by_model` rollup (tokens always; cost when priceable).
+     *
+     * The OpenInference-style `llm.*` duplicates (`llm.model_name` /
+     * `llm.model_provider` / `llm.usage.prompt_tokens_cost` /
+     * `llm.usage.completion_tokens_cost`) mirror the MLflow-convention attributes
+     * because the Databricks experiment Overview dashboard aggregates cost from
+     * those keys via `variant_get` on the UC span table — without them, Cost
+     * Breakdown / Cost Over Time render $0.00. The cost values must
+     * stay numbers: `variant_get(..., 'DOUBLE')` returns NULL for strings.
+     */
+    private enrichLLMSpan;
+    /**
+     * Write token-usage span attributes. The reserved `mlflow.chat.tokenUsage`
+     * attribute carries MLflow's `TokenUsageKey` shape extended with the two cache
+     * slots — `{ input_tokens, output_tokens, total_tokens, cache_read_input_tokens,
+     * cache_creation_input_tokens }` — so they survive Databricks ingestion and
+     * render in the UI usage breakdown. The non-reserved `gen_ai.usage.*`
+     * token-count attrs mirror the Python `_write_token_usage`.
+     */
+    private writeTokenUsage;
+    /**
+     * Write the trace-level cost rollups + `cost.unknown_model` tag (Python
+     * `_write_root_trace_rollup`). `mlflow.trace.cost` is the reserved MLflow
+     * `CostKey` shape — exactly `{ input_cost, output_cost, total_cost }`;
+     * Databricks strips any extra key on ingestion. `total_cost` already accounts
+     * for the cached discount and cache-creation premium; the per-line-item split
+     * lives on the spans.
+     *
+     * `cost.by_model` is the non-reserved per-model breakdown: emitted
+     * as a JSON string under a **trace tag** (like `cost.unknown_model`) rather
+     * than a custom metadata key — tags are proven to survive Databricks UC
+     * ingestion and stay queryable from the `*_trace_unified` view. Each bucket
+     * carries cost and tokens; the per-model totals reconcile with
+     * `mlflow.trace.cost` / `mlflow.chat.tokenUsage` within rounding. Models are
+     * emitted in sorted order, mirroring the Python tracer (the JSON parses to the
+     * same structure; the serialized bytes differ — Python `json.dumps` inserts
+     * `", "` / `": "` separators and renders integral zero floats as `0.0`).
+     *
+     * Must run before the root `span.end()` (see {@link onRunUpdate}).
+     */
+    private writeRootTraceRollup;
+    protected persistRun(_run: Run): Promise<void>;
+}
+//# sourceMappingURL=databricks-tracer.d.ts.map

package/dist/databricks-tracer.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"databricks-tracer.d.ts","sourceRoot":"","sources":["../src/databricks-tracer.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA+BG;AACH,OAAO,EAGL,QAAQ,EAGT,MAAM,cAAc,CAAC;AACtB,OAAO,EAAE,UAAU,EAAE,KAAK,GAAG,EAAE,MAAM,8BAA8B,CAAC;AACpE,OAAO,EAKL,KAAK,UAAU,EAChB,MAAM,WAAW,CAAC;AAInB,oFAAoF;AACpF,wBAAgB,mCAAmC,CAAC,OAAO,EAAE,MAAM,GAAG,QAAQ,CAa7E;AAED,mFAAmF;AACnF,wBAAgB,qBAAqB,CAAC,KAAK,EAAE,OAAO,GAAG,OAAO,CAO7D;AAED;;;;;;;;;;;;;;;;;;;;;;GAsBG;AAEH,8EAA8E;AAC9E,MAAM,WAAW,kBAAkB;IACjC,8FAA8F;IAC9F,KAAK,EAAE,MAAM,CAAC;IACd,+DAA+D;IAC/D,KAAK,EAAE,MAAM,CAAC;IACd;;;;;OAKG;IACH,QAAQ,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;CACnC;AAED;;;;GAIG;AACH,MAAM,MAAM,kBAAkB,GAAG,CAC/B,IAAI,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,EAC7B,GAAG,EAAE,kBAAkB,KACpB,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAAG,SAAS,GAAG,IAAI,CAAC;AAkNhD;;;;;GAKG;AACH,wBAAgB,qBAAqB,CAAC,QAAQ,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,SAAS,CAgB3E;AAmBD;;;;;;;;;;;;;;;;;GAiBG;AACH,wBAAgB,sBAAsB,CACpC,GAAG,EAAE,OAAO,EACZ,QAAQ,EAAE,kBAAkB,GAAG,SAAS,EACxC,QAAQ,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,EACjC,cAAc,CAAC,EAAE,MAAM,EACvB,mBAAmB,UAAO,GACzB,OAAO,CAGT;AA8ED;;;;;;;;;;;GAWG;AACH,wBAAgB,YAAY,CAAC,GAAG,EAAE,GAAG,GAAG,UAAU,GAAG,SAAS,CAqE7D;AAED;;;;;;;;;GASG;AACH,wBAAgB,YAAY,CAC1B,GAAG,EAAE,GAAG,EACR,OAAO,GAAE,CAAC,IAAI,EAAE,MAAM,KAAK,MAA8B,GACxD,MAAM,GAAG,SAAS,CAsBpB;AA6DD,qBAAa,yBAA0B,SAAQ,UAAU;IACvD,IAAI,SAA0B;IAE9B,OAAO,CAAC,QAAQ,CAAC,MAAM,CAAS;IAChC,OAAO,CAAC,QAAQ,CAAC,YAAY,CAA2B;IACxD,OAAO,CAAC,QAAQ,CAAC,kBAAkB,CAAiC;IACpE,OAAO,CAAC,QAAQ,CAAC,cAAc,CAAqB;IACpD,OAAO,CAAC,QAAQ,CAAC,OAAO,CAA+B;IACvD,OAAO,CAAC,QAAQ,CAAC,eAAe,CAAqC;IACrE,OAAO,CAAC,QAAQ,CAAC,WAAW,CAA6B;gBAGvD,OAAO,GAAE;QACP,MAAM,CAAC,EAAE,MAAM,CAAC;QAChB;;;;;;WAMG;QACH,kBAAkB,CAAC,EAAE,CAAC,IAAI,EAAE,MAAM,KAAK,MAAM,GAAG,SAAS,CAAC;QAC1D;;;;;WAKG;QACH,kBAAkB,CAAC,EAAE,kBAAkB,CAAC;QACxC;;;;;;WAMG;QACH,cAAc,CAAC,EAAE,MAAM,CAAC;KACpB;IAYF,WAAW,CAAC,GAAG,EAAE,GAAG,GAAG,OAAO,CAAC,IAAI,CAAC;IA+CpC,WAAW,CAAC,GAAG,EAAE,GAAG,GAAG,OAAO,CAAC,IAAI,CAAC;IA6E1C;;;;;;OAMG;IACH,OAAO,CAAC,oBAAoB;IA+B5B;;;;OAIG;IACH,OAAO,CAAC,0BAA0B;IAYlC;;;;;;;;;;;;;;;;;OAiBG;IACH,OAAO,CAAC,aAAa;IA+DrB;;;;;;;OAOG;IACH,OAAO,CAAC,eAAe;IAoBvB;;;;;;;;;;;;;;;;;;;OAmBG;IACH,OAAO,CAAC,oBAAoB;cA8BZ,UAAU,CAAC,IAAI,EAAE,GAAG,GAAG,OAAO,CAAC,IAAI,CAAC;CAIrD"}