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

package/dist/token-aggregate.d.ts

@@ -0,0 +1,29 @@
+/** Enriched 5-field trace-level token rollup (base 3 + cache-read/creation + reasoning). */
+export interface EnrichedTraceTokenUsage {
+    input_tokens: number;
+    output_tokens: number;
+    total_tokens: number;
+    cache_read_input_tokens: number;
+    cache_creation_input_tokens: number;
+    reasoning_tokens: number;
+}
+interface SpanLike {
+    spanId: string;
+    parentId: string | null;
+    attributes: Record<string, unknown> | undefined;
+}
+/**
+ * Replacement for `@mlflow/core`'s `aggregateUsageFromSpans` that returns the
+ * enriched rollup (cache-read / cache-creation / reasoning in addition to the
+ * base input/output/total sums). Reads `gen_ai.usage.input_tokens.cached`,
+ * `gen_ai.usage.input_tokens.cache_creation`, and
+ * `gen_ai.usage.output_tokens.reasoning` written by the tracer on LLM child
+ * spans. Preserves the SDK's double-count guard that skips a span whose parent
+ * also carries `TOKEN_USAGE`.
+ *
+ * Returns `null` when no span has `TOKEN_USAGE` — matching the SDK contract so
+ * the exporter's `if (aggregatedUsage)` guard still kicks in.
+ */
+export declare function enrichedAggregateUsageFromSpans(spans: readonly SpanLike[]): EnrichedTraceTokenUsage | null;
+export {};
+//# sourceMappingURL=token-aggregate.d.ts.map

package/dist/token-aggregate.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"token-aggregate.d.ts","sourceRoot":"","sources":["../src/token-aggregate.ts"],"names":[],"mappings":"AAoBA,4FAA4F;AAC5F,MAAM,WAAW,uBAAuB;IACtC,YAAY,EAAE,MAAM,CAAC;IACrB,aAAa,EAAE,MAAM,CAAC;IACtB,YAAY,EAAE,MAAM,CAAC;IACrB,uBAAuB,EAAE,MAAM,CAAC;IAChC,2BAA2B,EAAE,MAAM,CAAC;IACpC,gBAAgB,EAAE,MAAM,CAAC;CAC1B;AAED,UAAU,QAAQ;IAChB,MAAM,EAAE,MAAM,CAAC;IACf,QAAQ,EAAE,MAAM,GAAG,IAAI,CAAC;IACxB,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAAG,SAAS,CAAC;CACjD;AAMD;;;;;;;;;;;GAWG;AACH,wBAAgB,+BAA+B,CAC7C,KAAK,EAAE,SAAS,QAAQ,EAAE,GACzB,uBAAuB,GAAG,IAAI,CAsChC"}

package/dist/token-aggregate.js

@@ -0,0 +1,63 @@
+/**
+ * Enriched trace-level token rollup for the `@mlflow/core` SDK aggregator.
+ *
+ * This is a TypeScript-only enrichment with no Python parallel: the Python
+ * package relies on MLflow's native aggregation of the reserved per-span
+ * `mlflow.chat.tokenUsage` (which already carries the cache slots), whereas the
+ * `@mlflow/core` SDK's built-in `aggregateUsageFromSpans` hard-codes a 3-field
+ * trace rollup (input/output/total), dropping the cache + reasoning detail. We
+ * replace it (see `patchMlflowAggregator` in `index.ts`) so the Databricks
+ * "Usage breakdown" panel renders the `(cache read)` row and so the
+ * cache-creation / reasoning counts are queryable on the trace.
+ *
+ * NB: the enriched trace-level `mlflow.trace.tokenUsage` key it feeds is NOT
+ * part of the shared cross-language schema — the schema's token keys are all
+ * span-level. This module only enriches the SDK's own trace rollup.
+ */
+import { SpanAttrKey } from "./generated-keys.js";
+const TOKEN_USAGE = SpanAttrKey.TOKEN_USAGE; // "mlflow.chat.tokenUsage"
+function toNum(v) {
+    return typeof v === "number" && Number.isFinite(v) ? v : 0;
+}
+/**
+ * Replacement for `@mlflow/core`'s `aggregateUsageFromSpans` that returns the
+ * enriched rollup (cache-read / cache-creation / reasoning in addition to the
+ * base input/output/total sums). Reads `gen_ai.usage.input_tokens.cached`,
+ * `gen_ai.usage.input_tokens.cache_creation`, and
+ * `gen_ai.usage.output_tokens.reasoning` written by the tracer on LLM child
+ * spans. Preserves the SDK's double-count guard that skips a span whose parent
+ * also carries `TOKEN_USAGE`.
+ *
+ * Returns `null` when no span has `TOKEN_USAGE` — matching the SDK contract so
+ * the exporter's `if (aggregatedUsage)` guard still kicks in.
+ */
+export function enrichedAggregateUsageFromSpans(spans) {
+    const total = {
+        input_tokens: 0,
+        output_tokens: 0,
+        total_tokens: 0,
+        cache_read_input_tokens: 0,
+        cache_creation_input_tokens: 0,
+        reasoning_tokens: 0,
+    };
+    let hasUsage = false;
+    for (const span of spans) {
+        const tokenUsage = span.attributes?.[TOKEN_USAGE];
+        if (!tokenUsage)
+            continue;
+        if (span.parentId) {
+            const parent = spans.find((s) => s.spanId === span.parentId);
+            if (parent?.attributes?.[TOKEN_USAGE])
+                continue;
+        }
+        total.input_tokens += toNum(tokenUsage.input_tokens);
+        total.output_tokens += toNum(tokenUsage.output_tokens);
+        total.total_tokens += toNum(tokenUsage.total_tokens);
+        total.cache_read_input_tokens += toNum(span.attributes?.[SpanAttrKey.INPUT_TOKENS_CACHED]);
+        total.cache_creation_input_tokens += toNum(span.attributes?.[SpanAttrKey.INPUT_TOKENS_CACHE_CREATION]);
+        total.reasoning_tokens += toNum(span.attributes?.[SpanAttrKey.OUTPUT_TOKENS_REASONING]);
+        hasUsage = true;
+    }
+    return hasUsage ? total : null;
+}
+//# sourceMappingURL=token-aggregate.js.map

