insightfactory-databricks-langgraph-tracer

LangGraph tracer for Databricks MLflow. It uses a custom LangChain BaseTracer to write trace tags, metadata, token usage, and cost rollups before the root span ends.

Install

Published on PyPI:

uv add insightfactory-databricks-langgraph-tracer
# or: pip install insightfactory-databricks-langgraph-tracer

Requires Python 3.12 and resolves mlflow>=3.15.0,<4.

Quickstart

from databricks_langgraph_tracer import (
    configure_databricks_tracing,
    get_tracing_callbacks,
)

# 1. Bootstrap once at startup (reads env-first; kwargs override).
configure_databricks_tracing(experiment_id="<mlflow-experiment-id>", source="my-agent")

# 2. Attach the callbacks to your LangGraph / LangChain run.
graph = build_graph().with_config({"callbacks": get_tracing_callbacks()})
graph.invoke(state)

Traces appear in the configured Databricks MLflow experiment with the full shared schema: the source tag, session, user and thread metadata, model, provider, token usage and cost per span, the trace-level mlflow.trace.cost rollup, and the per-model cost.by_model rollup tag.

Configuration

Environment variables are read first. Any kwarg to configure_databricks_tracing(...) overrides the matching variable.

Auth is resolved by databricks_utils. Use a service principal through DATABRICKS_HOST, DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET, or a CLI profile through DATABRICKS_CONFIG_PROFILE or the profile= kwarg.

Missing required config raises DatabricksTracingConfigurationError at startup. Disable the tracer explicitly for local and dev runs.

Unity Catalog-backed tracing

For experiments whose traces live in Unity Catalog _otel_* tables, set uc_tracing=True or MLFLOW_TRACING_UC_BACKED=true and provide a SQL warehouse. The library validates the warehouse, resolves the experiment's UC trace location from its binding tag, and passes it to set_experiment so spans persist to the _otel_spans table. Without that step MLflow silently skips span export to UC. Classic workspace experiments need none of this and are the default. UC-backed tracing needs an MLflow release that provides the UnityCatalog trace-location API.

Multimodal inputs (image, PDF, file)

Chat-model spans record their inputs as structured messages, so the multimodal content parts a graph sends to the model survive: OpenAI and LangChain image_url, OpenAI file, and Anthropic image and document. Internally the tracer runs the LangChain BaseTracer in original+chat mode. The default mode would flatten chat messages to a text-only prompts string and drop every attachment.

