Skip to content

feat(llm): add optional include_trace to /rag for retrieval debugging#354

Open
Nishieee wants to merge 2 commits into
apache:mainfrom
Nishieee:feat/rag-include-trace
Open

feat(llm): add optional include_trace to /rag for retrieval debugging#354
Nishieee wants to merge 2 commits into
apache:mainfrom
Nishieee:feat/rag-include-trace

Conversation

@Nishieee
Copy link
Copy Markdown

@Nishieee Nishieee commented Jun 1, 2026

Summary

Closes #347

  • Adds opt-in include_trace on POST /rag for graph retrieval debugging.
  • Default /rag response is unchanged (no trace unless requested).
  • Vector/rerank trace deferred to follow-up, per the issue scope.

What changed

Area Change
RAGRequest New field include_trace: bool = False
RAGTrace Typed model with 6 graph-recall fields (same names as /rag/graph, minus query)
Flows RAGGraphOnlyFlow / RAGGraphVectorFlow copy allowlisted fields from WkFlowState when trace is enabled
flows/common.py GRAPH_RECALL_TRACE_KEYS (6) for /rag trace; GRAPH_TRACE_KEYS (7, incl. query) for /rag/graph
rag_api.py Builds response via serialize_rag_trace() (exclude_none=True); attaches trace only for graph_only or graph_vector_answer
rag_block.py Passes include_trace through the scheduler (Gradio UI unchanged)

Behavior

  • include_trace=false -> same as before: query + answer field(s).
  • include_trace=true + graph mode -> adds trace with e.g. keywords, match_vids, gremlin, graph_result, etc.
  • vector_only / raw_answer only -> no trace.
  • Secrets and prompts are excluded via the allowlist (not a full state dump).
  • Missing trace fields are omitted; empty lists are kept when recall ran but matched nothing, so clients can distinguish "didn't run" from "ran, no matches."

Test plan

  • Default /rag - no trace, same shape
  • include_trace=true + graph_only - trace returned
  • include_trace=true + graph_vector_answer - trace returned
  • include_trace=true + vector_only only - no trace
  • Allowlist excludes prompts/secrets
  • RAGGraphOnlyFlow / RAGGraphVectorFlow post_deal populate trace
  • serialize_rag_trace omits None fields
cd hugegraph-llm
SKIP_EXTERNAL_SERVICES=true uv run pytest src/tests/api/test_rag_api.py -v --tb=short

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. enhancement New feature or request labels Jun 1, 2026
@github-actions github-actions Bot added the llm label Jun 1, 2026
@imbajin imbajin requested a review from Copilot June 1, 2026 09:46
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 = False to RAGRequest and propagates it through WkFlowInput and 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 /rag response building to optionally attach a serialized trace for 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.

Comment thread hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py Outdated
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request llm size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] Add optional trace output to the RAG answer API for retrieval debugging

2 participants