Metadata-Version: 2.5
Name: insightfactory-databricks-langgraph-tracer
Version: 1.0.0.dev15
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. |
| 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.

### 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).