package/dist/token-aggregate.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"token-aggregate.js","sourceRoot":"","sources":["../src/token-aggregate.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;GAeG;AACH,OAAO,EAAE,WAAW,EAAE,MAAM,qBAAqB,CAAC;AAElD,MAAM,WAAW,GAAG,WAAW,CAAC,WAAW,CAAC,CAAC,2BAA2B;AAkBxE,SAAS,KAAK,CAAC,CAAU;IACvB,OAAO,OAAO,CAAC,KAAK,QAAQ,IAAI,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;AAC7D,CAAC;AAED;;;;;;;;;;;GAWG;AACH,MAAM,UAAU,+BAA+B,CAC7C,KAA0B;IAE1B,MAAM,KAAK,GAA4B;QACrC,YAAY,EAAE,CAAC;QACf,aAAa,EAAE,CAAC;QAChB,YAAY,EAAE,CAAC;QACf,uBAAuB,EAAE,CAAC;QAC1B,2BAA2B,EAAE,CAAC;QAC9B,gBAAgB,EAAE,CAAC;KACpB,CAAC;IACF,IAAI,QAAQ,GAAG,KAAK,CAAC;IAErB,KAAK,MAAM,IAAI,IAAI,KAAK,EAAE,CAAC;QACzB,MAAM,UAAU,GAAG,IAAI,CAAC,UAAU,EAAE,CAAC,WAAW,CAEnC,CAAC;QACd,IAAI,CAAC,UAAU;YAAE,SAAS;QAE1B,IAAI,IAAI,CAAC,QAAQ,EAAE,CAAC;YAClB,MAAM,MAAM,GAAG,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,MAAM,KAAK,IAAI,CAAC,QAAQ,CAAC,CAAC;YAC7D,IAAI,MAAM,EAAE,UAAU,EAAE,CAAC,WAAW,CAAC;gBAAE,SAAS;QAClD,CAAC;QAED,KAAK,CAAC,YAAY,IAAI,KAAK,CAAC,UAAU,CAAC,YAAY,CAAC,CAAC;QACrD,KAAK,CAAC,aAAa,IAAI,KAAK,CAAC,UAAU,CAAC,aAAa,CAAC,CAAC;QACvD,KAAK,CAAC,YAAY,IAAI,KAAK,CAAC,UAAU,CAAC,YAAY,CAAC,CAAC;QACrD,KAAK,CAAC,uBAAuB,IAAI,KAAK,CACpC,IAAI,CAAC,UAAU,EAAE,CAAC,WAAW,CAAC,mBAAmB,CAAC,CACnD,CAAC;QACF,KAAK,CAAC,2BAA2B,IAAI,KAAK,CACxC,IAAI,CAAC,UAAU,EAAE,CAAC,WAAW,CAAC,2BAA2B,CAAC,CAC3D,CAAC;QACF,KAAK,CAAC,gBAAgB,IAAI,KAAK,CAC7B,IAAI,CAAC,UAAU,EAAE,CAAC,WAAW,CAAC,uBAAuB,CAAC,CACvD,CAAC;QACF,QAAQ,GAAG,IAAI,CAAC;IAClB,CAAC;IAED,OAAO,QAAQ,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,IAAI,CAAC;AACjC,CAAC"}

package/dist/trace-metadata.d.ts

@@ -0,0 +1,14 @@
+/**
+ * Writes trace-level `tags` and/or `traceMetadata` on the in-memory trace for
+ * `traceId`. Both objects are mutated in place via `Object.assign` — writes must
+ * happen before the SDK flushes the trace to the exporter (the span processor
+ * pops the trace from the manager synchronously when the root span ends).
+ *
+ * Failures (missing trace, unexpected shape, etc.) are swallowed: the tracer
+ * must never throw from this path (resilient by design).
+ */
+export declare function writeTraceFields(traceId: string, fields: {
+    tags?: Record<string, string>;
+    metadata?: Record<string, string>;
+}): void;
+//# sourceMappingURL=trace-metadata.d.ts.map

package/dist/trace-metadata.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"trace-metadata.d.ts","sourceRoot":"","sources":["../src/trace-metadata.ts"],"names":[],"mappings":"AAUA;;;;;;;;GAQG;AACH,wBAAgB,gBAAgB,CAC9B,OAAO,EAAE,MAAM,EACf,MAAM,EAAE;IAAE,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;IAAC,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;CAAE,GAC3E,IAAI,CASN"}

package/dist/trace-metadata.js

@@ -0,0 +1,33 @@
+// Deliberate private-API wrapper around `@mlflow/core`'s `InMemoryTraceManager`.
+// `updateCurrentTrace` is a no-op from `startSpan`-style tracers (it keys off the
+// OTEL active span, which is not set when callers pass `parent` explicitly), so
+// trace-level tag/metadata writes have to reach into the in-memory manager. This
+// is the TypeScript parallel to the Python tracer's `_write_trace_fields`
+// (`InMemoryTraceManager.get_trace(trace_id)` + `.info.tags`/`.info.trace_metadata`
+// mutation). Centralized here so the rest of the tracer never imports from
+// `@mlflow/core/dist/...` paths directly.
+import { InMemoryTraceManager } from "@mlflow/core/dist/core/trace_manager.js";
+/**
+ * Writes trace-level `tags` and/or `traceMetadata` on the in-memory trace for
+ * `traceId`. Both objects are mutated in place via `Object.assign` — writes must
+ * happen before the SDK flushes the trace to the exporter (the span processor
+ * pops the trace from the manager synchronously when the root span ends).
+ *
+ * Failures (missing trace, unexpected shape, etc.) are swallowed: the tracer
+ * must never throw from this path (resilient by design).
+ */
+export function writeTraceFields(traceId, fields) {
+    try {
+        const trace = InMemoryTraceManager.getInstance().getTrace(traceId);
+        if (!trace)
+            return;
+        if (fields.tags)
+            Object.assign(trace.info.tags, fields.tags);
+        if (fields.metadata)
+            Object.assign(trace.info.traceMetadata, fields.metadata);
+    }
+    catch {
+        // never throw from a live run
+    }
+}
+//# sourceMappingURL=trace-metadata.js.map

