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 carrying the full shared schema: the source tag, session/user/thread metadata, per-span model / provider / token-usage + cost, the trace-level mlflow.trace.cost rollup, and the per-model cost.by_model cost/token rollup tag.

Configuration

Env-first; any kwarg to configure_databricks_tracing(...) overrides the matching env var.

Auth (resolved by databricks_utils): a service principal (DATABRICKS_HOST / DATABRICKS_CLIENT_ID / DATABRICKS_CLIENT_SECRET) or a CLI profile (DATABRICKS_CONFIG_PROFILE / profile=).

Missing required config raises DatabricksTracingConfigurationError (fail-fast) — disable explicitly for local/dev runs.

Unity-Catalog-backed tracing

For experiments whose traces are stored 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 actually persist to the _otel_spans table — without it, MLflow silently skips span export to UC. Classic (workspace) experiments need none of this and are the default. UC-backed tracing requires an MLflow release that provides the UnityCatalog trace-location API.

Multimodal inputs (image / PDF / file)

Chat-model spans record their inputs as structured messages, preserving the multimodal content parts a graph sends to the model — OpenAI / LangChain image_url, OpenAI file, and Anthropic image / document. (Internally the tracer runs the LangChain BaseTracer in original+chat mode; the default would otherwise 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: it is removed before recording and replaced with a lightweight reference, so the Unity Catalog trace tables stay readable (multi-MB data URIs previously pushed 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 (but see the text cap below — if enabled, it 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 — e.g. the Unity Catalog volume path the image was loaded from, so it can be re-fetched at runtime — 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 resolver is called once per multimodal part with the part and a ContentPartContext — index (position within the message content array, i.e. the page number for a one-image-per-page invoice), bytes (the inline payload length), and the run metadata. Return a dict to store as the reference, or None for the default placeholder. It is also accepted by DatabricksLangGraphTracer(content_ref_resolver=...) for per-graph wiring.

The library guarantees no inline image bytes are stored: if a resolver result re-introduces an inline payload (a data: URI, a recognized base64 content part, or the part's own payload echoed back under any key — anywhere in the returned object), it is rejected and the placeholder is used. Beyond that, keep the reference compact — a fabricated large string under a custom key is the consumer's responsibility (the library strips inline payloads but does not otherwise bound what a resolver returns). ctx.metadata is a shallow copy of the run metadata, so a resolver cannot corrupt run state by setting top-level keys (don't mutate its nested values, which are shared).

Capping large text (issue #23)

Multimodal externalization handles inline bytes, but large plain text can also push a trace past the SQL inline read limit — e.g. a classification vocabulary or aggregated result set threaded through every fan-out span's inputs and outputs. 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 compact placeholder:

{"_truncated": true, "chars": 812345, "bytes": 812345, "preview": "first 256 chars…"}

It is opt-in / 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. The cap also reaches text nested inside Pydantic models / dataclasses / tuples (e.g. a model returned by a node such as final_output), normalizing them to the same shape MLflow records. Also accepted by DatabricksLangGraphTracer(max_string_chars=...) for per-graph wiring.

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

The threshold counts characters (code points here; the TypeScript package counts UTF-16 units, so the two can differ on non-BMP text), while the inline limit is in bytes — multibyte text can be up to ~4× larger in bytes than characters, so for CJK/emoji-heavy content size the cap below limit / 4. The placeholder's bytes field always reports the exact UTF-8 size of the original.

Autolog fallback

mode="autolog" wires MLflow's built-in LangChain autologging plus compatibility patches and emits a reduced schema — everything except a tracer-computed mlflow.trace.cost rollup (the backend may still aggregate it server-side). The default mode="tracer" (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.