diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py index 433438eee..61f570cd6 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py +++ b/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.") diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_response.py b/hugegraph-llm/src/hugegraph_llm/api/models/rag_response.py index fe139eeb8..e37505f48 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_response.py +++ b/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) diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py index 76217251a..64795a14f 100644 --- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py +++ b/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)} diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py index d52d6005e..c2c558273 100644 --- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py +++ b/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", ""), diff --git a/hugegraph-llm/src/hugegraph_llm/flows/common.py b/hugegraph-llm/src/hugegraph_llm/flows/common.py index 8705a3008..51970363b 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/common.py +++ b/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): """ diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py index 66e8fb0be..755584758 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_only.py +++ b/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 diff --git a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py b/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py index 1d738b2ac..d90861038 100644 --- a/hugegraph-llm/src/hugegraph_llm/flows/rag_flow_graph_vector.py +++ b/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 diff --git a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py b/hugegraph-llm/src/hugegraph_llm/state/ai_state.py index 7dce7e857..3fbc67550 100644 --- a/hugegraph-llm/src/hugegraph_llm/state/ai_state.py +++ b/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 diff --git a/hugegraph-llm/src/tests/api/test_rag_api.py b/hugegraph-llm/src/tests/api/test_rag_api.py index 55fbd6794..09ea04784 100644 --- a/hugegraph-llm/src/tests/api/test_rag_api.py +++ b/hugegraph-llm/src/tests/api/test_rag_api.py @@ -20,7 +20,30 @@ from fastapi import APIRouter, FastAPI, status from fastapi.testclient import TestClient -from hugegraph_llm.api.rag_api import rag_http_api +from hugegraph_llm.api.models.rag_requests import RAGRequest +from hugegraph_llm.api.models.rag_response import RAGTrace, serialize_rag_trace +from hugegraph_llm.api.rag_api import build_rag_api_response, rag_http_api +from hugegraph_llm.flows.common import graph_trace_payload +from hugegraph_llm.flows.rag_flow_graph_only import RAGGraphOnlyFlow +from hugegraph_llm.flows.rag_flow_graph_vector import RAGGraphVectorFlow +from hugegraph_llm.state.ai_state import WkFlowInput, WkFlowState + + +def _rag_client(rag_answer_func): + router = APIRouter() + rag_http_api( + router, + rag_answer_func=rag_answer_func, + graph_rag_recall_func=Mock(), + apply_graph_conf=Mock(), + apply_llm_conf=Mock(), + apply_embedding_conf=Mock(), + apply_reranker_conf=Mock(), + gremlin_generate_selective_func=Mock(), + ) + app = FastAPI() + app.include_router(router) + return TestClient(app) def test_graph_config_api_passes_graph_field_to_apply_graph_conf(): @@ -60,3 +83,148 @@ def test_graph_config_api_passes_graph_field_to_apply_graph_conf(): "space_a", origin_call="http", ) + + +def test_rag_default_response_has_no_trace(): + rag_answer = Mock(return_value=("", "", "graph answer", "")) + response = _rag_client(rag_answer).post("/rag", json={"query": "hello", "graph_only": True}) + + assert response.status_code == status.HTTP_200_OK + body = response.json() + assert body == {"query": "hello", "graph_only": "graph answer"} + assert "trace" not in body + + +def test_rag_include_trace_returns_graph_debug_fields(): + rag_answer = Mock( + return_value={ + "graph_only_answer": "graph answer", + "trace": { + "keywords": ["hello"], + "match_vids": ["1:foo"], + "gremlin": "g.V()", + }, + } + ) + response = _rag_client(rag_answer).post( + "/rag", + json={"query": "hello", "graph_only": True, "include_trace": True}, + ) + + assert response.status_code == status.HTTP_200_OK + body = response.json() + assert body["graph_only"] == "graph answer" + assert body["trace"]["keywords"] == ["hello"] + assert body["trace"]["match_vids"] == ["1:foo"] + rag_answer.assert_called_once() + assert rag_answer.call_args.kwargs["include_trace"] is True + + +def test_rag_include_trace_returns_graph_debug_fields_for_graph_vector_answer(): + rag_answer = Mock( + return_value={ + "graph_vector_answer": "hybrid answer", + "trace": { + "keywords": ["hello"], + "match_vids": ["1:foo"], + "gremlin": "g.V()", + }, + } + ) + response = _rag_client(rag_answer).post( + "/rag", + json={ + "query": "hello", + "graph_only": False, + "graph_vector_answer": True, + "include_trace": True, + }, + ) + + assert response.status_code == status.HTTP_200_OK + body = response.json() + assert body["graph_vector_answer"] == "hybrid answer" + assert body["trace"]["keywords"] == ["hello"] + assert body["trace"]["gremlin"] == "g.V()" + + +def test_rag_include_trace_rejects_vector_only(): + rag_answer = Mock(return_value=("", "vector answer", "", "")) + response = _rag_client(rag_answer).post( + "/rag", + json={"query": "hello", "graph_only": False, "vector_only": True, "include_trace": True}, + ) + + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert "include_trace" in response.json()["detail"] + rag_answer.assert_not_called() + + +def test_graph_trace_payload_excludes_prompts_and_secrets(): + trace = graph_trace_payload( + { + "keywords": ["a"], + "gremlin": "g.V()", + "answer_prompt": "secret prompt", + "graph_pwd": "secret", + } + ) + assert trace == {"keywords": ["a"], "gremlin": "g.V()"} + + +def test_serialize_rag_trace_omits_none_fields(): + trace = serialize_rag_trace({"keywords": ["a"], "gremlin": None}) + assert trace == {"keywords": ["a"]} + assert RAGTrace(**trace).keywords == ["a"] + + +def test_build_rag_api_response_supports_legacy_tuple(): + req = RAGRequest(query="hello", graph_only=True) + response = build_rag_api_response(req, ("", "", "graph answer", "")) + assert response == {"query": "hello", "graph_only": "graph answer"} + + +def test_wkflow_state_assign_from_json_preserves_gremlin(): + state = WkFlowState() + state.assign_from_json({"gremlin": "g.V()"}) + assert state.gremlin == "g.V()" + assert graph_trace_payload(state.to_json())["gremlin"] == "g.V()" + + +def test_rag_graph_only_post_deal_include_trace(): + pipeline = Mock() + state = WkFlowState() + state.graph_only_answer = "answer" + state.assign_from_json({"keywords": ["kw"], "gremlin": "g.V()"}) + wk_input = WkFlowInput() + wk_input.include_trace = True + + def get_param(name): + return wk_input if name == "wkflow_input" else state + + pipeline.getGParamWithNoEmpty.side_effect = get_param + + result = RAGGraphOnlyFlow().post_deal(pipeline) + assert result["graph_only_answer"] == "answer" + assert result["trace"]["keywords"] == ["kw"] + assert result["trace"]["gremlin"] == "g.V()" + + +def test_rag_graph_vector_post_deal_include_trace(): + pipeline = Mock() + state = WkFlowState() + state.graph_vector_answer = "hybrid answer" + state.keywords = ["kw"] + state.match_vids = ["1:foo"] + wk_input = WkFlowInput() + wk_input.include_trace = True + + def get_param(name): + return wk_input if name == "wkflow_input" else state + + pipeline.getGParamWithNoEmpty.side_effect = get_param + + result = RAGGraphVectorFlow().post_deal(pipeline) + assert result["graph_vector_answer"] == "hybrid answer" + assert result["trace"]["keywords"] == ["kw"] + assert result["trace"]["match_vids"] == ["1:foo"]