package/dist/trace-metadata.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"trace-metadata.js","sourceRoot":"","sources":["../src/trace-metadata.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,kFAAkF;AAClF,gFAAgF;AAChF,iFAAiF;AACjF,0EAA0E;AAC1E,oFAAoF;AACpF,2EAA2E;AAC3E,0CAA0C;AAC1C,OAAO,EAAE,oBAAoB,EAAE,MAAM,yCAAyC,CAAC;AAE/E;;;;;;;;GAQG;AACH,MAAM,UAAU,gBAAgB,CAC9B,OAAe,EACf,MAA4E;IAE5E,IAAI,CAAC;QACH,MAAM,KAAK,GAAG,oBAAoB,CAAC,WAAW,EAAE,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC;QACnE,IAAI,CAAC,KAAK;YAAE,OAAO;QACnB,IAAI,MAAM,CAAC,IAAI;YAAE,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,IAAI,CAAC,IAAI,EAAE,MAAM,CAAC,IAAI,CAAC,CAAC;QAC7D,IAAI,MAAM,CAAC,QAAQ;YAAE,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,IAAI,CAAC,aAAa,EAAE,MAAM,CAAC,QAAQ,CAAC,CAAC;IAChF,CAAC;IAAC,MAAM,CAAC;QACP,8BAA8B;IAChC,CAAC;AACH,CAAC"}

package/dist/uc-export.d.ts

@@ -0,0 +1,48 @@
+/** Experiment tag carrying the UC destination (`catalog.schema.table_prefix`). */
+export declare const UC_DESTINATION_TAG = "mlflow.experiment.databricksTraceDestinationPath";
+/** Experiment tag naming the spans table outright (preferred over deriving). */
+export declare const UC_SPANS_TABLE_TAG = "mlflow.experiment.databricksTraceSpanStorageTable";
+/** Header naming the UC spans table for the OTLP ingest endpoint. */
+export declare const UC_TABLE_HEADER = "X-Databricks-UC-Table-Name";
+export interface UcDestination {
+    catalogName: string;
+    schemaName: string;
+    tablePrefix: string;
+    /** `catalog.schema.table_prefix` — the V4 location path segment. */
+    fullTablePrefix: string;
+    /** `catalog.schema.<prefix>_otel_spans` — the OTLP ingest header value. */
+    spansTableName: string;
+}
+/**
+ * Parse a `databricksTraceDestinationPath` tag value into a {@link UcDestination}.
+ * Throws when the value is not the expected three-part `catalog.schema.prefix`.
+ */
+export declare function parseUcDestinationPath(path: string): UcDestination;
+type HeadersProvider = () => Promise<Record<string, string>> | Record<string, string>;
+/**
+ * Fetch the experiment and resolve its UC destination from the
+ * `databricksTraceDestinationPath` tag. Returns `null` for classic
+ * (workspace-backed) experiments. SDK/network failures raise an actionable
+ * configuration error rather than a generic fetch failure, mirroring the
+ * Python package's hardened guard.
+ */
+export declare function resolveUcDestination(hostUrl: string, headersProvider: HeadersProvider, experimentId: string): Promise<UcDestination | null>;
+/**
+ * Databricks' OTLP ingest rejects batches over 10 MiB ("Batch exceeds maximum
+ * size of 10485760 bytes"). Stay comfortably under it — protobuf framing per
+ * chunk adds a little over the sum of span sizes.
+ */
+export declare const OTLP_MAX_BATCH_BYTES: number;
+/** Reset patch state. For testing only. */
+export declare function _resetUcExportForTests(): void;
+/**
+ * Patch `MlflowSpanExporter.exportTraceToBackend` with a UC-aware
+ * implementation. The destination is resolved once (memoized promise) from
+ * the experiment's tags on the first export — the SDK's client (and its
+ * auth headers) only exists at export time — and classic experiments
+ * delegate to the original implementation untouched. Resolution failures
+ * log loudly and fall back to classic export rather than dropping traces.
+ */
+export declare function patchExporterForUc(experimentId: string): void;
+export {};
+//# sourceMappingURL=uc-export.d.ts.map

package/dist/uc-export.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"uc-export.d.ts","sourceRoot":"","sources":["../src/uc-export.ts"],"names":[],"mappings":"AAiCA,kFAAkF;AAClF,eAAO,MAAM,kBAAkB,qDAAqD,CAAC;AAErF,gFAAgF;AAChF,eAAO,MAAM,kBAAkB,sDAAsD,CAAC;AAEtF,qEAAqE;AACrE,eAAO,MAAM,eAAe,+BAA+B,CAAC;AAM5D,MAAM,WAAW,aAAa;IAC5B,WAAW,EAAE,MAAM,CAAC;IACpB,UAAU,EAAE,MAAM,CAAC;IACnB,WAAW,EAAE,MAAM,CAAC;IACpB,oEAAoE;IACpE,eAAe,EAAE,MAAM,CAAC;IACxB,2EAA2E;IAC3E,cAAc,EAAE,MAAM,CAAC;CACxB;AAED;;;GAGG;AACH,wBAAgB,sBAAsB,CAAC,IAAI,EAAE,MAAM,GAAG,aAAa,CAgBlE;AAED,KAAK,eAAe,GAAG,MAAM,OAAO,CAAC,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC,GAAG,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;AAsCtF;;;;;;GAMG;AACH,wBAAsB,oBAAoB,CACxC,OAAO,EAAE,MAAM,EACf,eAAe,EAAE,eAAe,EAChC,YAAY,EAAE,MAAM,GACnB,OAAO,CAAC,aAAa,GAAG,IAAI,CAAC,CA+B/B;AAOD;;;;GAIG;AACH,eAAO,MAAM,oBAAoB,QAAkB,CAAC;AAsGpD,2CAA2C;AAC3C,wBAAgB,sBAAsB,IAAI,IAAI,CAE7C;AAED;;;;;;;GAOG;AACH,wBAAgB,kBAAkB,CAAC,YAAY,EAAE,MAAM,GAAG,IAAI,CAsD7D"}