The inline base64 of each such part is never stored in the trace. The tracer removes it before recording and replaces it with a small reference, so the Unity Catalog trace tables stay readable. Multi-MB data URIs used to push large invoice traces past the SQL inline read limit (issue #21). A remote http(s):// image URL is already a small reference, so this step keeps it verbatim. The text cap below, if enabled, still truncates any string over its threshold, URLs included. The transform runs on a copy of the inputs, so the live message sent to the model is untouched and prompt caching is unaffected.

By default a part becomes a {"type": ..., "_omitted": true, "bytes": N} placeholder. To store a meaningful reference instead, such as the Unity Catalog volume path the image was loaded from so it can be re-fetched later, pass a content_ref_resolver:

from databricks_langgraph_tracer import (
    ContentPartContext,
    configure_databricks_tracing,
)

def image_ref(part: dict, ctx: ContentPartContext) -> dict | None:
    # ctx.metadata is the run metadata. Pass per-run data (e.g. a source volume
    # path) via the invoke config's `metadata`, which propagates to the LLM run
    # alongside langgraph_node etc.
    path = ctx.metadata.get("encoded_images_path")
    if path:
        return {"type": part.get("type"), "ref": path, "page": ctx.index}
    return None  # fall back to the default placeholder

configure_databricks_tracing(experiment_id="...", content_ref_resolver=image_ref)

# ... then carry the per-run ref data on the invoke config metadata:
graph.invoke(state, config={
    "callbacks": get_tracing_callbacks(),
    "metadata": {"encoded_images_path": "/Volumes/cat/sch/vol/inv/pages.txt"},
})

The tracer calls the resolver once per multimodal part with the part and a ContentPartContext. The context carries index, the part's position in the message content array, which for a one-image-per-page invoice is the page number. It also carries bytes, the inline payload length, and the run metadata. Return a dict to store as the reference, or None for the default placeholder. DatabricksLangGraphTracer also accepts content_ref_resolver= for per-graph wiring.

No inline image bytes ever reach the trace. If a resolver result re-introduces an inline payload anywhere in the returned object, whether a data: URI, a recognized base64 content part, or the part's own payload echoed back under another key, the tracer rejects it and uses the placeholder. Beyond that, keep the reference small. The library strips inline payloads but does not otherwise bound what a resolver returns, so a fabricated large string under a custom key is the caller's problem. ctx.metadata is a shallow copy of the run metadata, so setting top-level keys in the resolver cannot corrupt run state. Its nested values are shared, so don't mutate those.

Capping large text

Multimodal externalization handles inline bytes, but large plain text can also push a trace past the SQL inline read limit (issue #23). A classification vocabulary or an aggregated result set threaded through every fan-out span's inputs and outputs is enough. Set max_string_chars, or DATABRICKS_TRACING_MAX_STRING_CHARS, to truncate it:

configure_databricks_tracing(experiment_id="...", max_string_chars=50_000)

When set, any string value longer than the threshold in a span's inputs or outputs is replaced with a placeholder:

{"_truncated": true, "chars": 812345, "bytes": 812345, "preview": "<the first 256 chars>"}

The cap is off by default. Generic truncation costs debuggability, so you choose the threshold. It caps string values only, never keys or structural fields, and runs on a copy, so the live messages are untouched. It also reaches text nested inside Pydantic models, dataclasses and tuples, such as a model a node returns as final_output, by normalizing them to the same shape MLflow records. DatabricksLangGraphTracer also accepts max_string_chars= for per-graph wiring.

This is a separate knob from the multimodal handling above. content_ref_resolver chooses references for inline image, PDF and file bytes. max_string_chars caps arbitrary text and never truncates a resolver's reference. Because it is generic, it also truncates any other string over the threshold, including a remote http(s):// image URL the multimodal step keeps verbatim, so set the threshold well above your reference and URL lengths. It is a per-leaf limit, not a per-trace byte budget. Enough sub-threshold leaves can still sum past the limit, so for the heaviest spans also record less. Pass ids or references through node state rather than full payloads.

The threshold counts characters. This package counts code points and the TypeScript package counts UTF-16 units, so the two can differ on non-BMP text. The inline limit is in bytes, and multibyte text can be up to four times larger in bytes than in characters, so for CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's bytes field always reports the exact UTF-8 size of the original.

Spans from outside the tracer

A helper decorated with @mlflow.trace, or an MLflow autolog integration, starts its span through MLflow's context rather than through LangChain's callbacks. The tracer publishes each live span into that context, so those spans nest under whichever node is running instead of opening a root trace of their own:

import mlflow

@mlflow.trace(span_type="TOOL", name="lakebase.execute_query")
def execute_query(sql: str): ...

Calling that from a graph node produces one trace, with the tool span parented to the node. Without it you get a second root trace per call, carrying none of the root tags, session metadata or cost rollup the tracer writes.

Attaching the context requires the tracer's callbacks to run in the caller's own context, so run_inline is set on the handler whenever propagation is on. Only LangChain's async callback manager reads that flag: without it an async run sends the handler through an executor with a copied context, which the node body never sees. The sync manager calls handlers in the caller's thread either way.

The tracer still parents its own spans from LangChain's run tree, never from the context, so concurrent branches of a graph cannot cross-parent. Runs that overlap rather than nest, two stream(...) generators created before either is consumed, release their attaches in the order OpenTelemetry requires: a run that ends while another is still stacked above it holds its attach until that one ends, so a finished span is never restored over a live one.

The price of never releasing out of order is that a run whose end callback never fires pins the attaches beneath it. A node that abandons an inner stream(...) part-consumed leaves its span active for the life of that context, and spans started there parent to a run that has already finished. Sweeping the stack when the outermost run ends would clear it and bring back the overlapping-run bug above, so the tracer reports it instead: past a depth no real graph reaches, it logs a warning naming the count.

Only one handler attaches per run. Attach a second context-propagating tracer to the same graph and it records its own spans as before, but leaves the context to the first.

Three limits are worth knowing:

• It only works in one direction. A root run is always created without a parent and without consulting the context, so a graph invoked from inside an @mlflow.trace function still opens its own root trace rather than nesting into the caller's. This is the one place the TypeScript package behaves differently: MLflow's JavaScript startSpan falls back to the active context, so a graph there does nest into an enclosing span.
• Foreign spans are recorded by MLflow, not by this tracer, so neither the text cap nor the multimodal reference substitution applies to them. A traced query helper returning a large rowset can push a trace back over the inline read limit that max_string_chars was set to stay under.
• The tracer's callbacks run inline on the caller's thread, which on an async graph means the event loop. Keep content_ref_resolver fast and non-blocking.
• Another handler on the same run that also publishes to MLflow's context, which in practice means MlflowLangchainTracer from mlflow.langchain.autolog(), cannot be coordinated with. LangChain calls the two handlers' end hooks in registration order, so the attaches nest but the releases run front to back, and whichever releases last restores an ended span. Two of MLflow's own tracers on one run leak the same way with this tracer absent, so no release order fixes it here. The tracer warns once when it detects the conflict and names both settings to turn off; turn one of them off. Detection asks whether MLflow can still resolve the span that should be active, so a process that has shut tracing down with flush_trace_async_logging(terminate=True), or disabled it, gets no warning. Nothing can be adopted into a finalized trace in that state either. Issue #57 tracks it.

To turn the whole behaviour off, pass context_propagation=False or set DATABRICKS_TRACING_CONTEXT_PROPAGATION=false (0, no and off also work). That also restores LangChain's default callback dispatch.

The TypeScript package reaches the same outcome through getActiveTracingSpan() and withTracingContext(), named at the call site, because Node has no ambient equivalent. See ../typescript/README.md.

Autolog fallback

mode="autolog" wires MLflow's built-in LangChain autologging plus compatibility patches. It emits a reduced schema, everything except a tracer-computed mlflow.trace.cost rollup. The backend may still aggregate that server-side. The default mode="tracer", the custom BaseTracer, is the full-parity path.

Development

cd python
uv sync                                            # latest allowed mlflow (ceiling)
uv run pytest                                      # unit + lifecycle tests
uv run ruff check src tests
uv run ty check
uv run python scripts/generate_keys.py --check     # schema/keys parity

# mlflow floor matrix (CI runs both cells via UV_RESOLUTION):
UV_RESOLUTION=lowest-direct uv sync && UV_RESOLUTION=lowest-direct uv run pytest

Tests use a local sqlite MLflow tracking backend. Checks that need a live Databricks backend are marked integration.

Changelog

See CHANGELOG.md.