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

hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py

@@ -56,6 +56,10 @@ class RAGRequest(BaseModel):
                                    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 graph retrieval trace/debug info (graph_only or graph_vector_answer only).",
+    )
 
     # Keep prompt params in the end
     answer_prompt: Optional[str] = Query(prompt.answer_prompt, description="Prompt to guide the answer generation.")

hugegraph-llm/src/hugegraph_llm/api/models/rag_response.py

@@ -15,9 +15,24 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from typing import Any, Dict, List, Optional
+
 from pydantic import BaseModel
 
 
 class RAGResponse(BaseModel):
     status_code: int = -1
     message: str = ""
+
+
+class RAGTrace(BaseModel):
+    keywords: Optional[List[Any]] = None
+    match_vids: Optional[List[Any]] = None
+    graph_result_flag: Optional[int] = None
+    gremlin: Optional[str] = None
+    graph_result: Optional[List[Any]] = None
+    vertex_degree_list: Optional[List[Any]] = None
+
+
+def serialize_rag_trace(data: Dict[str, Any]) -> Dict[str, Any]:
+    return RAGTrace(**data).model_dump(exclude_none=True)

hugegraph-llm/src/hugegraph_llm/api/rag_api.py

@@ -28,11 +28,38 @@
     RAGRequest,
     RerankerConfigRequest,
 )
-from hugegraph_llm.api.models.rag_response import RAGResponse
+from hugegraph_llm.api.models.rag_response import RAGResponse, serialize_rag_trace
 from hugegraph_llm.config import huge_settings, llm_settings, prompt
+from hugegraph_llm.flows.common import GRAPH_TRACE_KEYS
 from hugegraph_llm.utils.graph_index_utils import get_vertex_details
 from hugegraph_llm.utils.log import log
 
+RAG_ANSWER_FIELDS = (
+    ("raw_answer", "raw_answer"),
+    ("vector_only", "vector_only_answer"),
+    ("graph_only", "graph_only_answer"),
+    ("graph_vector_answer", "graph_vector_answer"),
+)
+
+
+def build_rag_api_response(req: RAGRequest, result):
+    response = {"query": req.query}
+    if isinstance(result, dict):
+        for req_key, res_key in RAG_ANSWER_FIELDS:
+            if getattr(req, req_key):
+                response[req_key] = result.get(res_key, "")
+    else:
+        for (req_key, _), value in zip(RAG_ANSWER_FIELDS, result):
+            if getattr(req, req_key):
+                response[req_key] = value
+    if req.include_trace and (req.graph_only or req.graph_vector_answer):
+        raw_trace = result.get("trace") if isinstance(result, dict) else None
+        if raw_trace:
+            trace = serialize_rag_trace(raw_trace)
+            if trace:
+                response["trace"] = trace
+    return response
+
 
 # pylint: disable=too-many-statements
 def rag_http_api(
@@ -55,6 +82,11 @@ def rag_answer_api(req: RAGRequest):
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail="Query must not be empty.",
             )
+        if req.include_trace and not (req.graph_only or req.graph_vector_answer):
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="include_trace is only supported for graph_only or graph_vector_answer.",
+            )
 
         result = rag_answer_func(
             text=req.query,
@@ -75,19 +107,9 @@ def rag_answer_api(req: RAGRequest):
             answer_prompt=req.answer_prompt or prompt.answer_prompt,
             keywords_extract_prompt=req.keywords_extract_prompt or prompt.keywords_extract_prompt,
             gremlin_prompt=req.gremlin_prompt or prompt.gremlin_generate_prompt,
+            include_trace=req.include_trace,
         )
-        # TODO: we need more info in the response for users to understand the query logic
-        return {
-            "query": req.query,
-            **{
-                key: value
-                for key, value in zip(
-                    ["raw_answer", "vector_only", "graph_only", "graph_vector_answer"],
-                    result,
-                )
-                if getattr(req, key)
-            },
-        }
+        return build_rag_api_response(req, result)
 
     def set_graph_config(req):
         if req.client_config:
@@ -129,16 +151,7 @@ def graph_rag_recall_api(req: GraphRAGRequest):
                     result["match_vids"] = vertex_details
 
             if isinstance(result, dict):
-                params = [
-                    "query",
-                    "keywords",
-                    "match_vids",
-                    "graph_result_flag",
-                    "gremlin",
-                    "graph_result",
-                    "vertex_degree_list",
-                ]
-                user_result = {key: result[key] for key in params if key in result}
+                user_result = {key: result[key] for key in GRAPH_TRACE_KEYS if key in result}
                 return {"graph_recall": user_result}
             return {"graph_recall": json.dumps(result)}
 

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py