package/dist/uc-export.js

@@ -0,0 +1,243 @@
+/**
+ * UC table-backed span export for Databricks MLflow experiments.
+ *
+ * `@mlflow/core@0.2.0` exports traces via `StartTraceV3` + artifact-store
+ * upload — a path UC table-backed experiments silently ignore (the calls
+ * succeed but the spans never reach the `<prefix>_otel_spans` Delta table and
+ * V4 `search_traces` returns nothing). This module mirrors the Python
+ * package's UC fix (mlflow_utils.py `_resolve_uc_trace_location`) plus the
+ * protocol MLflow >= 3.10 Python uses natively:
+ *
+ * - Spans: OTLP `ExportTraceServiceRequest` protobuf POSTed to
+ *   `/api/2.0/otel/v1/traces` with the `X-Databricks-UC-Table-Name` header
+ *   naming the fully-qualified `<catalog>.<schema>.<prefix>_otel_spans` table.
+ * - Trace info: `CreateTraceInfo` JSON POSTed to
+ *   `/api/4.0/mlflow/traces/<location>/<otelTraceId>/info` with a
+ *   `UC_TABLE_PREFIX` trace location and `mlflow.trace_schema.version = "4"`.
+ *
+ * Like the aggregator patch in `index.ts`, the seam is a `createRequire`
+ * monkeypatch of the SDK's `MlflowSpanExporter` — the exporter pops the trace
+ * from `InMemoryTraceManager` synchronously inside root-span `end()`, so
+ * there is no public post-export hook to intercept instead.
+ *
+ * UC detection is resolved lazily (and memoized) from the experiment's
+ * `databricksTraceDestinationPath` tag on the first export, because the
+ * public `initDatabricksTracing` surface is synchronous and the SDK's
+ * authenticated client only exists at export time. Experiments without the
+ * tag fall through to the SDK's classic export path unchanged.
+ */
+import { createRequire } from "node:module";
+import { ProtobufTraceSerializer } from "@opentelemetry/otlp-transformer";
+import { DatabricksTracingConfigurationError } from "./errors.js";
+/** Experiment tag carrying the UC destination (`catalog.schema.table_prefix`). */
+export const UC_DESTINATION_TAG = "mlflow.experiment.databricksTraceDestinationPath";
+/** Experiment tag naming the spans table outright (preferred over deriving). */
+export const UC_SPANS_TABLE_TAG = "mlflow.experiment.databricksTraceSpanStorageTable";
+/** Header naming the UC spans table for the OTLP ingest endpoint. */
+export const UC_TABLE_HEADER = "X-Databricks-UC-Table-Name";
+/** V4 schema version stamped into UC trace metadata. */
+const TRACE_SCHEMA_VERSION_KEY = "mlflow.trace_schema.version";
+const UC_TRACE_SCHEMA_VERSION = "4";
+/**
+ * Parse a `databricksTraceDestinationPath` tag value into a {@link UcDestination}.
+ * Throws when the value is not the expected three-part `catalog.schema.prefix`.
+ */
+export function parseUcDestinationPath(path) {
+    const parts = path.split(".");
+    const [catalogName, schemaName, tablePrefix] = parts;
+    if (parts.length !== 3 || !catalogName || !schemaName || !tablePrefix) {
+        throw new DatabricksTracingConfigurationError(`Unexpected ${UC_DESTINATION_TAG} value '${path}' ` +
+            "(expected 'catalog.schema.table_prefix').");
+    }
+    return {
+        catalogName,
+        schemaName,
+        tablePrefix,
+        fullTablePrefix: path,
+        spansTableName: `${catalogName}.${schemaName}.${tablePrefix}_otel_spans`,
+    };
+}
+/**
+ * Fetch the experiment and resolve its UC destination from the
+ * `databricksTraceDestinationPath` tag. Returns `null` for classic
+ * (workspace-backed) experiments. SDK/network failures raise an actionable
+ * configuration error rather than a generic fetch failure, mirroring the
+ * Python package's hardened guard.
+ */
+export async function resolveUcDestination(hostUrl, headersProvider, experimentId) {
+    const url = `${hostUrl.replace(/\/$/, "")}/api/2.0/mlflow/experiments/get?experiment_id=${encodeURIComponent(experimentId)}`;
+    let response;
+    try {
+        response = await fetch(url, { headers: await headersProvider() });
+    }
+    catch (err) {
+        throw new DatabricksTracingConfigurationError(`Failed to fetch experiment '${experimentId}' while resolving the UC trace ` +
+            `destination: ${err.message}`);
+    }
+    if (!response.ok) {
+        throw new DatabricksTracingConfigurationError(`Failed to fetch experiment '${experimentId}' while resolving the UC trace ` +
+            `destination: HTTP ${response.status}`);
+    }
+    const body = (await response.json());
+    const tags = body.experiment?.tags ?? [];
+    const destinationTag = tags.find((t) => t.key === UC_DESTINATION_TAG);
+    if (!destinationTag?.value)
+        return null;
+    const destination = parseUcDestinationPath(destinationTag.value);
+    // Prefer the explicit spans-table tag when present; the derived
+    // `<prefix>_otel_spans` name is the documented default fallback.
+    const spansTableTag = tags.find((t) => t.key === UC_SPANS_TABLE_TAG);
+    if (spansTableTag?.value) {
+        destination.spansTableName = spansTableTag.value;
+    }
+    return destination;
+}
+/** Strip the `tr-` prefix from an MLflow trace id, yielding the OTel hex id. */
+function toOtelTraceId(mlflowTraceId) {
+    return mlflowTraceId.startsWith("tr-") ? mlflowTraceId.slice(3) : mlflowTraceId;
+}
+/**
+ * Databricks' OTLP ingest rejects batches over 10 MiB ("Batch exceeds maximum
+ * size of 10485760 bytes"). Stay comfortably under it — protobuf framing per
+ * chunk adds a little over the sum of span sizes.
+ */
+export const OTLP_MAX_BATCH_BYTES = 9 * 1024 * 1024;
+/**
+ * POST one OTLP `ExportTraceServiceRequest` to the Databricks ingest endpoint,
+ * addressed at the UC spans table.
+ */
+async function postOtlpBatch(client, destination, payload) {
+    const url = `${client.hostUrl.replace(/\/$/, "")}/api/2.0/otel/v1/traces`;
+    const response = await fetch(url, {
+        method: "POST",
+        headers: {
+            ...(await client.headersProvider()),
+            "Content-Type": "application/x-protobuf",
+            [UC_TABLE_HEADER]: destination.spansTableName,
+        },
+        body: payload,
+    });
+    if (!response.ok) {
+        const text = await response.text().catch(() => "");
+        throw new Error(`OTLP span export failed: HTTP ${response.status} ${text.slice(0, 500)}`);
+    }
+}
+/**
+ * POST the trace's spans to the Databricks OTLP ingest endpoint, splitting
+ * into multiple requests when the serialized batch exceeds the ingest size
+ * cap (SQL-heavy agent turns routinely serialize past 10 MiB — one validation
+ * run saw a 98-span turn produce a 16.8 MB batch).
+ * Spans are append-only rows in the UC table, so a trace's spans may arrive
+ * across several requests. Binary-splits recursively (log N serializations);
+ * a single span that alone exceeds the cap is sent anyway and surfaces the
+ * server's 400 — truncating span payloads is the tracer's job, not the
+ * exporter's.
+ */
+async function exportSpansToUcTable(client, destination, spans, maxBatchBytes = OTLP_MAX_BATCH_BYTES) {
+    if (spans.length === 0)
+        return;
+    const payload = ProtobufTraceSerializer.serializeRequest(spans);
+    if (!payload) {
+        throw new Error("OTLP serialization produced no payload");
+    }
+    if (payload.byteLength > maxBatchBytes && spans.length > 1) {
+        const mid = Math.ceil(spans.length / 2);
+        await exportSpansToUcTable(client, destination, spans.slice(0, mid), maxBatchBytes);
+        await exportSpansToUcTable(client, destination, spans.slice(mid), maxBatchBytes);
+        return;
+    }
+    await postOtlpBatch(client, destination, payload);
+}
+/**
+ * POST the trace info to the V4 `CreateTraceInfo` endpoint with a
+ * `UC_TABLE_PREFIX` location and V4 trace-id format
+ * (`trace:/<location>/<otelTraceId>`).
+ */
+async function exportTraceInfoV4(client, destination, info) {
+    const otelTraceId = toOtelTraceId(info.traceId);
+    const body = info.toJson();
+    body.trace_id = `trace:/${destination.fullTablePrefix}/${otelTraceId}`;
+    body.trace_location = {
+        type: "UC_TABLE_PREFIX",
+        uc_table_prefix: {
+            catalog_name: destination.catalogName,
+            schema_name: destination.schemaName,
+            table_prefix: destination.tablePrefix,
+        },
+    };
+    body.trace_metadata = {
+        ...body.trace_metadata,
+        [TRACE_SCHEMA_VERSION_KEY]: UC_TRACE_SCHEMA_VERSION,
+    };
+    const url = `${client.hostUrl.replace(/\/$/, "")}/api/4.0/mlflow/traces/` +
+        `${destination.fullTablePrefix}/${otelTraceId}/info`;
+    const response = await fetch(url, {
+        method: "POST",
+        headers: {
+            ...(await client.headersProvider()),
+            "Content-Type": "application/json",
+        },
+        body: JSON.stringify(body),
+    });
+    if (!response.ok) {
+        const text = await response.text().catch(() => "");
+        throw new Error(`V4 trace-info export failed: HTTP ${response.status} ${text.slice(0, 500)}`);
+    }
+}
+let exporterPatched = false;
+/** Reset patch state. For testing only. */
+export function _resetUcExportForTests() {
+    exporterPatched = false;
+}
+/**
+ * Patch `MlflowSpanExporter.exportTraceToBackend` with a UC-aware
+ * implementation. The destination is resolved once (memoized promise) from
+ * the experiment's tags on the first export — the SDK's client (and its
+ * auth headers) only exists at export time — and classic experiments
+ * delegate to the original implementation untouched. Resolution failures
+ * log loudly and fall back to classic export rather than dropping traces.
+ */
+export function patchExporterForUc(experimentId) {
+    if (exporterPatched)
+        return;
+    const nodeRequire = createRequire(import.meta.url);
+    const exporterModule = nodeRequire("@mlflow/core/dist/exporters/mlflow.js");
+    const proto = exporterModule.MlflowSpanExporter.prototype;
+    const originalExport = proto.exportTraceToBackend;
+    let destinationPromise = null;
+    const resolveOnce = (client) => {
+        destinationPromise ??= resolveUcDestination(client.hostUrl, client.headersProvider, experimentId).catch((err) => {
+            console.error("[databricks-langgraph-tracer] UC destination resolution failed:", err);
+            return null;
+        });
+        return destinationPromise;
+    };
+    proto.exportTraceToBackend = async function (trace) {
+        const destination = await resolveOnce(this._client);
+        if (!destination) {
+            return originalExport.call(this, trace);
+        }
+        // UC path: OTLP spans -> UC table; trace info -> V4 endpoint. No
+        // artifact upload (UC experiments do not read the artifact store).
+        // Span export goes first — trace info without spans renders as an
+        // empty trace, while spans without info are recoverable on retry.
+        const spans = trace.data.spans.map((s) => s._span);
+        try {
+            await exportSpansToUcTable(this._client, destination, spans);
+            await exportTraceInfoV4(this._client, destination, trace.info);
+        }
+        catch (err) {
+            console.error(`[databricks-langgraph-tracer] UC export failed for ${trace.info.traceId}, ` +
+                "falling back to classic export:", err);
+            return originalExport.call(this, trace);
+        }
+        // Mirror the SDK's `finally` cleanup: `export()` records this promise in
+        // `_pendingExports` and the original implementation deletes the entry on
+        // completion. Only the UC success path needs this here — both delegation
+        // paths above end in `originalExport`, whose own `finally` cleans up
+        // (deleting earlier would hide an in-flight fallback from `forceFlush`).
+        delete this._pendingExports[trace.info.traceId];
+    };
+    exporterPatched = true;
+}
+//# sourceMappingURL=uc-export.js.map

