feat(trace): generalize the RAG trace to all agents and export it to a file#42
Draft
pjmalandrino wants to merge 1 commit into
Draft
feat(trace): generalize the RAG trace to all agents and export it to a file#42pjmalandrino wants to merge 1 commit into
pjmalandrino wants to merge 1 commit into
Conversation
Contributor
|
✅ DCO Check Passed Thanks @pjmalandrino, all your commits are properly signed off. 🎉 |
Contributor
Merge Protections🔴 1 of 2 protections blocking · waiting on 👀 reviews
🔴 Require two reviewer for test updatesWaiting for
This rule is failing.When test data is updated, we require two reviewers
Show 1 satisfied protection🟢 Enforce conventional commitMake sure that we follow https://www.conventionalcommits.org/en/v1.0.0/
|
…a file Generalize the run_with_trace pattern from docling-project#39 into a generic, typed AgentTrace that every agent produces, and let the orchestrator compose the traces of the sub-agents it dispatches to into a tree that can be exported to a single JSON file for debugging. - New AgentStep / AgentTrace value objects (agent_trace.py). AgentTrace carries ordered steps, nested children (sub-agent traces), timing, model id and the produced document on `output` (excluded from serialization; result_name is the persisted pointer). SerializeAsAny on children preserves subclass fields. - BaseDoclingAgent.run_with_trace(): concrete default so every agent exposes a trace (timing + model + result) with no bespoke code; agents override to add steps. run() stays the source of truth for the document. - RAGTrace becomes a subclass of AgentTrace, so docling-project#39 is preserved (same fields, same construction, covariant return) and a RAG run nests into the tree. DoclingRAGAgent builds the answer doc onto output; run() returns it. - DoclingOrchestratorAgent.run_task_with_trace() composes the tree by recording each dispatched sub-agent trace; run_task() is unchanged and incurs no overhead. - LoggingConfig.trace_path + CLI export. Public re-exports: AgentTrace, AgentStep. No global state, no logging coupling: the trace is a returned value object. Design doc: docs/design/37-export-session-trace.md. Closes docling-project#37 Signed-off-by: Pier-Jean Malandrino <pierjean.malandrino@scub.net>
7cb2ec2 to
518f071
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements #37 ("Logs: export session logs") by generalizing the
run_with_tracepattern from #39 to every agent and letting the orchestrator compose the sub-agent traces into a tree that exports to a single JSON file for debugging.This is the "similar tracing capability across all agents, exportable to a file" follow-up suggested on #39 — built as a value object (no global state, no coupling to the logging subsystem).
Design doc:
docs/design/37-export-session-trace.md.The logical path from #39
#39gave RAG three layers (RAGIteration→RAGResult→RAGTrace) +run_with_trace(). This PR lifts that shape to the base class:AgentStep/AgentTrace(newagent_trace.py) — generic counterparts.AgentTraceholds orderedsteps, nestedchildren(sub-agent traces), timing, model id, and the produced document onoutput(kept in-memory, excluded from serialization;result_nameis the persisted pointer).childrenusesSerializeAsAnyso nested subclass fields (e.g. RAG) survive serialization.BaseDoclingAgent.run_with_trace()— concrete default: timesrun()and wraps the result, so every agent exposes a trace (timing + model + result) with no bespoke code.run()stays the source of truth for the document.RAGTracebecomes a subclass ofAgentTrace—#39is fully preserved (same fields, same construction, covariant return); a RAG run nests straight into the tree. RAG now builds the answer doc ontooutputandrun()returns it.DoclingOrchestratorAgent.run_task_with_trace()— composes the tree by recording each dispatched sub-agent's trace.run_task()is unchanged and incurs no overhead when tracing is off.LoggingConfig.trace_path+ CLI export; public re-exportsAgentTrace/AgentStep.Why this shape (vs the first revision)
The first revision captured raw LLM I/O via an ambient global recorder hooked into the logging functions. It worked for the CLI but introduced global mutable state and concurrency hazards for a server consumer (Docling Studio), and coupled tracing to logging. This revision is a pure value-object tree — same generic reach, none of those problems — and is literally "the same as #39, generalized". Raw prompt/response capture is deferred and becomes two optional fields on
AgentStepif wanted later.Scope / honesty
stepsfor enricher/writer/editor/extractor are an additive follow-up (RAG already has the richest form).run()/run_task().Tests
tests/test_agent_trace.py: base default, run/run_with_trace equivalence, RAG-as-AgentTrace,outputexcluded from serialization, save round-trip, nested RAGTrace subclass fields preserved, orchestrator tree composition, no-op when off.#39'stest_rag_trace.pyuntouched and green; orchestrator fakes widened to the realrunsignature.uv run pytest→ 116 passeduv run ruff check docling_agent tests→ cleanuv run mypy docling_agent→ clean (30 files)Closes #37