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

| Setting | Env var | Notes |
|---------|---------|-------|
| Tracking URI | `MLFLOW_TRACKING_URI` | `databricks` or `databricks://<profile>` (required) |
| Experiment | `MLFLOW_EXPERIMENT_ID` | by **id** only (required) |
| Source tag | — | `source=` kwarg, default `langgraph` |
| Multimodal refs | — | inline image/PDF/file bytes are externalized to a reference; `content_ref_resolver=` chooses it — see below |
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | opt-in; `max_string_chars=` truncates over-long plain-text span content — see below |
| Disable | `TESTING` / `BUILDING` = `true`, or `enabled=False` | no-op — the only non-raising path |
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | see below |
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID` → `DATABRICKS_WAREHOUSE_ID` | required when UC-backed |

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

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

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

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

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