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.

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.