feat(llm): add optional include_trace to /rag for retrieval debugging#354
Open
Nishieee wants to merge 2 commits into
Open
feat(llm): add optional include_trace to /rag for retrieval debugging#354Nishieee wants to merge 2 commits into
Nishieee wants to merge 2 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds an opt-in include_trace flag to POST /rag to return a structured graph-retrieval debug trace (while keeping the default response shape unchanged), supporting retrieval debugging without leaking prompts/secrets via a full state dump.
Changes:
- Adds
include_trace: bool = FalsetoRAGRequestand propagates it throughWkFlowInputand RAG flows. - Introduces
RAGTrace+serialize_rag_trace()and allowlisted trace key sets (GRAPH_RECALL_TRACE_KEYS/GRAPH_TRACE_KEYS) to control what can be returned. - Updates
/ragresponse building to optionally attach a serializedtracefor graph modes, with tests covering default behavior, trace inclusion, and allowlisting.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| hugegraph-llm/src/tests/api/test_rag_api.py | Adds API + flow unit tests validating include_trace behavior and trace allowlisting/serialization. |
| hugegraph-llm/src/hugegraph_llm/state/ai_state.py | Adds include_trace to WkFlowInput and resets it in reset(). |
| hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py | Passes include_trace into the flow input and attaches allowlisted trace in post_deal(). |
| hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py | Passes include_trace into the flow input and attaches allowlisted trace in post_deal() (non-recall path). |
| hugegraph-llm/src/hugegraph_llm/flows/common.py | Centralizes answer payload + trace allowlist keys and helper payload builders. |
| hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py | Threads include_trace through scheduler calls and optionally returns the raw dict result when enabled. |
| hugegraph-llm/src/hugegraph_llm/api/rag_api.py | Adds build_rag_api_response() to preserve legacy tuple support and conditionally attaches serialized trace. |
| hugegraph-llm/src/hugegraph_llm/api/models/rag_response.py | Adds RAGTrace model and serialize_rag_trace() helper (exclude None). |
| hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py | Adds include_trace request field with default False. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
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
Closes #347
include_traceonPOST /ragfor graph retrieval debugging./ragresponse is unchanged (no trace unless requested).What changed
RAGRequestinclude_trace: bool = FalseRAGTrace/rag/graph, minusquery)RAGGraphOnlyFlow/RAGGraphVectorFlowcopy allowlisted fields fromWkFlowStatewhen trace is enabledflows/common.pyGRAPH_RECALL_TRACE_KEYS(6) for/ragtrace;GRAPH_TRACE_KEYS(7, incl.query) for/rag/graphrag_api.pyserialize_rag_trace()(exclude_none=True); attachestraceonly forgraph_onlyorgraph_vector_answerrag_block.pyinclude_tracethrough the scheduler (Gradio UI unchanged)Behavior
include_trace=false-> same as before:query+ answer field(s).include_trace=true+ graph mode -> addstracewith e.g.keywords,match_vids,gremlin,graph_result, etc.vector_only/raw_answeronly -> no trace.Test plan
/rag- no trace, same shapeinclude_trace=true+graph_only- trace returnedinclude_trace=true+graph_vector_answer- trace returnedinclude_trace=true+vector_onlyonly - no traceRAGGraphOnlyFlow/RAGGraphVectorFlowpost_dealpopulate traceserialize_rag_traceomitsNonefieldscd hugegraph-llm SKIP_EXTERNAL_SERVICES=true uv run pytest src/tests/api/test_rag_api.py -v --tb=short