Skip to content

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

Open
Nishieee wants to merge 3 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 3 commits into
apache:mainfrom
Nishieee:feat/rag-include-trace

Conversation

@Nishieee

@Nishieee Nishieee commented Jun 1, 2026

Copy link
Copy Markdown
Contributor

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

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>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

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.

Block conflicts:
image

@imbajin imbajin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Reviewed the include_trace changes excluding the current branch conflict. I found one trace correctness issue and one API contract clarity issue.

Comment thread hugegraph-llm/src/hugegraph_llm/flows/common.py
Comment thread hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py Outdated
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

3 participants