package/dist/uc-export.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"uc-export.js","sourceRoot":"","sources":["../src/uc-export.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;GA2BG;AACH,OAAO,EAAE,aAAa,EAAE,MAAM,aAAa,CAAC;AAC5C,OAAO,EAAE,uBAAuB,EAAE,MAAM,iCAAiC,CAAC;AAE1E,OAAO,EAAE,mCAAmC,EAAE,MAAM,aAAa,CAAC;AAElE,kFAAkF;AAClF,MAAM,CAAC,MAAM,kBAAkB,GAAG,kDAAkD,CAAC;AAErF,gFAAgF;AAChF,MAAM,CAAC,MAAM,kBAAkB,GAAG,mDAAmD,CAAC;AAEtF,qEAAqE;AACrE,MAAM,CAAC,MAAM,eAAe,GAAG,4BAA4B,CAAC;AAE5D,wDAAwD;AACxD,MAAM,wBAAwB,GAAG,6BAA6B,CAAC;AAC/D,MAAM,uBAAuB,GAAG,GAAG,CAAC;AAYpC;;;GAGG;AACH,MAAM,UAAU,sBAAsB,CAAC,IAAY;IACjD,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;IAC9B,MAAM,CAAC,WAAW,EAAE,UAAU,EAAE,WAAW,CAAC,GAAG,KAAK,CAAC;IACrD,IAAI,KAAK,CAAC,MAAM,KAAK,CAAC,IAAI,CAAC,WAAW,IAAI,CAAC,UAAU,IAAI,CAAC,WAAW,EAAE,CAAC;QACtE,MAAM,IAAI,mCAAmC,CAC3C,cAAc,kBAAkB,WAAW,IAAI,IAAI;YACjD,2CAA2C,CAC9C,CAAC;IACJ,CAAC;IACD,OAAO;QACL,WAAW;QACX,UAAU;QACV,WAAW;QACX,eAAe,EAAE,IAAI;QACrB,cAAc,EAAE,GAAG,WAAW,IAAI,UAAU,IAAI,WAAW,aAAa;KACzE,CAAC;AACJ,CAAC;AAwCD;;;;;;GAMG;AACH,MAAM,CAAC,KAAK,UAAU,oBAAoB,CACxC,OAAe,EACf,eAAgC,EAChC,YAAoB;IAEpB,MAAM,GAAG,GAAG,GAAG,OAAO,CAAC,OAAO,CAAC,KAAK,EAAE,EAAE,CAAC,iDAAiD,kBAAkB,CAAC,YAAY,CAAC,EAAE,CAAC;IAC7H,IAAI,QAAkB,CAAC;IACvB,IAAI,CAAC;QACH,QAAQ,GAAG,MAAM,KAAK,CAAC,GAAG,EAAE,EAAE,OAAO,EAAE,MAAM,eAAe,EAAE,EAAE,CAAC,CAAC;IACpE,CAAC;IAAC,OAAO,GAAG,EAAE,CAAC;QACb,MAAM,IAAI,mCAAmC,CAC3C,+BAA+B,YAAY,iCAAiC;YAC1E,gBAAiB,GAAa,CAAC,OAAO,EAAE,CAC3C,CAAC;IACJ,CAAC;IACD,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;QACjB,MAAM,IAAI,mCAAmC,CAC3C,+BAA+B,YAAY,iCAAiC;YAC1E,qBAAqB,QAAQ,CAAC,MAAM,EAAE,CACzC,CAAC;IACJ,CAAC;IACD,MAAM,IAAI,GAAG,CAAC,MAAM,QAAQ,CAAC,IAAI,EAAE,CAElC,CAAC;IACF,MAAM,IAAI,GAAG,IAAI,CAAC,UAAU,EAAE,IAAI,IAAI,EAAE,CAAC;IACzC,MAAM,cAAc,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,KAAK,kBAAkB,CAAC,CAAC;IACtE,IAAI,CAAC,cAAc,EAAE,KAAK;QAAE,OAAO,IAAI,CAAC;IACxC,MAAM,WAAW,GAAG,sBAAsB,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC;IACjE,gEAAgE;IAChE,iEAAiE;IACjE,MAAM,aAAa,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,KAAK,kBAAkB,CAAC,CAAC;IACrE,IAAI,aAAa,EAAE,KAAK,EAAE,CAAC;QACzB,WAAW,CAAC,cAAc,GAAG,aAAa,CAAC,KAAK,CAAC;IACnD,CAAC;IACD,OAAO,WAAW,CAAC;AACrB,CAAC;AAED,gFAAgF;AAChF,SAAS,aAAa,CAAC,aAAqB;IAC1C,OAAO,aAAa,CAAC,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,aAAa,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,aAAa,CAAC;AAClF,CAAC;AAED;;;;GAIG;AACH,MAAM,CAAC,MAAM,oBAAoB,GAAG,CAAC,GAAG,IAAI,GAAG,IAAI,CAAC;AAEpD;;;GAGG;AACH,KAAK,UAAU,aAAa,CAC1B,MAA6B,EAC7B,WAA0B,EAC1B,OAAmB;IAEnB,MAAM,GAAG,GAAG,GAAG,MAAM,CAAC,OAAO,CAAC,OAAO,CAAC,KAAK,EAAE,EAAE,CAAC,yBAAyB,CAAC;IAC1E,MAAM,QAAQ,GAAG,MAAM,KAAK,CAAC,GAAG,EAAE;QAChC,MAAM,EAAE,MAAM;QACd,OAAO,EAAE;YACP,GAAG,CAAC,MAAM,MAAM,CAAC,eAAe,EAAE,CAAC;YACnC,cAAc,EAAE,wBAAwB;YACxC,CAAC,eAAe,CAAC,EAAE,WAAW,CAAC,cAAc;SAC9C;QACD,IAAI,EAAE,OAAO;KACd,CAAC,CAAC;IACH,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;QACjB,MAAM,IAAI,GAAG,MAAM,QAAQ,CAAC,IAAI,EAAE,CAAC,KAAK,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,CAAC;QACnD,MAAM,IAAI,KAAK,CAAC,iCAAiC,QAAQ,CAAC,MAAM,IAAI,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,CAAC,EAAE,CAAC,CAAC;IAC5F,CAAC;AACH,CAAC;AAED;;;;;;;;;;GAUG;AACH,KAAK,UAAU,oBAAoB,CACjC,MAA6B,EAC7B,WAA0B,EAC1B,KAAqB,EACrB,gBAAwB,oBAAoB;IAE5C,IAAI,KAAK,CAAC,MAAM,KAAK,CAAC;QAAE,OAAO;IAC/B,MAAM,OAAO,GAAG,uBAAuB,CAAC,gBAAgB,CAAC,KAAK,CAAC,CAAC;IAChE,IAAI,CAAC,OAAO,EAAE,CAAC;QACb,MAAM,IAAI,KAAK,CAAC,wCAAwC,CAAC,CAAC;IAC5D,CAAC;IACD,IAAI,OAAO,CAAC,UAAU,GAAG,aAAa,IAAI,KAAK,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;QAC3D,MAAM,GAAG,GAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC;QACxC,MAAM,oBAAoB,CAAC,MAAM,EAAE,WAAW,EAAE,KAAK,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,CAAC,EAAE,aAAa,CAAC,CAAC;QACpF,MAAM,oBAAoB,CAAC,MAAM,EAAE,WAAW,EAAE,KAAK,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,aAAa,CAAC,CAAC;QACjF,OAAO;IACT,CAAC;IACD,MAAM,aAAa,CAAC,MAAM,EAAE,WAAW,EAAE,OAAO,CAAC,CAAC;AACpD,CAAC;AAED;;;;GAIG;AACH,KAAK,UAAU,iBAAiB,CAC9B,MAA6B,EAC7B,WAA0B,EAC1B,IAAwB;IAExB,MAAM,WAAW,GAAG,aAAa,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;IAChD,MAAM,IAAI,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC;IAC3B,IAAI,CAAC,QAAQ,GAAG,UAAU,WAAW,CAAC,eAAe,IAAI,WAAW,EAAE,CAAC;IACvE,IAAI,CAAC,cAAc,GAAG;QACpB,IAAI,EAAE,iBAAiB;QACvB,eAAe,EAAE;YACf,YAAY,EAAE,WAAW,CAAC,WAAW;YACrC,WAAW,EAAE,WAAW,CAAC,UAAU;YACnC,YAAY,EAAE,WAAW,CAAC,WAAW;SACtC;KACF,CAAC;IACF,IAAI,CAAC,cAAc,GAAG;QACpB,GAAI,IAAI,CAAC,cAAyC;QAClD,CAAC,wBAAwB,CAAC,EAAE,uBAAuB;KACpD,CAAC;IACF,MAAM,GAAG,GACP,GAAG,MAAM,CAAC,OAAO,CAAC,OAAO,CAAC,KAAK,EAAE,EAAE,CAAC,yBAAyB;QAC7D,GAAG,WAAW,CAAC,eAAe,IAAI,WAAW,OAAO,CAAC;IACvD,MAAM,QAAQ,GAAG,MAAM,KAAK,CAAC,GAAG,EAAE;QAChC,MAAM,EAAE,MAAM;QACd,OAAO,EAAE;YACP,GAAG,CAAC,MAAM,MAAM,CAAC,eAAe,EAAE,CAAC;YACnC,cAAc,EAAE,kBAAkB;SACnC;QACD,IAAI,EAAE,IAAI,CAAC,SAAS,CAAC,IAAI,CAAC;KAC3B,CAAC,CAAC;IACH,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;QACjB,MAAM,IAAI,GAAG,MAAM,QAAQ,CAAC,IAAI,EAAE,CAAC,KAAK,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,CAAC;QACnD,MAAM,IAAI,KAAK,CAAC,qCAAqC,QAAQ,CAAC,MAAM,IAAI,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,CAAC,EAAE,CAAC,CAAC;IAChG,CAAC;AACH,CAAC;AAED,IAAI,eAAe,GAAG,KAAK,CAAC;AAE5B,2CAA2C;AAC3C,MAAM,UAAU,sBAAsB;IACpC,eAAe,GAAG,KAAK,CAAC;AAC1B,CAAC;AAED;;;;;;;GAOG;AACH,MAAM,UAAU,kBAAkB,CAAC,YAAoB;IACrD,IAAI,eAAe;QAAE,OAAO;IAC5B,MAAM,WAAW,GAAG,aAAa,CAAC,MAAM,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;IACnD,MAAM,cAAc,GAAG,WAAW,CAChC,uCAAuC,CAChB,CAAC;IAC1B,MAAM,KAAK,GAAG,cAAc,CAAC,kBAAkB,CAAC,SAAS,CAAC;IAC1D,MAAM,cAAc,GAAG,KAAK,CAAC,oBAAoB,CAAC;IAElD,IAAI,kBAAkB,GAAyC,IAAI,CAAC;IACpE,MAAM,WAAW,GAAG,CAAC,MAA6B,EAAiC,EAAE;QACnF,kBAAkB,KAAK,oBAAoB,CACzC,MAAM,CAAC,OAAO,EACd,MAAM,CAAC,eAAe,EACtB,YAAY,CACb,CAAC,KAAK,CAAC,CAAC,GAAG,EAAE,EAAE;YACd,OAAO,CAAC,KAAK,CAAC,iEAAiE,EAAE,GAAG,CAAC,CAAC;YACtF,OAAO,IAAI,CAAC;QACd,CAAC,CAAC,CAAC;QACH,OAAO,kBAAkB,CAAC;IAC5B,CAAC,CAAC;IAEF,KAAK,CAAC,oBAAoB,GAAG,KAAK,WAEhC,KAAqB;QAErB,MAAM,WAAW,GAAG,MAAM,WAAW,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpD,IAAI,CAAC,WAAW,EAAE,CAAC;YACjB,OAAO,cAAc,CAAC,IAAI,CAAC,IAAI,EAAE,KAAK,CAAkB,CAAC;QAC3D,CAAC;QACD,iEAAiE;QACjE,mEAAmE;QACnE,kEAAkE;QAClE,kEAAkE;QAClE,MAAM,KAAK,GAAG,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;QACnD,IAAI,CAAC;YACH,MAAM,oBAAoB,CAAC,IAAI,CAAC,OAAO,EAAE,WAAW,EAAE,KAAK,CAAC,CAAC;YAC7D,MAAM,iBAAiB,CAAC,IAAI,CAAC,OAAO,EAAE,WAAW,EAAE,KAAK,CAAC,IAAI,CAAC,CAAC;QACjE,CAAC;QAAC,OAAO,GAAG,EAAE,CAAC;YACb,OAAO,CAAC,KAAK,CACX,sDAAsD,KAAK,CAAC,IAAI,CAAC,OAAO,IAAI;gBAC1E,iCAAiC,EACnC,GAAG,CACJ,CAAC;YACF,OAAO,cAAc,CAAC,IAAI,CAAC,IAAI,EAAE,KAAK,CAAkB,CAAC;QAC3D,CAAC;QACD,yEAAyE;QACzE,yEAAyE;QACzE,yEAAyE;QACzE,qEAAqE;QACrE,yEAAyE;QACzE,OAAO,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;IAClD,CAAC,CAAC;IACF,eAAe,GAAG,IAAI,CAAC;AACzB,CAAC"}

