Metadata-Version: 2.5
Name: insightfactory-databricks-langgraph-tracer
Version: 1.0.1.dev21
Summary: LangGraph tracing for Databricks MLflow
Project-URL: Homepage, https://insightfactory.ai
Author-email: "insightfactory.ai" <support@insightfactory.ai>
License-Expression: LicenseRef-Proprietary
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES
Keywords: databricks,langchain,langgraph,mlflow,tracing
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: License :: Other/Proprietary License
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.12
Classifier: Typing :: Typed
Requires-Python: <3.13,>=3.12
Requires-Dist: databricks-sdk>=0.82.0
Requires-Dist: langchain-core>=1.0.0
Requires-Dist: langchain>=1.3.9
Requires-Dist: langgraph>=1.0.0
Requires-Dist: mlflow<4.0.0,>=3.15.0
Description-Content-Type: text/markdown

# 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:

```bash
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

```python
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.

| Setting | Env var | Notes |
|---------|---------|-------|
| Tracking URI | `MLFLOW_TRACKING_URI` | `databricks` or `databricks://<profile>`. Required. |
| Experiment | `MLFLOW_EXPERIMENT_ID` | By id only. Required. |
| Source tag | none | `source=` kwarg, default `langgraph` |
| Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
| Context propagation | `DATABRICKS_TRACING_CONTEXT_PROPAGATION` | On by default. Set to `false`, `0`, `no` or `off`, or pass `context_propagation=False`, to stop publishing spans into MLflow's context. See below. |
| Disable | `TESTING` or `BUILDING` set to `true`, or `enabled=False` | Turns the tracer into a no-op. This is the only path that does not raise. |
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |

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`:

```python
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:

```python
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:

```json
{"_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:

```python
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`](../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

```bash
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`](CHANGELOG.md).