@@ -49,7 +49,8 @@ def rag_answer(
     topk_return_results=20,
     vector_dis_threshold=0.9,
     topk_per_keyword=1,
-) -> Tuple:
+    include_trace: bool = False,
+) -> tuple[str, str, str, str] | dict[str, object]:
     """
     Generate an answer using the RAG (Retrieval-Augmented Generation) pipeline.
     Fetch the Scheduler to deal with the request
@@ -103,9 +104,12 @@ def rag_answer(
             topk_return_results=topk_return_results,
             vector_dis_threshold=vector_dis_threshold,
             topk_per_keyword=topk_per_keyword,
+            include_trace=include_trace,
         )
         if res.get("switch_to_bleu"):
             gr.Warning("Online reranker fails, automatically switches to local bleu rerank.")
+        if include_trace:
+            return res
         return (
             res.get("raw_answer", ""),
             res.get("vector_only_answer", ""),

hugegraph-llm/src/hugegraph_llm/flows/common.py

@@ -19,6 +19,30 @@
 from hugegraph_llm.state.ai_state import WkFlowInput
 from hugegraph_llm.utils.log import log
 
+GRAPH_RECALL_TRACE_KEYS = (
+    "keywords",
+    "match_vids",
+    "graph_result_flag",
+    "gremlin",
+    "graph_result",
+    "vertex_degree_list",
+)
+
+GRAPH_TRACE_KEYS = ("query",) + GRAPH_RECALL_TRACE_KEYS
+
+
+def rag_answer_payload(res: Dict[str, Any]) -> Dict[str, str]:
+    return {
+        "raw_answer": res.get("raw_answer", ""),
+        "vector_only_answer": res.get("vector_only_answer", ""),
+        "graph_only_answer": res.get("graph_only_answer", ""),
+        "graph_vector_answer": res.get("graph_vector_answer", ""),
+    }
+
+
+def graph_trace_payload(res: Dict[str, Any]) -> Dict[str, Any]:
+    return {key: res[key] for key in GRAPH_RECALL_TRACE_KEYS if key in res}
+
 
 class BaseFlow(ABC):
     """

hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py

@@ -19,7 +19,7 @@
 from pycgraph import GCondition, GPipeline, GRegion
 
 from hugegraph_llm.config import huge_settings, prompt
-from hugegraph_llm.flows.common import BaseFlow
+from hugegraph_llm.flows.common import BaseFlow, graph_trace_payload, rag_answer_payload
 from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode
 from hugegraph_llm.nodes.hugegraph_node.graph_query_node import GraphQueryNode
 from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode
@@ -95,6 +95,7 @@ def prepare(
 
         prepared_input.is_graph_rag_recall = is_graph_rag_recall
         prepared_input.is_vector_only = is_vector_only
+        prepared_input.include_trace = kwargs.get("include_trace", False)
         prepared_input.data_json = {
             "query": query,
             "vector_search": vector_search,
@@ -141,13 +142,9 @@ def post_deal(self, pipeline=None, **kwargs):
             return {"error": "No pipeline provided"}
         res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json()
         log.info("RAGGraphOnlyFlow post processing success")
-        return (
-            {
-                "raw_answer": res.get("raw_answer", ""),
-                "vector_only_answer": res.get("vector_only_answer", ""),
-                "graph_only_answer": res.get("graph_only_answer", ""),
-                "graph_vector_answer": res.get("graph_vector_answer", ""),
-            }
-            if not res.get("is_graph_rag_recall", False)
-            else res
-        )
+        if res.get("is_graph_rag_recall", False):
+            return res
+        out = rag_answer_payload(res)
+        if pipeline.getGParamWithNoEmpty("wkflow_input").include_trace:
+            out["trace"] = graph_trace_payload(res)
+        return out

hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py

@@ -19,7 +19,7 @@
 from pycgraph import GPipeline
 
 from hugegraph_llm.config import huge_settings, prompt
-from hugegraph_llm.flows.common import BaseFlow
+from hugegraph_llm.flows.common import BaseFlow, graph_trace_payload, rag_answer_payload
 from hugegraph_llm.nodes.common_node.merge_rerank_node import MergeRerankNode
 from hugegraph_llm.nodes.hugegraph_node.graph_query_node import GraphQueryNode
 from hugegraph_llm.nodes.hugegraph_node.schema import SchemaNode
@@ -81,6 +81,7 @@ def prepare(
         prepared_input.answer_prompt = answer_prompt or prompt.answer_prompt
         prepared_input.custom_related_information = custom_related_information
         prepared_input.schema = huge_settings.graph_name
+        prepared_input.include_trace = kwargs.get("include_trace", False)
 
         prepared_input.data_json = {
             "query": query,
@@ -121,9 +122,7 @@ def post_deal(self, pipeline=None, **kwargs):
             return {"error": "No pipeline provided"}
         res = pipeline.getGParamWithNoEmpty("wkflow_state").to_json()
         log.info("RAGGraphVectorFlow post processing success")
-        return {
-            "raw_answer": res.get("raw_answer", ""),
-            "vector_only_answer": res.get("vector_only_answer", ""),
-            "graph_only_answer": res.get("graph_only_answer", ""),
-            "graph_vector_answer": res.get("graph_vector_answer", ""),
-        }
+        out = rag_answer_payload(res)
+        if pipeline.getGParamWithNoEmpty("wkflow_input").include_trace:
+            out["trace"] = graph_trace_payload(res)
+        return out

hugegraph-llm/src/hugegraph_llm/state/ai_state.py

@@ -76,6 +76,7 @@ class WkFlowInput(GParam):
     # used for rag_recall api
     is_graph_rag_recall: bool = False
     is_vector_only: bool = False
+    include_trace: bool = False
 
     # used for build text2gremin index
     examples: Optional[List[Dict[str, str]]] = None
@@ -132,6 +133,7 @@ def reset(self, _: CStatus) -> None:
         self.examples = None
         self.is_graph_rag_recall = False
         self.is_vector_only = False
+        self.include_trace = False
 
 
 class WkFlowState(GParam):
@@ -185,6 +187,7 @@ class WkFlowState(GParam):
     stream_generator: Optional[AsyncGenerator] = None
 
     graph_result_flag: Optional[int] = None
+    gremlin: Optional[str] = None
     vertex_degree_list: Optional[List] = None
     knowledge_with_degree: Optional[Dict] = None
     graph_context_head: Optional[str] = None
@@ -243,6 +246,7 @@ def setup(self) -> CStatus:
 
         self.stream_generator = None
         self.graph_result_flag = None
+        self.gremlin = None
         self.vertex_degree_list = None
         self.knowledge_with_degree = None
         self.graph_context_head = None