package/package.json

@@ -0,0 +1,68 @@
+{
+  "name": "@insightfactory.ai/insightfactory-databricks-langgraph-tracer",
+  "version": "1.0.0-dev.0",
+  "description": "LangGraph tracing for Databricks MLflow (TypeScript)",
+  "type": "module",
+  "packageManager": "yarn@4.12.0",
+  "license": "SEE LICENSE IN LICENSE",
+  "author": {
+    "name": "insightfactory.ai",
+    "email": "support@insightfactory.ai"
+  },
+  "homepage": "https://insightfactory.ai",
+  "keywords": [
+    "databricks",
+    "langchain",
+    "langgraph",
+    "mlflow",
+    "tracing"
+  ],
+  "engines": {
+    "node": ">=20"
+  },
+  "publishConfig": {
+    "registry": "https://registry.npmjs.org/",
+    "access": "public",
+    "provenance": true
+  },
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
+  "files": [
+    "dist",
+    "src/pricing",
+    "CHANGELOG.md",
+    "LICENSE",
+    "THIRD_PARTY_NOTICES",
+    "README.md"
+  ],
+  "scripts": {
+    "build": "tsc -p tsconfig.build.json && node scripts/copy-pricing.mjs",
+    "test": "vitest run",
+    "keys:generate": "tsx scripts/generate-keys.ts",
+    "keys:check": "tsx scripts/generate-keys.ts --check",
+    "pricing:update": "tsx scripts/update-pricing.ts"
+  },
+  "dependencies": {
+    "@opentelemetry/otlp-transformer": "^0.205.0",
+    "@opentelemetry/sdk-trace-base": "^2.1.0"
+  },
+  "peerDependencies": {
+    "@langchain/core": "^1.1.47",
+    "@mlflow/core": "^0.2.0"
+  },
+  "devDependencies": {
+    "@langchain/core": "^1.1.47",
+    "@langchain/langgraph": "^1.2.6",
+    "@mlflow/core": "^0.2.0",
+    "@types/node": "^25.5.0",
+    "tsx": "^4.21.0",
+    "typescript": "^6.0.2",
+    "vitest": "^4.1.4"
+  }
+}

package/src/pricing/model-pricing-overrides.json

@@ -0,0 +1,28 @@
+{
+  "azure_ai/deepseek-v4-flash": {
+    "input_cost_per_token": 1.9e-7,
+    "output_cost_per_token": 5.1e-7,
+    "litellm_provider": "azure_ai",
+    "max_input_tokens": 131072,
+    "max_output_tokens": 131072,
+    "max_tokens": 131072,
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_reasoning": true,
+    "supports_tool_choice": true,
+    "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-deepseek-v4-flash-and-v4-pro-in-microsoft-foundry/4515174"
+  },
+  "azure_ai/deepseek-v4-pro": {
+    "input_cost_per_token": 1.74e-6,
+    "output_cost_per_token": 3.48e-6,
+    "litellm_provider": "azure_ai",
+    "max_input_tokens": 131072,
+    "max_output_tokens": 131072,
+    "max_tokens": 131072,
+    "mode": "chat",
+    "supports_function_calling": true,
+    "supports_reasoning": true,
+    "supports_tool_choice": true,
+    "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-deepseek-v4-flash-and-v4-pro-in-microsoft-foundry/4515174"
+  }
+}