feat: structured source document references in graph-rag responses (#1035)

trustgraph-base/trustgraph/api/async_socket_client.py

@@ -267,6 +267,7 @@ class AsyncSocketClient:
                 in_token=resp.get("in_token"),
                 out_token=resp.get("out_token"),
                 model=resp.get("model"),
+                sources=resp.get("sources", []),
             )
 
     async def aclose(self):

trustgraph-base/trustgraph/api/flow.py

@@ -414,6 +414,7 @@ class FlowInstance:
             in_token=result.get("in_token"),
             out_token=result.get("out_token"),
             model=result.get("model"),
+            sources=result.get("sources", []),
         )
 
     def document_rag(

trustgraph-base/trustgraph/api/socket_client.py

@@ -451,6 +451,7 @@ class SocketClient:
                 in_token=resp.get("in_token"),
                 out_token=resp.get("out_token"),
                 model=resp.get("model"),
+                sources=resp.get("sources", []),
             )
 
     def _build_provenance_event(self, resp: Dict[str, Any]) -> ProvenanceEvent:
@@ -715,6 +716,7 @@ class SocketFlowInstance:
                 in_token=result.get("in_token"),
                 out_token=result.get("out_token"),
                 model=result.get("model"),
+                sources=result.get("sources", []),
             )
 
     def graph_rag_explain(

trustgraph-base/trustgraph/api/types.py

@@ -205,6 +205,8 @@ class RAGChunk(StreamingChunk):
         in_token: Input token count (populated on the final chunk, 0 otherwise)
         out_token: Output token count (populated on the final chunk, 0 otherwise)
         model: Model identifier (populated on the final chunk, empty otherwise)
+        sources: Source document references as uri/title dicts (populated
+            on the final chunk, empty otherwise)
         message_type: Always "rag"
     """
     message_type: str = "rag"
@@ -213,6 +215,7 @@ class RAGChunk(StreamingChunk):
     in_token: Optional[int] = None
     out_token: Optional[int] = None
     model: Optional[str] = None
+    sources: List[Dict[str, str]] = dataclasses.field(default_factory=list)
 
 @dataclasses.dataclass
 class TextCompletionResult:
@@ -228,11 +231,14 @@ class TextCompletionResult:
         in_token: Input token count (None if not available)
         out_token: Output token count (None if not available)
         model: Model identifier (None if not available)
+        sources: Source document references as uri/title dicts (graph RAG
+            only, empty otherwise)
     """
     text: Optional[str]
     in_token: Optional[int] = None
     out_token: Optional[int] = None
     model: Optional[str] = None
+    sources: List[Dict[str, str]] = dataclasses.field(default_factory=list)
 
 @dataclasses.dataclass
 class ProvenanceEvent:

trustgraph-base/trustgraph/messaging/translators/retrieval.py

@@ -160,6 +160,13 @@ class GraphRagResponseTranslator(MessageTranslator):
                 self.triple_translator.encode(t) for t in explain_triples
             ]
 
+        # Include source document references (final message only)
+        sources = getattr(obj, "sources", [])
+        if sources:
+            result["sources"] = [
+                {"uri": s.uri, "title": s.title} for s in sources
+            ]
+
         # Include end_of_stream flag (LLM stream complete)
         result["end_of_stream"] = getattr(obj, "end_of_stream", False)
 

trustgraph-base/trustgraph/schema/services/retrieval.py

@@ -19,6 +19,11 @@ class GraphRagQuery:
     streaming: bool = False
     parent_uri: str = ""
 
+@dataclass
+class Source:
+    uri: str = ""     # Source document URI
+    title: str = ""   # Document title (empty when the document has none)
+
 @dataclass
 class GraphRagResponse:
     error: Error | None = None
@@ -32,6 +37,7 @@ class GraphRagResponse:
     in_token: int | None = None
     out_token: int | None = None
     model: str | None = None
+    sources: list[Source] = field(default_factory=list)  # Source documents, on the final message
 
 ############################################################################