Set a readable root-span request and response - #31
Conversation
Signed-off-by: Adam Gurary <guraryadam@gmail.com>
span.set_outputs() inside a @mlflow.trace body is overwritten by the
decorator after the function exits (last-write-wins). Empirically
confirmed on MLflow 3.11.1.dev0: the old example left root span outputs
at 58k chars. Two correct patterns now shown: update_current_trace
(request_preview/response_preview) for compact Trace list UI previews
while keeping the decorator, and a manual start_span entry point for
compact storage on the root span itself. Also fixes None-safety on
note_text slice: (state.get("note_text") or "")[:200].
Signed-off-by: Adam Gurary <guraryadam@gmail.com>
1aa0c12 to
ef4ef3e
Compare
|
@B-Step62 ready for review. |
|
|
||
| **Span types**: `LLM`, `CHAIN`, `TOOL`, `AGENT`, `RETRIEVER`, `EMBEDDING`, `RERANKER`, `PARSER`, `UNKNOWN` | ||
|
|
||
| **Caution: the decorator auto-captures raw function arguments AND return value.** `span.set_inputs()` inside the body correctly overrides the auto-captured inputs. `span.set_outputs()` inside the body does not: the decorator calls `span.set_outputs(return_value)` after the function exits, so the return value always wins. |
There was a problem hiding this comment.
@adamgurary Do you have a screenshot of the inputs/outputs for the current root span? Any way to make it render better by default rather than asking users to set them manually?
Asking this because we already have some parsing logic on frontend which handles cases like chat format.
There was a problem hiding this comment.
Good question, and it pushed me to check the actual behavior rather than assume it.
The auto-derivation does work, but only for OpenAI-shaped payloads. In mlflow/tracing/utils/truncation.py, set_request_response_preview() fills request_preview / response_preview only when the user has not set them, and _try_extract_messages() recognizes exactly three shapes: a top-level messages list, choices[0].message, and a Responses-API input list. When one matches, you get the last user/assistant message's text content, which renders well with no user action.
Anything else falls through to content or request_or_response, which is json.dumps(obj) truncated to the max length. That is the case this section is about: a LangGraph agent whose root span input/output is a state dict ({"customer": ..., "date": ..., "note_text": ..., "output_path": ...}). It is not chat-shaped, so there are no messages to extract, and the preview is the head of the serialized dict, cut mid-JSON.
So on "can it render better by default": for the chat-format case it already does, and this guidance should not tell people to override it. For arbitrary agent state, the default cannot do much better without guessing which keys matter, and getting that wrong is worse than a plain truncation. Picking the significant fields is a judgment call only the author can make.
I have pushed a revision that does two things:
- Scopes the guidance explicitly to non-chat-shaped inputs/outputs, and says that OpenAI chat-format payloads already get a readable preview automatically, so
update_current_traceis unnecessary there. - Drops the framing that implied previews always need manual setting.
That keeps the genuinely useful part (the set_outputs() asymmetry, which I did verify against fluent.py: the decorator calls span.set_inputs(inputs) before the body and span.set_outputs(result) after it, so an in-body set_inputs() survives and an in-body set_outputs() is overwritten) without duplicating what the frontend already handles.
On the screenshot: I do not have one I can share here. The reproduction was on an internal deployment, and I would rather describe the shape than post a capture with workload details in it. The state-dict shape above is the relevant part.
MLflow already derives readable trace-list previews for OpenAI chat-shaped payloads (messages, choices[0].message, Responses-API input) in mlflow/tracing/utils/truncation.py, so telling users to set previews manually there is wrong. Limit the guidance to root spans whose input or output is not chat-shaped, where the preview falls back to truncated JSON. Co-authored-by: Isaac Signed-off-by: Adam Gurary <guraryadam@gmail.com>
2 participants
Fixes #29
Summary
When a coding agent used the MLflow instrumentation skills to trace a LangGraph agent, the top of the trace showed an unreadable blob instead of the run inputs and outputs. The `@mlflow.trace` decorator on the entry point auto-recorded the compiled LangGraph object as the request and the entire final state dict as the response. The root span showed neither the customer name going in nor the note coming out.
Steps to reproduce
Actual behavior
The root-span request was the Python repr of the compiled graph object. The root-span response was the full accumulated state dictionary (transcript plus all research plus draft). Both were unreadable in the UI.
Root cause
The skill applied `@mlflow.trace` to a function whose input is a large framework object and whose output is a large state dict, without guiding the agent to record a compact, human-readable summary instead.
Expected behavior and fix
The MLflow instrumentation skills should teach the coding agent to set a compact, readable request and response on the root span, rather than letting `@mlflow.trace` capture raw framework objects. For a LangGraph agent, that means the meaningful inputs (customer and date) and the meaningful outputs (output path and a short preview). The goal is that the trace reads well in the UI by default.
Test setup
Verified on: MLflow 3.14.0, LangGraph 1.2.7, LangChain 1.3.11, databricks-langchain 0.17.0, Python 3.12.