insightfactory-databricks-langgraph-tracer 1.0.0.dev10__tar.gz → 1.0.0.dev15__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- insightfactory_databricks_langgraph_tracer-1.0.0.dev15/PKG-INFO +221 -0
- insightfactory_databricks_langgraph_tracer-1.0.0.dev15/README.md +197 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/pyproject.toml +1 -1
- insightfactory_databricks_langgraph_tracer-1.0.0.dev10/PKG-INFO +0 -224
- insightfactory_databricks_langgraph_tracer-1.0.0.dev10/README.md +0 -200
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/.gitignore +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/CHANGELOG.md +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/LICENSE +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/THIRD_PARTY_NOTICES +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/__init__.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/_generated_keys.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/aggregate.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/autolog_compat.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/cost.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/databricks_utils.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/errors.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/mlflow_utils.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/pricing/model-pricing-overrides.json +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/pricing/model-pricing.json +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/py.typed +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/tracer.py +0 -0
- {insightfactory_databricks_langgraph_tracer-1.0.0.dev10 → insightfactory_databricks_langgraph_tracer-1.0.0.dev15}/src/databricks_langgraph_tracer/uc_export.py +0 -0
|
@@ -0,0 +1,221 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: insightfactory-databricks-langgraph-tracer
|
|
3
|
+
Version: 1.0.0.dev15
|
|
4
|
+
Summary: LangGraph tracing for Databricks MLflow
|
|
5
|
+
Project-URL: Homepage, https://insightfactory.ai
|
|
6
|
+
Author-email: "insightfactory.ai" <support@insightfactory.ai>
|
|
7
|
+
License-Expression: LicenseRef-Proprietary
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
License-File: THIRD_PARTY_NOTICES
|
|
10
|
+
Keywords: databricks,langchain,langgraph,mlflow,tracing
|
|
11
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: Other/Proprietary License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Typing :: Typed
|
|
17
|
+
Requires-Python: <3.13,>=3.12
|
|
18
|
+
Requires-Dist: databricks-sdk>=0.82.0
|
|
19
|
+
Requires-Dist: langchain-core>=1.0.0
|
|
20
|
+
Requires-Dist: langchain>=1.3.9
|
|
21
|
+
Requires-Dist: langgraph>=1.0.0
|
|
22
|
+
Requires-Dist: mlflow<4.0.0,>=3.15.0
|
|
23
|
+
Description-Content-Type: text/markdown
|
|
24
|
+
|
|
25
|
+
# insightfactory-databricks-langgraph-tracer
|
|
26
|
+
|
|
27
|
+
LangGraph tracer for Databricks MLflow. It uses a custom LangChain `BaseTracer` to write
|
|
28
|
+
trace tags, metadata, token usage, and cost rollups before the root span ends.
|
|
29
|
+
|
|
30
|
+
## Install
|
|
31
|
+
|
|
32
|
+
Published on PyPI:
|
|
33
|
+
|
|
34
|
+
```bash
|
|
35
|
+
uv add insightfactory-databricks-langgraph-tracer
|
|
36
|
+
# or: pip install insightfactory-databricks-langgraph-tracer
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
Requires Python 3.12 and resolves `mlflow>=3.15.0,<4`.
|
|
40
|
+
|
|
41
|
+
## Quickstart
|
|
42
|
+
|
|
43
|
+
```python
|
|
44
|
+
from databricks_langgraph_tracer import (
|
|
45
|
+
configure_databricks_tracing,
|
|
46
|
+
get_tracing_callbacks,
|
|
47
|
+
)
|
|
48
|
+
|
|
49
|
+
# 1. Bootstrap once at startup (reads env-first; kwargs override).
|
|
50
|
+
configure_databricks_tracing(experiment_id="<mlflow-experiment-id>", source="my-agent")
|
|
51
|
+
|
|
52
|
+
# 2. Attach the callbacks to your LangGraph / LangChain run.
|
|
53
|
+
graph = build_graph().with_config({"callbacks": get_tracing_callbacks()})
|
|
54
|
+
graph.invoke(state)
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
Traces appear in the configured Databricks MLflow experiment with the full shared schema:
|
|
58
|
+
the `source` tag, session, user and thread metadata, model, provider, token usage and cost
|
|
59
|
+
per span, the trace-level `mlflow.trace.cost` rollup, and the per-model `cost.by_model`
|
|
60
|
+
rollup tag.
|
|
61
|
+
|
|
62
|
+
## Configuration
|
|
63
|
+
|
|
64
|
+
Environment variables are read first. Any kwarg to `configure_databricks_tracing(...)`
|
|
65
|
+
overrides the matching variable.
|
|
66
|
+
|
|
67
|
+
| Setting | Env var | Notes |
|
|
68
|
+
|---------|---------|-------|
|
|
69
|
+
| Tracking URI | `MLFLOW_TRACKING_URI` | `databricks` or `databricks://<profile>`. Required. |
|
|
70
|
+
| Experiment | `MLFLOW_EXPERIMENT_ID` | By id only. Required. |
|
|
71
|
+
| Source tag | none | `source=` kwarg, default `langgraph` |
|
|
72
|
+
| Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
|
|
73
|
+
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
|
|
74
|
+
| 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. |
|
|
75
|
+
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
|
|
76
|
+
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |
|
|
77
|
+
|
|
78
|
+
Auth is resolved by `databricks_utils`. Use a service principal through `DATABRICKS_HOST`,
|
|
79
|
+
`DATABRICKS_CLIENT_ID` and `DATABRICKS_CLIENT_SECRET`, or a CLI profile through
|
|
80
|
+
`DATABRICKS_CONFIG_PROFILE` or the `profile=` kwarg.
|
|
81
|
+
|
|
82
|
+
Missing required config raises `DatabricksTracingConfigurationError` at startup. Disable the
|
|
83
|
+
tracer explicitly for local and dev runs.
|
|
84
|
+
|
|
85
|
+
### Unity Catalog-backed tracing
|
|
86
|
+
|
|
87
|
+
For experiments whose traces live in Unity Catalog `_otel_*` tables, set `uc_tracing=True`
|
|
88
|
+
or `MLFLOW_TRACING_UC_BACKED=true` and provide a SQL warehouse. The library validates the
|
|
89
|
+
warehouse, resolves the experiment's UC trace location from its binding tag, and passes it
|
|
90
|
+
to `set_experiment` so spans persist to the `_otel_spans` table. Without that step MLflow
|
|
91
|
+
silently skips span export to UC. Classic workspace experiments need none of this and are
|
|
92
|
+
the default. UC-backed tracing needs an MLflow release that provides the `UnityCatalog`
|
|
93
|
+
trace-location API.
|
|
94
|
+
|
|
95
|
+
### Multimodal inputs (image, PDF, file)
|
|
96
|
+
|
|
97
|
+
Chat-model spans record their inputs as structured messages, so the multimodal content
|
|
98
|
+
parts a graph sends to the model survive: OpenAI and LangChain `image_url`, OpenAI `file`,
|
|
99
|
+
and Anthropic `image` and `document`. Internally the tracer runs the LangChain `BaseTracer`
|
|
100
|
+
in `original+chat` mode. The default mode would flatten chat messages to a text-only
|
|
101
|
+
`prompts` string and drop every attachment.
|
|
102
|
+
|
|
103
|
+
The inline base64 of each such part is never stored in the trace. The tracer removes it
|
|
104
|
+
before recording and replaces it with a small reference, so the Unity Catalog trace tables
|
|
105
|
+
stay readable. Multi-MB data URIs used to push large invoice traces past the SQL inline read
|
|
106
|
+
limit (issue #21). A remote `http(s)://` image URL is already a small reference, so this step
|
|
107
|
+
keeps it verbatim. The text cap below, if enabled, still truncates any string over its
|
|
108
|
+
threshold, URLs included. The transform runs on a copy of the inputs, so the live message
|
|
109
|
+
sent to the model is untouched and prompt caching is unaffected.
|
|
110
|
+
|
|
111
|
+
By default a part becomes a `{"type": ..., "_omitted": true, "bytes": N}` placeholder. To
|
|
112
|
+
store a meaningful reference instead, such as the Unity Catalog volume path the image was
|
|
113
|
+
loaded from so it can be re-fetched later, pass a `content_ref_resolver`:
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
from databricks_langgraph_tracer import (
|
|
117
|
+
ContentPartContext,
|
|
118
|
+
configure_databricks_tracing,
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
def image_ref(part: dict, ctx: ContentPartContext) -> dict | None:
|
|
122
|
+
# ctx.metadata is the run metadata. Pass per-run data (e.g. a source volume
|
|
123
|
+
# path) via the invoke config's `metadata`, which propagates to the LLM run
|
|
124
|
+
# alongside langgraph_node etc.
|
|
125
|
+
path = ctx.metadata.get("encoded_images_path")
|
|
126
|
+
if path:
|
|
127
|
+
return {"type": part.get("type"), "ref": path, "page": ctx.index}
|
|
128
|
+
return None # fall back to the default placeholder
|
|
129
|
+
|
|
130
|
+
configure_databricks_tracing(experiment_id="...", content_ref_resolver=image_ref)
|
|
131
|
+
|
|
132
|
+
# ... then carry the per-run ref data on the invoke config metadata:
|
|
133
|
+
graph.invoke(state, config={
|
|
134
|
+
"callbacks": get_tracing_callbacks(),
|
|
135
|
+
"metadata": {"encoded_images_path": "/Volumes/cat/sch/vol/inv/pages.txt"},
|
|
136
|
+
})
|
|
137
|
+
```
|
|
138
|
+
|
|
139
|
+
The tracer calls the resolver once per multimodal part with the part and a
|
|
140
|
+
`ContentPartContext`. The context carries `index`, the part's position in the message
|
|
141
|
+
content array, which for a one-image-per-page invoice is the page number. It also carries
|
|
142
|
+
`bytes`, the inline payload length, and the run `metadata`. Return a dict to store as the
|
|
143
|
+
reference, or `None` for the default placeholder. `DatabricksLangGraphTracer` also accepts
|
|
144
|
+
`content_ref_resolver=` for per-graph wiring.
|
|
145
|
+
|
|
146
|
+
No inline image bytes ever reach the trace. If a resolver result re-introduces an inline
|
|
147
|
+
payload anywhere in the returned object, whether a `data:` URI, a recognized base64 content
|
|
148
|
+
part, or the part's own payload echoed back under another key, the tracer rejects it and
|
|
149
|
+
uses the placeholder. Beyond that, keep the reference small. The library strips inline
|
|
150
|
+
payloads but does not otherwise bound what a resolver returns, so a fabricated large string
|
|
151
|
+
under a custom key is the caller's problem. `ctx.metadata` is a shallow copy of the run
|
|
152
|
+
metadata, so setting top-level keys in the resolver cannot corrupt run state. Its nested
|
|
153
|
+
values are shared, so don't mutate those.
|
|
154
|
+
|
|
155
|
+
### Capping large text
|
|
156
|
+
|
|
157
|
+
Multimodal externalization handles inline bytes, but large plain text can also push a trace
|
|
158
|
+
past the SQL inline read limit (issue #23). A classification vocabulary or an aggregated
|
|
159
|
+
result set threaded through every fan-out span's inputs and outputs is enough. Set
|
|
160
|
+
`max_string_chars`, or `DATABRICKS_TRACING_MAX_STRING_CHARS`, to truncate it:
|
|
161
|
+
|
|
162
|
+
```python
|
|
163
|
+
configure_databricks_tracing(experiment_id="...", max_string_chars=50_000)
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
When set, any string value longer than the threshold in a span's inputs or outputs is
|
|
167
|
+
replaced with a placeholder:
|
|
168
|
+
|
|
169
|
+
```json
|
|
170
|
+
{"_truncated": true, "chars": 812345, "bytes": 812345, "preview": "<the first 256 chars>"}
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
The cap is off by default. Generic truncation costs debuggability, so you choose the
|
|
174
|
+
threshold. It caps string values only, never keys or structural fields, and runs on a copy,
|
|
175
|
+
so the live messages are untouched. It also reaches text nested inside Pydantic models,
|
|
176
|
+
dataclasses and tuples, such as a model a node returns as `final_output`, by normalizing
|
|
177
|
+
them to the same shape MLflow records. `DatabricksLangGraphTracer` also accepts
|
|
178
|
+
`max_string_chars=` for per-graph wiring.
|
|
179
|
+
|
|
180
|
+
This is a separate knob from the multimodal handling above. `content_ref_resolver` chooses
|
|
181
|
+
references for inline image, PDF and file bytes. `max_string_chars` caps arbitrary text and
|
|
182
|
+
never truncates a resolver's reference. Because it is generic, it also truncates any other
|
|
183
|
+
string over the threshold, including a remote `http(s)://` image URL the multimodal step
|
|
184
|
+
keeps verbatim, so set the threshold well above your reference and URL lengths. It is a
|
|
185
|
+
per-leaf limit, not a per-trace byte budget. Enough sub-threshold leaves can still sum past
|
|
186
|
+
the limit, so for the heaviest spans also record less. Pass ids or references through node
|
|
187
|
+
state rather than full payloads.
|
|
188
|
+
|
|
189
|
+
The threshold counts characters. This package counts code points and the TypeScript
|
|
190
|
+
package counts UTF-16 units, so the two can differ on non-BMP text. The inline limit is in
|
|
191
|
+
bytes, and multibyte text can be up to four times larger in bytes than in characters, so for
|
|
192
|
+
CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's
|
|
193
|
+
`bytes` field always reports the exact UTF-8 size of the original.
|
|
194
|
+
|
|
195
|
+
### Autolog fallback
|
|
196
|
+
|
|
197
|
+
`mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches. It
|
|
198
|
+
emits a reduced schema, everything except a tracer-computed `mlflow.trace.cost` rollup. The
|
|
199
|
+
backend may still aggregate that server-side. The default `mode="tracer"`, the custom
|
|
200
|
+
`BaseTracer`, is the full-parity path.
|
|
201
|
+
|
|
202
|
+
## Development
|
|
203
|
+
|
|
204
|
+
```bash
|
|
205
|
+
cd python
|
|
206
|
+
uv sync # latest allowed mlflow (ceiling)
|
|
207
|
+
uv run pytest # unit + lifecycle tests
|
|
208
|
+
uv run ruff check src tests
|
|
209
|
+
uv run ty check
|
|
210
|
+
uv run python scripts/generate_keys.py --check # schema/keys parity
|
|
211
|
+
|
|
212
|
+
# mlflow floor matrix (CI runs both cells via UV_RESOLUTION):
|
|
213
|
+
UV_RESOLUTION=lowest-direct uv sync && UV_RESOLUTION=lowest-direct uv run pytest
|
|
214
|
+
```
|
|
215
|
+
|
|
216
|
+
Tests use a local sqlite MLflow tracking backend. Checks that need a live Databricks backend
|
|
217
|
+
are marked `integration`.
|
|
218
|
+
|
|
219
|
+
## Changelog
|
|
220
|
+
|
|
221
|
+
See [`CHANGELOG.md`](CHANGELOG.md).
|
|
@@ -0,0 +1,197 @@
|
|
|
1
|
+
# insightfactory-databricks-langgraph-tracer
|
|
2
|
+
|
|
3
|
+
LangGraph tracer for Databricks MLflow. It uses a custom LangChain `BaseTracer` to write
|
|
4
|
+
trace tags, metadata, token usage, and cost rollups before the root span ends.
|
|
5
|
+
|
|
6
|
+
## Install
|
|
7
|
+
|
|
8
|
+
Published on PyPI:
|
|
9
|
+
|
|
10
|
+
```bash
|
|
11
|
+
uv add insightfactory-databricks-langgraph-tracer
|
|
12
|
+
# or: pip install insightfactory-databricks-langgraph-tracer
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
Requires Python 3.12 and resolves `mlflow>=3.15.0,<4`.
|
|
16
|
+
|
|
17
|
+
## Quickstart
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from databricks_langgraph_tracer import (
|
|
21
|
+
configure_databricks_tracing,
|
|
22
|
+
get_tracing_callbacks,
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
# 1. Bootstrap once at startup (reads env-first; kwargs override).
|
|
26
|
+
configure_databricks_tracing(experiment_id="<mlflow-experiment-id>", source="my-agent")
|
|
27
|
+
|
|
28
|
+
# 2. Attach the callbacks to your LangGraph / LangChain run.
|
|
29
|
+
graph = build_graph().with_config({"callbacks": get_tracing_callbacks()})
|
|
30
|
+
graph.invoke(state)
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
Traces appear in the configured Databricks MLflow experiment with the full shared schema:
|
|
34
|
+
the `source` tag, session, user and thread metadata, model, provider, token usage and cost
|
|
35
|
+
per span, the trace-level `mlflow.trace.cost` rollup, and the per-model `cost.by_model`
|
|
36
|
+
rollup tag.
|
|
37
|
+
|
|
38
|
+
## Configuration
|
|
39
|
+
|
|
40
|
+
Environment variables are read first. Any kwarg to `configure_databricks_tracing(...)`
|
|
41
|
+
overrides the matching variable.
|
|
42
|
+
|
|
43
|
+
| Setting | Env var | Notes |
|
|
44
|
+
|---------|---------|-------|
|
|
45
|
+
| Tracking URI | `MLFLOW_TRACKING_URI` | `databricks` or `databricks://<profile>`. Required. |
|
|
46
|
+
| Experiment | `MLFLOW_EXPERIMENT_ID` | By id only. Required. |
|
|
47
|
+
| Source tag | none | `source=` kwarg, default `langgraph` |
|
|
48
|
+
| Multimodal refs | none | Inline image, PDF and file bytes become a reference. `content_ref_resolver=` chooses it. See below. |
|
|
49
|
+
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | Opt-in. `max_string_chars=` truncates long plain-text span content. See below. |
|
|
50
|
+
| 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. |
|
|
51
|
+
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | See below. |
|
|
52
|
+
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID`, falling back to `DATABRICKS_WAREHOUSE_ID` | Required when UC-backed. |
|
|
53
|
+
|
|
54
|
+
Auth is resolved by `databricks_utils`. Use a service principal through `DATABRICKS_HOST`,
|
|
55
|
+
`DATABRICKS_CLIENT_ID` and `DATABRICKS_CLIENT_SECRET`, or a CLI profile through
|
|
56
|
+
`DATABRICKS_CONFIG_PROFILE` or the `profile=` kwarg.
|
|
57
|
+
|
|
58
|
+
Missing required config raises `DatabricksTracingConfigurationError` at startup. Disable the
|
|
59
|
+
tracer explicitly for local and dev runs.
|
|
60
|
+
|
|
61
|
+
### Unity Catalog-backed tracing
|
|
62
|
+
|
|
63
|
+
For experiments whose traces live in Unity Catalog `_otel_*` tables, set `uc_tracing=True`
|
|
64
|
+
or `MLFLOW_TRACING_UC_BACKED=true` and provide a SQL warehouse. The library validates the
|
|
65
|
+
warehouse, resolves the experiment's UC trace location from its binding tag, and passes it
|
|
66
|
+
to `set_experiment` so spans persist to the `_otel_spans` table. Without that step MLflow
|
|
67
|
+
silently skips span export to UC. Classic workspace experiments need none of this and are
|
|
68
|
+
the default. UC-backed tracing needs an MLflow release that provides the `UnityCatalog`
|
|
69
|
+
trace-location API.
|
|
70
|
+
|
|
71
|
+
### Multimodal inputs (image, PDF, file)
|
|
72
|
+
|
|
73
|
+
Chat-model spans record their inputs as structured messages, so the multimodal content
|
|
74
|
+
parts a graph sends to the model survive: OpenAI and LangChain `image_url`, OpenAI `file`,
|
|
75
|
+
and Anthropic `image` and `document`. Internally the tracer runs the LangChain `BaseTracer`
|
|
76
|
+
in `original+chat` mode. The default mode would flatten chat messages to a text-only
|
|
77
|
+
`prompts` string and drop every attachment.
|
|
78
|
+
|
|
79
|
+
The inline base64 of each such part is never stored in the trace. The tracer removes it
|
|
80
|
+
before recording and replaces it with a small reference, so the Unity Catalog trace tables
|
|
81
|
+
stay readable. Multi-MB data URIs used to push large invoice traces past the SQL inline read
|
|
82
|
+
limit (issue #21). A remote `http(s)://` image URL is already a small reference, so this step
|
|
83
|
+
keeps it verbatim. The text cap below, if enabled, still truncates any string over its
|
|
84
|
+
threshold, URLs included. The transform runs on a copy of the inputs, so the live message
|
|
85
|
+
sent to the model is untouched and prompt caching is unaffected.
|
|
86
|
+
|
|
87
|
+
By default a part becomes a `{"type": ..., "_omitted": true, "bytes": N}` placeholder. To
|
|
88
|
+
store a meaningful reference instead, such as the Unity Catalog volume path the image was
|
|
89
|
+
loaded from so it can be re-fetched later, pass a `content_ref_resolver`:
|
|
90
|
+
|
|
91
|
+
```python
|
|
92
|
+
from databricks_langgraph_tracer import (
|
|
93
|
+
ContentPartContext,
|
|
94
|
+
configure_databricks_tracing,
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
def image_ref(part: dict, ctx: ContentPartContext) -> dict | None:
|
|
98
|
+
# ctx.metadata is the run metadata. Pass per-run data (e.g. a source volume
|
|
99
|
+
# path) via the invoke config's `metadata`, which propagates to the LLM run
|
|
100
|
+
# alongside langgraph_node etc.
|
|
101
|
+
path = ctx.metadata.get("encoded_images_path")
|
|
102
|
+
if path:
|
|
103
|
+
return {"type": part.get("type"), "ref": path, "page": ctx.index}
|
|
104
|
+
return None # fall back to the default placeholder
|
|
105
|
+
|
|
106
|
+
configure_databricks_tracing(experiment_id="...", content_ref_resolver=image_ref)
|
|
107
|
+
|
|
108
|
+
# ... then carry the per-run ref data on the invoke config metadata:
|
|
109
|
+
graph.invoke(state, config={
|
|
110
|
+
"callbacks": get_tracing_callbacks(),
|
|
111
|
+
"metadata": {"encoded_images_path": "/Volumes/cat/sch/vol/inv/pages.txt"},
|
|
112
|
+
})
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
The tracer calls the resolver once per multimodal part with the part and a
|
|
116
|
+
`ContentPartContext`. The context carries `index`, the part's position in the message
|
|
117
|
+
content array, which for a one-image-per-page invoice is the page number. It also carries
|
|
118
|
+
`bytes`, the inline payload length, and the run `metadata`. Return a dict to store as the
|
|
119
|
+
reference, or `None` for the default placeholder. `DatabricksLangGraphTracer` also accepts
|
|
120
|
+
`content_ref_resolver=` for per-graph wiring.
|
|
121
|
+
|
|
122
|
+
No inline image bytes ever reach the trace. If a resolver result re-introduces an inline
|
|
123
|
+
payload anywhere in the returned object, whether a `data:` URI, a recognized base64 content
|
|
124
|
+
part, or the part's own payload echoed back under another key, the tracer rejects it and
|
|
125
|
+
uses the placeholder. Beyond that, keep the reference small. The library strips inline
|
|
126
|
+
payloads but does not otherwise bound what a resolver returns, so a fabricated large string
|
|
127
|
+
under a custom key is the caller's problem. `ctx.metadata` is a shallow copy of the run
|
|
128
|
+
metadata, so setting top-level keys in the resolver cannot corrupt run state. Its nested
|
|
129
|
+
values are shared, so don't mutate those.
|
|
130
|
+
|
|
131
|
+
### Capping large text
|
|
132
|
+
|
|
133
|
+
Multimodal externalization handles inline bytes, but large plain text can also push a trace
|
|
134
|
+
past the SQL inline read limit (issue #23). A classification vocabulary or an aggregated
|
|
135
|
+
result set threaded through every fan-out span's inputs and outputs is enough. Set
|
|
136
|
+
`max_string_chars`, or `DATABRICKS_TRACING_MAX_STRING_CHARS`, to truncate it:
|
|
137
|
+
|
|
138
|
+
```python
|
|
139
|
+
configure_databricks_tracing(experiment_id="...", max_string_chars=50_000)
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
When set, any string value longer than the threshold in a span's inputs or outputs is
|
|
143
|
+
replaced with a placeholder:
|
|
144
|
+
|
|
145
|
+
```json
|
|
146
|
+
{"_truncated": true, "chars": 812345, "bytes": 812345, "preview": "<the first 256 chars>"}
|
|
147
|
+
```
|
|
148
|
+
|
|
149
|
+
The cap is off by default. Generic truncation costs debuggability, so you choose the
|
|
150
|
+
threshold. It caps string values only, never keys or structural fields, and runs on a copy,
|
|
151
|
+
so the live messages are untouched. It also reaches text nested inside Pydantic models,
|
|
152
|
+
dataclasses and tuples, such as a model a node returns as `final_output`, by normalizing
|
|
153
|
+
them to the same shape MLflow records. `DatabricksLangGraphTracer` also accepts
|
|
154
|
+
`max_string_chars=` for per-graph wiring.
|
|
155
|
+
|
|
156
|
+
This is a separate knob from the multimodal handling above. `content_ref_resolver` chooses
|
|
157
|
+
references for inline image, PDF and file bytes. `max_string_chars` caps arbitrary text and
|
|
158
|
+
never truncates a resolver's reference. Because it is generic, it also truncates any other
|
|
159
|
+
string over the threshold, including a remote `http(s)://` image URL the multimodal step
|
|
160
|
+
keeps verbatim, so set the threshold well above your reference and URL lengths. It is a
|
|
161
|
+
per-leaf limit, not a per-trace byte budget. Enough sub-threshold leaves can still sum past
|
|
162
|
+
the limit, so for the heaviest spans also record less. Pass ids or references through node
|
|
163
|
+
state rather than full payloads.
|
|
164
|
+
|
|
165
|
+
The threshold counts characters. This package counts code points and the TypeScript
|
|
166
|
+
package counts UTF-16 units, so the two can differ on non-BMP text. The inline limit is in
|
|
167
|
+
bytes, and multibyte text can be up to four times larger in bytes than in characters, so for
|
|
168
|
+
CJK or emoji-heavy content size the cap below a quarter of the limit. The placeholder's
|
|
169
|
+
`bytes` field always reports the exact UTF-8 size of the original.
|
|
170
|
+
|
|
171
|
+
### Autolog fallback
|
|
172
|
+
|
|
173
|
+
`mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches. It
|
|
174
|
+
emits a reduced schema, everything except a tracer-computed `mlflow.trace.cost` rollup. The
|
|
175
|
+
backend may still aggregate that server-side. The default `mode="tracer"`, the custom
|
|
176
|
+
`BaseTracer`, is the full-parity path.
|
|
177
|
+
|
|
178
|
+
## Development
|
|
179
|
+
|
|
180
|
+
```bash
|
|
181
|
+
cd python
|
|
182
|
+
uv sync # latest allowed mlflow (ceiling)
|
|
183
|
+
uv run pytest # unit + lifecycle tests
|
|
184
|
+
uv run ruff check src tests
|
|
185
|
+
uv run ty check
|
|
186
|
+
uv run python scripts/generate_keys.py --check # schema/keys parity
|
|
187
|
+
|
|
188
|
+
# mlflow floor matrix (CI runs both cells via UV_RESOLUTION):
|
|
189
|
+
UV_RESOLUTION=lowest-direct uv sync && UV_RESOLUTION=lowest-direct uv run pytest
|
|
190
|
+
```
|
|
191
|
+
|
|
192
|
+
Tests use a local sqlite MLflow tracking backend. Checks that need a live Databricks backend
|
|
193
|
+
are marked `integration`.
|
|
194
|
+
|
|
195
|
+
## Changelog
|
|
196
|
+
|
|
197
|
+
See [`CHANGELOG.md`](CHANGELOG.md).
|
|
@@ -1,224 +0,0 @@
|
|
|
1
|
-
Metadata-Version: 2.5
|
|
2
|
-
Name: insightfactory-databricks-langgraph-tracer
|
|
3
|
-
Version: 1.0.0.dev10
|
|
4
|
-
Summary: LangGraph tracing for Databricks MLflow
|
|
5
|
-
Project-URL: Homepage, https://insightfactory.ai
|
|
6
|
-
Author-email: "insightfactory.ai" <support@insightfactory.ai>
|
|
7
|
-
License-Expression: LicenseRef-Proprietary
|
|
8
|
-
License-File: LICENSE
|
|
9
|
-
License-File: THIRD_PARTY_NOTICES
|
|
10
|
-
Keywords: databricks,langchain,langgraph,mlflow,tracing
|
|
11
|
-
Classifier: Development Status :: 5 - Production/Stable
|
|
12
|
-
Classifier: Intended Audience :: Developers
|
|
13
|
-
Classifier: License :: Other/Proprietary License
|
|
14
|
-
Classifier: Programming Language :: Python :: 3 :: Only
|
|
15
|
-
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
-
Classifier: Typing :: Typed
|
|
17
|
-
Requires-Python: <3.13,>=3.12
|
|
18
|
-
Requires-Dist: databricks-sdk>=0.82.0
|
|
19
|
-
Requires-Dist: langchain-core>=1.0.0
|
|
20
|
-
Requires-Dist: langchain>=1.3.9
|
|
21
|
-
Requires-Dist: langgraph>=1.0.0
|
|
22
|
-
Requires-Dist: mlflow<4.0.0,>=3.15.0
|
|
23
|
-
Description-Content-Type: text/markdown
|
|
24
|
-
|
|
25
|
-
# insightfactory-databricks-langgraph-tracer
|
|
26
|
-
|
|
27
|
-
LangGraph tracer for Databricks MLflow. It uses a custom LangChain `BaseTracer` to write
|
|
28
|
-
trace tags, metadata, token usage, and cost rollups before the root span ends.
|
|
29
|
-
|
|
30
|
-
## Install
|
|
31
|
-
|
|
32
|
-
Published on PyPI:
|
|
33
|
-
|
|
34
|
-
```bash
|
|
35
|
-
uv add insightfactory-databricks-langgraph-tracer
|
|
36
|
-
# or: pip install insightfactory-databricks-langgraph-tracer
|
|
37
|
-
```
|
|
38
|
-
|
|
39
|
-
Requires Python 3.12 and resolves `mlflow>=3.15.0,<4`.
|
|
40
|
-
|
|
41
|
-
## Quickstart
|
|
42
|
-
|
|
43
|
-
```python
|
|
44
|
-
from databricks_langgraph_tracer import (
|
|
45
|
-
configure_databricks_tracing,
|
|
46
|
-
get_tracing_callbacks,
|
|
47
|
-
)
|
|
48
|
-
|
|
49
|
-
# 1. Bootstrap once at startup (reads env-first; kwargs override).
|
|
50
|
-
configure_databricks_tracing(experiment_id="<mlflow-experiment-id>", source="my-agent")
|
|
51
|
-
|
|
52
|
-
# 2. Attach the callbacks to your LangGraph / LangChain run.
|
|
53
|
-
graph = build_graph().with_config({"callbacks": get_tracing_callbacks()})
|
|
54
|
-
graph.invoke(state)
|
|
55
|
-
```
|
|
56
|
-
|
|
57
|
-
Traces appear in the configured Databricks MLflow experiment carrying the full shared schema:
|
|
58
|
-
the `source` tag, session/user/thread metadata, per-span model / provider / token-usage + cost,
|
|
59
|
-
the trace-level `mlflow.trace.cost` rollup, and the per-model `cost.by_model` cost/token rollup tag.
|
|
60
|
-
|
|
61
|
-
## Configuration
|
|
62
|
-
|
|
63
|
-
Env-first; any kwarg to `configure_databricks_tracing(...)` overrides the matching env var.
|
|
64
|
-
|
|
65
|
-
| Setting | Env var | Notes |
|
|
66
|
-
|---------|---------|-------|
|
|
67
|
-
| Tracking URI | `MLFLOW_TRACKING_URI` | `databricks` or `databricks://<profile>` (required) |
|
|
68
|
-
| Experiment | `MLFLOW_EXPERIMENT_ID` | by **id** only (required) |
|
|
69
|
-
| Source tag | — | `source=` kwarg, default `langgraph` |
|
|
70
|
-
| Multimodal refs | — | inline image/PDF/file bytes are externalized to a reference; `content_ref_resolver=` chooses it — see below |
|
|
71
|
-
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | opt-in; `max_string_chars=` truncates over-long plain-text span content — see below |
|
|
72
|
-
| Disable | `TESTING` / `BUILDING` = `true`, or `enabled=False` | no-op — the only non-raising path |
|
|
73
|
-
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | see below |
|
|
74
|
-
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID` → `DATABRICKS_WAREHOUSE_ID` | required when UC-backed |
|
|
75
|
-
|
|
76
|
-
**Auth** (resolved by `databricks_utils`): a service principal
|
|
77
|
-
(`DATABRICKS_HOST` / `DATABRICKS_CLIENT_ID` / `DATABRICKS_CLIENT_SECRET`) **or** a CLI profile
|
|
78
|
-
(`DATABRICKS_CONFIG_PROFILE` / `profile=`).
|
|
79
|
-
|
|
80
|
-
Missing required config **raises** `DatabricksTracingConfigurationError` (fail-fast) — disable
|
|
81
|
-
explicitly for local/dev runs.
|
|
82
|
-
|
|
83
|
-
### Unity-Catalog-backed tracing
|
|
84
|
-
|
|
85
|
-
For experiments whose traces are stored in **Unity Catalog** `_otel_*` tables, set
|
|
86
|
-
`uc_tracing=True` (or `MLFLOW_TRACING_UC_BACKED=true`) and provide a SQL warehouse. The library
|
|
87
|
-
validates the warehouse, resolves the experiment's UC trace location from its binding tag, and
|
|
88
|
-
passes it to `set_experiment` so spans actually persist to the `_otel_spans` table — without it,
|
|
89
|
-
MLflow silently skips span export to UC. Classic (workspace) experiments need none of this and
|
|
90
|
-
are the default. UC-backed tracing requires an MLflow release that provides the
|
|
91
|
-
`UnityCatalog` trace-location API.
|
|
92
|
-
|
|
93
|
-
### Multimodal inputs (image / PDF / file)
|
|
94
|
-
|
|
95
|
-
Chat-model spans record their inputs as **structured messages**, preserving the
|
|
96
|
-
multimodal content parts a graph sends to the model — OpenAI / LangChain
|
|
97
|
-
`image_url`, OpenAI `file`, and Anthropic `image` / `document`. (Internally the
|
|
98
|
-
tracer runs the LangChain `BaseTracer` in `original+chat` mode; the default would
|
|
99
|
-
otherwise flatten chat messages to a text-only `prompts` string and drop every
|
|
100
|
-
attachment.)
|
|
101
|
-
|
|
102
|
-
The inline base64 of each such part is **never stored** in the trace: it is removed
|
|
103
|
-
before recording and replaced with a lightweight reference, so the Unity Catalog
|
|
104
|
-
trace tables stay readable (multi-MB data URIs previously pushed large invoice
|
|
105
|
-
traces past the SQL inline read limit — issue #21). A remote `http(s)://` image URL
|
|
106
|
-
is already a small reference, so this step keeps it verbatim (but see the text cap
|
|
107
|
-
below — if enabled, it still truncates *any* string over its threshold, URLs
|
|
108
|
-
included). The transform runs on a copy of the inputs, so the live message sent to
|
|
109
|
-
the model is untouched and prompt caching is unaffected.
|
|
110
|
-
|
|
111
|
-
By default a part becomes a `{"type": ..., "_omitted": true, "bytes": N}`
|
|
112
|
-
placeholder. To store a meaningful reference instead — e.g. the Unity Catalog volume
|
|
113
|
-
path the image was loaded from, so it can be re-fetched at runtime — pass a
|
|
114
|
-
`content_ref_resolver`:
|
|
115
|
-
|
|
116
|
-
```python
|
|
117
|
-
from databricks_langgraph_tracer import (
|
|
118
|
-
ContentPartContext,
|
|
119
|
-
configure_databricks_tracing,
|
|
120
|
-
)
|
|
121
|
-
|
|
122
|
-
def image_ref(part: dict, ctx: ContentPartContext) -> dict | None:
|
|
123
|
-
# ctx.metadata is the run metadata — pass per-run data (e.g. a source volume
|
|
124
|
-
# path) via the invoke config's `metadata`, which propagates to the LLM run
|
|
125
|
-
# alongside langgraph_node etc.
|
|
126
|
-
path = ctx.metadata.get("encoded_images_path")
|
|
127
|
-
if path:
|
|
128
|
-
return {"type": part.get("type"), "ref": path, "page": ctx.index}
|
|
129
|
-
return None # fall back to the default placeholder
|
|
130
|
-
|
|
131
|
-
configure_databricks_tracing(experiment_id="...", content_ref_resolver=image_ref)
|
|
132
|
-
|
|
133
|
-
# ... then carry the per-run ref data on the invoke config metadata:
|
|
134
|
-
graph.invoke(state, config={
|
|
135
|
-
"callbacks": get_tracing_callbacks(),
|
|
136
|
-
"metadata": {"encoded_images_path": "/Volumes/cat/sch/vol/inv/pages.txt"},
|
|
137
|
-
})
|
|
138
|
-
```
|
|
139
|
-
|
|
140
|
-
The resolver is called once per multimodal part with the part and a
|
|
141
|
-
`ContentPartContext` — `index` (position within the message content array, i.e. the
|
|
142
|
-
page number for a one-image-per-page invoice), `bytes` (the inline payload length),
|
|
143
|
-
and the run `metadata`. Return a dict to store as the reference, or `None` for the
|
|
144
|
-
default placeholder. It is also accepted by
|
|
145
|
-
`DatabricksLangGraphTracer(content_ref_resolver=...)` for per-graph wiring.
|
|
146
|
-
|
|
147
|
-
The library guarantees no inline image bytes are stored: if a resolver result
|
|
148
|
-
re-introduces an inline payload (a `data:` URI, a recognized base64 content part, or
|
|
149
|
-
the part's own payload echoed back under any key — anywhere in the returned object),
|
|
150
|
-
it is rejected and the placeholder is used. Beyond that, keep the reference
|
|
151
|
-
**compact** — a *fabricated* large string under a custom key is the consumer's
|
|
152
|
-
responsibility (the library strips inline payloads but does not otherwise bound what
|
|
153
|
-
a resolver returns).
|
|
154
|
-
`ctx.metadata` is a shallow copy of the run metadata, so a resolver cannot corrupt run
|
|
155
|
-
state by setting top-level keys (don't mutate its nested values, which are shared).
|
|
156
|
-
|
|
157
|
-
### Capping large text (issue #23)
|
|
158
|
-
|
|
159
|
-
Multimodal externalization handles inline *bytes*, but large **plain text** can also
|
|
160
|
-
push a trace past the SQL inline read limit — e.g. a classification vocabulary or
|
|
161
|
-
aggregated result set threaded through every fan-out span's inputs and outputs. Set
|
|
162
|
-
`max_string_chars` (or `DATABRICKS_TRACING_MAX_STRING_CHARS`) to truncate it:
|
|
163
|
-
|
|
164
|
-
```python
|
|
165
|
-
configure_databricks_tracing(experiment_id="...", max_string_chars=50_000)
|
|
166
|
-
```
|
|
167
|
-
|
|
168
|
-
When set, any string **value** longer than the threshold — in a span's inputs *or*
|
|
169
|
-
outputs — is replaced with a compact placeholder:
|
|
170
|
-
|
|
171
|
-
```json
|
|
172
|
-
{"_truncated": true, "chars": 812345, "bytes": 812345, "preview": "first 256 chars…"}
|
|
173
|
-
```
|
|
174
|
-
|
|
175
|
-
It is **opt-in / off by default** (generic truncation costs debuggability, so you
|
|
176
|
-
choose the threshold). It caps string values only — never keys or structural fields —
|
|
177
|
-
and runs on a copy, so the live messages are untouched. The cap also reaches text
|
|
178
|
-
nested inside **Pydantic models / dataclasses / tuples** (e.g. a model returned by a
|
|
179
|
-
node such as `final_output`), normalizing them to the same shape MLflow records. Also
|
|
180
|
-
accepted by `DatabricksLangGraphTracer(max_string_chars=...)` for per-graph wiring.
|
|
181
|
-
|
|
182
|
-
This is a **distinct** knob from the multimodal handling above: `content_ref_resolver`
|
|
183
|
-
chooses references for inline image/PDF/file bytes; `max_string_chars` caps arbitrary
|
|
184
|
-
text (and never truncates a resolver's reference). Because it is generic, it also
|
|
185
|
-
truncates any other over-threshold string — including a remote `http(s)://` image URL
|
|
186
|
-
the multimodal step keeps verbatim — so set the threshold comfortably above your
|
|
187
|
-
reference / URL lengths. Note it is a **per-leaf** mitigation, not a hard per-trace
|
|
188
|
-
byte budget — enough sub-threshold leaves can still sum past the limit — so for the
|
|
189
|
-
heaviest spans also reduce what you record (pass ids/references through node state
|
|
190
|
-
rather than full payloads).
|
|
191
|
-
|
|
192
|
-
The threshold counts **characters** (code points here; the TypeScript package counts
|
|
193
|
-
UTF-16 units, so the two can differ on non-BMP text), while the inline limit is in
|
|
194
|
-
bytes — multibyte text can be up to ~4× larger in bytes than characters, so for
|
|
195
|
-
CJK/emoji-heavy content size the cap below `limit / 4`. The placeholder's `bytes`
|
|
196
|
-
field always reports the exact UTF-8 size of the original.
|
|
197
|
-
|
|
198
|
-
### Autolog fallback
|
|
199
|
-
|
|
200
|
-
`mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches and
|
|
201
|
-
emits a **reduced** schema — everything except a tracer-computed
|
|
202
|
-
`mlflow.trace.cost` rollup (the backend may still aggregate it server-side). The default
|
|
203
|
-
`mode="tracer"` (custom `BaseTracer`) is the full-parity path.
|
|
204
|
-
|
|
205
|
-
## Development
|
|
206
|
-
|
|
207
|
-
```bash
|
|
208
|
-
cd python
|
|
209
|
-
uv sync # latest allowed mlflow (ceiling)
|
|
210
|
-
uv run pytest # unit + lifecycle tests
|
|
211
|
-
uv run ruff check src tests
|
|
212
|
-
uv run ty check
|
|
213
|
-
uv run python scripts/generate_keys.py --check # schema/keys parity
|
|
214
|
-
|
|
215
|
-
# mlflow floor matrix (CI runs both cells via UV_RESOLUTION):
|
|
216
|
-
UV_RESOLUTION=lowest-direct uv sync && UV_RESOLUTION=lowest-direct uv run pytest
|
|
217
|
-
```
|
|
218
|
-
|
|
219
|
-
Tests use a local sqlite MLflow tracking backend; checks that need a live Databricks backend are
|
|
220
|
-
marked `integration`.
|
|
221
|
-
|
|
222
|
-
## Changelog
|
|
223
|
-
|
|
224
|
-
See [`CHANGELOG.md`](CHANGELOG.md).
|
|
@@ -1,200 +0,0 @@
|
|
|
1
|
-
# insightfactory-databricks-langgraph-tracer
|
|
2
|
-
|
|
3
|
-
LangGraph tracer for Databricks MLflow. It uses a custom LangChain `BaseTracer` to write
|
|
4
|
-
trace tags, metadata, token usage, and cost rollups before the root span ends.
|
|
5
|
-
|
|
6
|
-
## Install
|
|
7
|
-
|
|
8
|
-
Published on PyPI:
|
|
9
|
-
|
|
10
|
-
```bash
|
|
11
|
-
uv add insightfactory-databricks-langgraph-tracer
|
|
12
|
-
# or: pip install insightfactory-databricks-langgraph-tracer
|
|
13
|
-
```
|
|
14
|
-
|
|
15
|
-
Requires Python 3.12 and resolves `mlflow>=3.15.0,<4`.
|
|
16
|
-
|
|
17
|
-
## Quickstart
|
|
18
|
-
|
|
19
|
-
```python
|
|
20
|
-
from databricks_langgraph_tracer import (
|
|
21
|
-
configure_databricks_tracing,
|
|
22
|
-
get_tracing_callbacks,
|
|
23
|
-
)
|
|
24
|
-
|
|
25
|
-
# 1. Bootstrap once at startup (reads env-first; kwargs override).
|
|
26
|
-
configure_databricks_tracing(experiment_id="<mlflow-experiment-id>", source="my-agent")
|
|
27
|
-
|
|
28
|
-
# 2. Attach the callbacks to your LangGraph / LangChain run.
|
|
29
|
-
graph = build_graph().with_config({"callbacks": get_tracing_callbacks()})
|
|
30
|
-
graph.invoke(state)
|
|
31
|
-
```
|
|
32
|
-
|
|
33
|
-
Traces appear in the configured Databricks MLflow experiment carrying the full shared schema:
|
|
34
|
-
the `source` tag, session/user/thread metadata, per-span model / provider / token-usage + cost,
|
|
35
|
-
the trace-level `mlflow.trace.cost` rollup, and the per-model `cost.by_model` cost/token rollup tag.
|
|
36
|
-
|
|
37
|
-
## Configuration
|
|
38
|
-
|
|
39
|
-
Env-first; any kwarg to `configure_databricks_tracing(...)` overrides the matching env var.
|
|
40
|
-
|
|
41
|
-
| Setting | Env var | Notes |
|
|
42
|
-
|---------|---------|-------|
|
|
43
|
-
| Tracking URI | `MLFLOW_TRACKING_URI` | `databricks` or `databricks://<profile>` (required) |
|
|
44
|
-
| Experiment | `MLFLOW_EXPERIMENT_ID` | by **id** only (required) |
|
|
45
|
-
| Source tag | — | `source=` kwarg, default `langgraph` |
|
|
46
|
-
| Multimodal refs | — | inline image/PDF/file bytes are externalized to a reference; `content_ref_resolver=` chooses it — see below |
|
|
47
|
-
| Text cap | `DATABRICKS_TRACING_MAX_STRING_CHARS` | opt-in; `max_string_chars=` truncates over-long plain-text span content — see below |
|
|
48
|
-
| Disable | `TESTING` / `BUILDING` = `true`, or `enabled=False` | no-op — the only non-raising path |
|
|
49
|
-
| UC-backed tracing | `MLFLOW_TRACING_UC_BACKED=true` or `uc_tracing=True` | see below |
|
|
50
|
-
| SQL warehouse | `MLFLOW_TRACING_SQL_WAREHOUSE_ID` → `DATABRICKS_WAREHOUSE_ID` | required when UC-backed |
|
|
51
|
-
|
|
52
|
-
**Auth** (resolved by `databricks_utils`): a service principal
|
|
53
|
-
(`DATABRICKS_HOST` / `DATABRICKS_CLIENT_ID` / `DATABRICKS_CLIENT_SECRET`) **or** a CLI profile
|
|
54
|
-
(`DATABRICKS_CONFIG_PROFILE` / `profile=`).
|
|
55
|
-
|
|
56
|
-
Missing required config **raises** `DatabricksTracingConfigurationError` (fail-fast) — disable
|
|
57
|
-
explicitly for local/dev runs.
|
|
58
|
-
|
|
59
|
-
### Unity-Catalog-backed tracing
|
|
60
|
-
|
|
61
|
-
For experiments whose traces are stored in **Unity Catalog** `_otel_*` tables, set
|
|
62
|
-
`uc_tracing=True` (or `MLFLOW_TRACING_UC_BACKED=true`) and provide a SQL warehouse. The library
|
|
63
|
-
validates the warehouse, resolves the experiment's UC trace location from its binding tag, and
|
|
64
|
-
passes it to `set_experiment` so spans actually persist to the `_otel_spans` table — without it,
|
|
65
|
-
MLflow silently skips span export to UC. Classic (workspace) experiments need none of this and
|
|
66
|
-
are the default. UC-backed tracing requires an MLflow release that provides the
|
|
67
|
-
`UnityCatalog` trace-location API.
|
|
68
|
-
|
|
69
|
-
### Multimodal inputs (image / PDF / file)
|
|
70
|
-
|
|
71
|
-
Chat-model spans record their inputs as **structured messages**, preserving the
|
|
72
|
-
multimodal content parts a graph sends to the model — OpenAI / LangChain
|
|
73
|
-
`image_url`, OpenAI `file`, and Anthropic `image` / `document`. (Internally the
|
|
74
|
-
tracer runs the LangChain `BaseTracer` in `original+chat` mode; the default would
|
|
75
|
-
otherwise flatten chat messages to a text-only `prompts` string and drop every
|
|
76
|
-
attachment.)
|
|
77
|
-
|
|
78
|
-
The inline base64 of each such part is **never stored** in the trace: it is removed
|
|
79
|
-
before recording and replaced with a lightweight reference, so the Unity Catalog
|
|
80
|
-
trace tables stay readable (multi-MB data URIs previously pushed large invoice
|
|
81
|
-
traces past the SQL inline read limit — issue #21). A remote `http(s)://` image URL
|
|
82
|
-
is already a small reference, so this step keeps it verbatim (but see the text cap
|
|
83
|
-
below — if enabled, it still truncates *any* string over its threshold, URLs
|
|
84
|
-
included). The transform runs on a copy of the inputs, so the live message sent to
|
|
85
|
-
the model is untouched and prompt caching is unaffected.
|
|
86
|
-
|
|
87
|
-
By default a part becomes a `{"type": ..., "_omitted": true, "bytes": N}`
|
|
88
|
-
placeholder. To store a meaningful reference instead — e.g. the Unity Catalog volume
|
|
89
|
-
path the image was loaded from, so it can be re-fetched at runtime — pass a
|
|
90
|
-
`content_ref_resolver`:
|
|
91
|
-
|
|
92
|
-
```python
|
|
93
|
-
from databricks_langgraph_tracer import (
|
|
94
|
-
ContentPartContext,
|
|
95
|
-
configure_databricks_tracing,
|
|
96
|
-
)
|
|
97
|
-
|
|
98
|
-
def image_ref(part: dict, ctx: ContentPartContext) -> dict | None:
|
|
99
|
-
# ctx.metadata is the run metadata — pass per-run data (e.g. a source volume
|
|
100
|
-
# path) via the invoke config's `metadata`, which propagates to the LLM run
|
|
101
|
-
# alongside langgraph_node etc.
|
|
102
|
-
path = ctx.metadata.get("encoded_images_path")
|
|
103
|
-
if path:
|
|
104
|
-
return {"type": part.get("type"), "ref": path, "page": ctx.index}
|
|
105
|
-
return None # fall back to the default placeholder
|
|
106
|
-
|
|
107
|
-
configure_databricks_tracing(experiment_id="...", content_ref_resolver=image_ref)
|
|
108
|
-
|
|
109
|
-
# ... then carry the per-run ref data on the invoke config metadata:
|
|
110
|
-
graph.invoke(state, config={
|
|
111
|
-
"callbacks": get_tracing_callbacks(),
|
|
112
|
-
"metadata": {"encoded_images_path": "/Volumes/cat/sch/vol/inv/pages.txt"},
|
|
113
|
-
})
|
|
114
|
-
```
|
|
115
|
-
|
|
116
|
-
The resolver is called once per multimodal part with the part and a
|
|
117
|
-
`ContentPartContext` — `index` (position within the message content array, i.e. the
|
|
118
|
-
page number for a one-image-per-page invoice), `bytes` (the inline payload length),
|
|
119
|
-
and the run `metadata`. Return a dict to store as the reference, or `None` for the
|
|
120
|
-
default placeholder. It is also accepted by
|
|
121
|
-
`DatabricksLangGraphTracer(content_ref_resolver=...)` for per-graph wiring.
|
|
122
|
-
|
|
123
|
-
The library guarantees no inline image bytes are stored: if a resolver result
|
|
124
|
-
re-introduces an inline payload (a `data:` URI, a recognized base64 content part, or
|
|
125
|
-
the part's own payload echoed back under any key — anywhere in the returned object),
|
|
126
|
-
it is rejected and the placeholder is used. Beyond that, keep the reference
|
|
127
|
-
**compact** — a *fabricated* large string under a custom key is the consumer's
|
|
128
|
-
responsibility (the library strips inline payloads but does not otherwise bound what
|
|
129
|
-
a resolver returns).
|
|
130
|
-
`ctx.metadata` is a shallow copy of the run metadata, so a resolver cannot corrupt run
|
|
131
|
-
state by setting top-level keys (don't mutate its nested values, which are shared).
|
|
132
|
-
|
|
133
|
-
### Capping large text (issue #23)
|
|
134
|
-
|
|
135
|
-
Multimodal externalization handles inline *bytes*, but large **plain text** can also
|
|
136
|
-
push a trace past the SQL inline read limit — e.g. a classification vocabulary or
|
|
137
|
-
aggregated result set threaded through every fan-out span's inputs and outputs. Set
|
|
138
|
-
`max_string_chars` (or `DATABRICKS_TRACING_MAX_STRING_CHARS`) to truncate it:
|
|
139
|
-
|
|
140
|
-
```python
|
|
141
|
-
configure_databricks_tracing(experiment_id="...", max_string_chars=50_000)
|
|
142
|
-
```
|
|
143
|
-
|
|
144
|
-
When set, any string **value** longer than the threshold — in a span's inputs *or*
|
|
145
|
-
outputs — is replaced with a compact placeholder:
|
|
146
|
-
|
|
147
|
-
```json
|
|
148
|
-
{"_truncated": true, "chars": 812345, "bytes": 812345, "preview": "first 256 chars…"}
|
|
149
|
-
```
|
|
150
|
-
|
|
151
|
-
It is **opt-in / off by default** (generic truncation costs debuggability, so you
|
|
152
|
-
choose the threshold). It caps string values only — never keys or structural fields —
|
|
153
|
-
and runs on a copy, so the live messages are untouched. The cap also reaches text
|
|
154
|
-
nested inside **Pydantic models / dataclasses / tuples** (e.g. a model returned by a
|
|
155
|
-
node such as `final_output`), normalizing them to the same shape MLflow records. Also
|
|
156
|
-
accepted by `DatabricksLangGraphTracer(max_string_chars=...)` for per-graph wiring.
|
|
157
|
-
|
|
158
|
-
This is a **distinct** knob from the multimodal handling above: `content_ref_resolver`
|
|
159
|
-
chooses references for inline image/PDF/file bytes; `max_string_chars` caps arbitrary
|
|
160
|
-
text (and never truncates a resolver's reference). Because it is generic, it also
|
|
161
|
-
truncates any other over-threshold string — including a remote `http(s)://` image URL
|
|
162
|
-
the multimodal step keeps verbatim — so set the threshold comfortably above your
|
|
163
|
-
reference / URL lengths. Note it is a **per-leaf** mitigation, not a hard per-trace
|
|
164
|
-
byte budget — enough sub-threshold leaves can still sum past the limit — so for the
|
|
165
|
-
heaviest spans also reduce what you record (pass ids/references through node state
|
|
166
|
-
rather than full payloads).
|
|
167
|
-
|
|
168
|
-
The threshold counts **characters** (code points here; the TypeScript package counts
|
|
169
|
-
UTF-16 units, so the two can differ on non-BMP text), while the inline limit is in
|
|
170
|
-
bytes — multibyte text can be up to ~4× larger in bytes than characters, so for
|
|
171
|
-
CJK/emoji-heavy content size the cap below `limit / 4`. The placeholder's `bytes`
|
|
172
|
-
field always reports the exact UTF-8 size of the original.
|
|
173
|
-
|
|
174
|
-
### Autolog fallback
|
|
175
|
-
|
|
176
|
-
`mode="autolog"` wires MLflow's built-in LangChain autologging plus compatibility patches and
|
|
177
|
-
emits a **reduced** schema — everything except a tracer-computed
|
|
178
|
-
`mlflow.trace.cost` rollup (the backend may still aggregate it server-side). The default
|
|
179
|
-
`mode="tracer"` (custom `BaseTracer`) is the full-parity path.
|
|
180
|
-
|
|
181
|
-
## Development
|
|
182
|
-
|
|
183
|
-
```bash
|
|
184
|
-
cd python
|
|
185
|
-
uv sync # latest allowed mlflow (ceiling)
|
|
186
|
-
uv run pytest # unit + lifecycle tests
|
|
187
|
-
uv run ruff check src tests
|
|
188
|
-
uv run ty check
|
|
189
|
-
uv run python scripts/generate_keys.py --check # schema/keys parity
|
|
190
|
-
|
|
191
|
-
# mlflow floor matrix (CI runs both cells via UV_RESOLUTION):
|
|
192
|
-
UV_RESOLUTION=lowest-direct uv sync && UV_RESOLUTION=lowest-direct uv run pytest
|
|
193
|
-
```
|
|
194
|
-
|
|
195
|
-
Tests use a local sqlite MLflow tracking backend; checks that need a live Databricks backend are
|
|
196
|
-
marked `integration`.
|
|
197
|
-
|
|
198
|
-
## Changelog
|
|
199
|
-
|
|
200
|
-
See [`CHANGELOG.md`](CHANGELOG.md).
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|