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

@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 <[email protected]>

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.

"keywords",
"match_vids",
"graph_result_flag",
"gremlin",

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.

‼️ gremlin is advertised as part of the new trace contract, but the real flow drops it before graph_trace_payload() can see it. GraphQueryNode writes context["gremlin"], then BaseNode persists node output through WkFlowState.assign_from_json(), which only keeps fields declared on WkFlowState. Since WkFlowState declares nearby graph trace fields such as graph_result_flag and vertex_degree_list but not gremlin, real /rag responses with include_trace=true will omit trace.gremlin even when Gremlin was generated.

The current test masks this by dynamically assigning state.gremlin = "g.V()", bypassing assign_from_json(). Please either add gremlin: Optional[str] to WkFlowState and reset it in setup(), or remove gremlin from the public trace contract. The regression test should exercise assign_from_json() or an actual flow-state path.

extracted from the query, by default only the most similar one is returned.",
)
client_config: Optional[GraphConfigRequest] = Query(None, description="hugegraph server config.")
include_trace: bool = Query(False, description="Include retrieval trace/debug info in the response.")

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.

⚠️ include_trace is exposed as a generic retrieval/debug switch, but the implementation only returns trace for graph_only / graph_vector_answer; vector_only + include_trace silently omits trace. That makes it hard for clients to distinguish "trace unsupported for this mode" from "no trace was produced", and vector-only is also a retrieval mode.

Please either narrow the request contract/description to graph-backed trace and reject unsupported combinations with a clear 400, or implement vector trace data as part of the same contract.

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