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

docs/tech-specs/query-time-explainability.md

@@ -167,6 +167,11 @@ Values are absent (not zero) when token counts are unavailable.
 ## GraphRagResponse Schema
 
 ```python
+@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
@@ -177,6 +182,10 @@ class GraphRagResponse:
     explain_triples: list[Triple] = field(default_factory=list)
     message_type: str = ""  # "chunk" or "explain"
     end_of_session: bool = False
+    in_token: int | None = None
+    out_token: int | None = None
+    model: str | None = None
+    sources: list[Source] = field(default_factory=list)
 ```
 
 ### Message Types
@@ -227,6 +236,15 @@ Selected edges can be traced back to source documents:
 2. Follow `prov:wasDerivedFrom` chain to root document
 3. Each step in chain: chunk → page → document
 
+### Source References in the Response
+
+GraphRAG performs this walk on every query to enrich the synthesis
+prompt with document metadata. The same walk also produces structured
+`sources` entries (`uri` plus `title` from `dc:title`/`rdfs:label`),
+deduplicated and sorted by URI, attached to the final response message
+(`end_of_session=True`) at no additional query cost. Clients can display
+citations without re-running the traversal against the knowledge graph.
+
 ### Cassandra Quoted Triple Support
 
 The Cassandra query service supports matching quoted triples:

specs/api/components/schemas/rag/GraphRagResponse.yaml

@@ -23,6 +23,13 @@ properties:
     description: Provenance triples for this explain event (inline, no follow-up query needed)
     items:
       $ref: '../common/Triple.yaml'
+  sources:
+    type: array
+    description: |
+      Source documents the answer was derived from, deduplicated and sorted
+      by URI. Present on the final message (end_of_session true).
+    items:
+      $ref: './Source.yaml'
   end_of_stream:
     type: boolean
     description: Indicates LLM response stream is complete

specs/api/components/schemas/rag/Source.yaml

@@ -0,0 +1,15 @@
+type: object
+description: |
+  Source document reference. Produced by tracing the graph edges used for
+  retrieval back to the documents they were extracted from.
+properties:
+  uri:
+    type: string
+    description: Source document URI
+    example: urn:document:5a90a175-9906-4dcb-b482-a8c1b6cbf9e0
+  title:
+    type: string
+    description: Document title (empty when the document has none)
+    example: Quantum Mechanics Primer
+required:
+  - uri

specs/api/paths/flow/graph-rag.yaml

@@ -118,6 +118,9 @@ post:
                   - Quantum information theory
                   - Computational complexity theory
                 end-of-stream: false
+                sources:
+                  - uri: urn:document:5a90a175-9906-4dcb-b482-a8c1b6cbf9e0
+                    title: Quantum Mechanics Primer
             streamingChunk:
               summary: Streaming response chunk
               value:

tests/contract/test_translator_completion_flags.py

@@ -8,7 +8,7 @@ based on message fields like end_of_stream and end_of_dialog.
 import pytest
 
 from trustgraph.schema import (
-    GraphRagResponse, DocumentRagResponse, AgentResponse, Error
+    GraphRagResponse, DocumentRagResponse, AgentResponse, Error, Source
 )
 from trustgraph.messaging import TranslatorRegistry
 
@@ -110,6 +110,57 @@ class TestRAGTranslatorCompletionFlags:
         assert response_dict["end_of_stream"] is True
         assert response_dict["end_of_session"] is False
 
+    def test_graph_rag_translator_encodes_sources_on_final_message(self):
+        """
+        Test that GraphRagResponseTranslator encodes source references
+        as uri/title dicts on the final message.
+        """
+        # Arrange
+        translator = TranslatorRegistry.get_response_translator("graph-rag")
+        response = GraphRagResponse(
+            response="A small domesticated mammal.",
+            message_type="chunk",
+            end_of_stream=True,
+            end_of_session=True,
+            error=None,
+            sources=[
+                Source(uri="urn:document:alpha",
+                       title="Quantum Mechanics Primer"),
+                Source(uri="urn:document:beta", title=""),
+            ]
+        )
+
+        # Act
+        response_dict, is_final = translator.encode_with_completion(response)
+
+        # Assert
+        assert is_final is True
+        assert response_dict["sources"] == [
+            {"uri": "urn:document:alpha",
+             "title": "Quantum Mechanics Primer"},
+            {"uri": "urn:document:beta", "title": ""},
+        ]
+
+    def test_graph_rag_translator_omits_empty_sources(self):
+        """
+        Test that the sources key is omitted when there are no sources.
+        """
+        # Arrange
+        translator = TranslatorRegistry.get_response_translator("graph-rag")
+        response = GraphRagResponse(
+            response="Chunk 1",
+            message_type="chunk",
+            end_of_stream=False,
+            end_of_session=False,
+            error=None
+        )
+
+        # Act
+        response_dict, is_final = translator.encode_with_completion(response)
+
+        # Assert
+        assert "sources" not in response_dict
+
     def test_document_rag_translator_is_final_with_end_of_session_true(self):
         """
         Test that DocumentRagResponseTranslator returns is_final=True

tests/integration/test_graph_rag_integration.py

@@ -175,7 +175,7 @@ class TestGraphRagIntegration:
         assert mock_prompt_client.prompt.call_count == 2
 
         # Verify final response
-        response, usage = response
+        response, usage, sources = response
         assert response is not None
         assert isinstance(response, str)
         assert "machine learning" in response.lower()

tests/integration/test_graph_rag_streaming_integration.py

@@ -127,7 +127,7 @@ class TestGraphRagStreaming:
         )
 
         # Assert
-        response, usage = response
+        response, usage, sources = response
         assert_streaming_chunks_valid(collector.chunks, min_chunks=1)
         assert_callback_invoked(AsyncMock(call_count=len(collector.chunks)), min_calls=1)
 
@@ -175,8 +175,8 @@ class TestGraphRagStreaming:
         )
 
         # Assert - Results should be equivalent
-        non_streaming_text, _ = non_streaming_response
+        non_streaming_text, _, _ = non_streaming_response
-        streaming_text, _ = streaming_response
+        streaming_text, _, _ = streaming_response
         assert streaming_text == non_streaming_text
         assert len(streaming_chunks) > 0
         assert "".join(streaming_chunks) == streaming_text
@@ -216,7 +216,7 @@ class TestGraphRagStreaming:
 
         # Assert - Should complete without error
         assert response is not None
-        response_text, usage = response
+        response_text, usage, sources = response
         assert isinstance(response_text, str)
 
     @pytest.mark.asyncio

tests/unit/test_api/test_rag_chunk_sources.py

@@ -0,0 +1,100 @@
+"""
+Tests that the socket clients propagate the sources field from the
+wire format to RAGChunk, and that graph_rag results carry it.
+"""
+
+import pytest
+
+from trustgraph.api.socket_client import SocketClient
+from trustgraph.api.async_socket_client import AsyncSocketClient
+from trustgraph.api.types import RAGChunk
+
+WIRE_SOURCES = [
+    {"uri": "urn:document:alpha", "title": "Quantum Mechanics Primer"},
+    {"uri": "urn:document:beta", "title": ""},
+]
+
+
+@pytest.fixture
+def client():
+    # We only need _parse_chunk — don't connect
+    c = object.__new__(SocketClient)
+    return c
+
+
+@pytest.fixture
+def async_client():
+    c = object.__new__(AsyncSocketClient)
+    return c
+
+
+class TestParseChunkSources:
+
+    def test_final_chunk_carries_sources(self, client):
+        resp = {
+            "message_type": "chunk",
+            "response": "",
+            "end_of_stream": False,
+            "end_of_session": True,
+            "sources": WIRE_SOURCES,
+        }
+        chunk = client._parse_chunk(resp)
+        assert isinstance(chunk, RAGChunk)
+        assert chunk.sources == WIRE_SOURCES
+
+    def test_intermediate_chunk_has_empty_sources(self, client):
+        resp = {
+            "message_type": "chunk",
+            "response": "partial text",
+            "end_of_stream": False,
+        }
+        chunk = client._parse_chunk(resp)
+        assert isinstance(chunk, RAGChunk)
+        assert chunk.sources == []
+
+    def test_async_final_chunk_carries_sources(self, async_client):
+        resp = {
+            "message_type": "chunk",
+            "response": "",
+            "end_of_session": True,
+            "sources": WIRE_SOURCES,
+        }
+        chunk = async_client._parse_chunk(resp)
+        assert isinstance(chunk, RAGChunk)
+        assert chunk.sources == WIRE_SOURCES
+
+    def test_async_intermediate_chunk_has_empty_sources(self, async_client):
+        resp = {
+            "message_type": "chunk",
+            "response": "partial text",
+        }
+        chunk = async_client._parse_chunk(resp)
+        assert isinstance(chunk, RAGChunk)
+        assert chunk.sources == []
+
+
+class TestRestGraphRagSources:
+
+    def test_graph_rag_result_carries_sources(self):
+        from trustgraph.api.flow import FlowInstance
+
+        instance = object.__new__(FlowInstance)
+        instance.request = lambda path, request: {
+            "response": "The answer.",
+            "sources": WIRE_SOURCES,
+        }
+
+        result = instance.graph_rag(query="What is quantum computing?")
+        assert result.text == "The answer."
+        assert result.sources == WIRE_SOURCES
+
+    def test_graph_rag_result_defaults_to_empty_sources(self):
+        from trustgraph.api.flow import FlowInstance
+
+        instance = object.__new__(FlowInstance)
+        instance.request = lambda path, request: {
+            "response": "The answer.",
+        }
+
+        result = instance.graph_rag(query="What is quantum computing?")
+        assert result.sources == []

tests/unit/test_retrieval/test_graph_rag.py

@@ -555,7 +555,7 @@ class TestQuery:
                 explain_callback=collect_provenance,
             )
 
-            response_text, usage = response
+            response_text, usage, sources = response
             assert response_text == expected_response
 
             # 5 events: question, grounding, exploration, focus, synthesis

tests/unit/test_retrieval/test_graph_rag_provenance_integration.py

@@ -14,6 +14,8 @@ from dataclasses import dataclass
 from trustgraph.retrieval.graph_rag.graph_rag import GraphRag, edge_id
 from trustgraph.schema import Triple as SchemaTriple, Term, IRI, LITERAL
 from trustgraph.base import PromptResult
+from trustgraph.base.triples_client import Triple as ClientTriple
+from trustgraph.knowledge import Uri, Literal
 
 from trustgraph.provenance.namespaces import (
     RDF_TYPE, PROV_ENTITY, PROV_WAS_DERIVED_FROM,
@@ -21,6 +23,7 @@ from trustgraph.provenance.namespaces import (
     TG_FOCUS, TG_SYNTHESIS, TG_ANSWER_TYPE,
     TG_QUERY, TG_CONCEPT, TG_ENTITY, TG_EDGE_COUNT,
     TG_SELECTED_EDGE, TG_EDGE, TG_SCORE, TG_EDGE_SELECTION,
+    TG_CONTAINS, DC_TITLE, RDFS_LABEL,
 )
 
 
@@ -423,7 +426,7 @@ class TestGraphRagQueryProvenance:
         async def explain_callback(triples, explain_id):
             events.append({"triples": triples, "explain_id": explain_id})
 
-        result_text, usage = await rag.query(
+        result_text, usage, sources = await rag.query(
             query="What is quantum computing?",
             explain_callback=explain_callback,
 
@@ -460,7 +463,7 @@ class TestGraphRagQueryProvenance:
         clients = build_mock_clients()
         rag = GraphRag(*clients)
 
-        result_text, usage = await rag.query(
+        result_text, usage, sources = await rag.query(
             query="What is quantum computing?",
 
         )
@@ -490,3 +493,165 @@ class TestGraphRagQueryProvenance:
                     f"Triple {t.s.iri} {t.p.iri} should be in "
                     f"urn:graph:retrieval, got {t.g}"
                 )
+
+
+# ---------------------------------------------------------------------------
+# Source document tracing
+# ---------------------------------------------------------------------------
+
+# Provenance chains served by the mock triples client:
+#   EDGE_1, EDGE_2 -> SUBGRAPH_A -> chunk/a -> page/a -> DOC_ALPHA
+#   EDGE_3         -> SUBGRAPH_B -> chunk/b -> page/b -> DOC_BETA + DOC_GAMMA
+SUBGRAPH_A = "http://trustgraph.ai/sg/aaa"
+SUBGRAPH_B = "http://trustgraph.ai/sg/bbb"
+DOC_ALPHA = "urn:document:alpha"
+DOC_BETA = "urn:document:beta"
+DOC_GAMMA = "urn:document:gamma"
+TG_MIME_TYPE = "http://trustgraph.ai/ns/provenance/mimeType"
+
+DERIVATIONS = {
+    SUBGRAPH_A: ["http://trustgraph.ai/chunk/a"],
+    "http://trustgraph.ai/chunk/a": ["http://trustgraph.ai/page/a"],
+    "http://trustgraph.ai/page/a": [DOC_ALPHA],
+    SUBGRAPH_B: ["http://trustgraph.ai/chunk/b"],
+    "http://trustgraph.ai/chunk/b": ["http://trustgraph.ai/page/b"],
+    "http://trustgraph.ai/page/b": [DOC_BETA, DOC_GAMMA],
+}
+
+# alpha has both dc:title and rdfs:label (dc:title preferred), beta has
+# only rdfs:label (fallback), gamma has no title at all (empty string)
+DOC_METADATA = {
+    DOC_ALPHA: [
+        ClientTriple(Uri(DOC_ALPHA), Uri(RDFS_LABEL),
+                     Literal("alpha label")),
+        ClientTriple(Uri(DOC_ALPHA), Uri(DC_TITLE),
+                     Literal("Quantum Mechanics Primer")),
+        ClientTriple(Uri(DOC_ALPHA), Uri(TG_MIME_TYPE),
+                     Literal("application/pdf")),
+    ],
+    DOC_BETA: [
+        ClientTriple(Uri(DOC_BETA), Uri(RDFS_LABEL),
+                     Literal("Physics Notes")),
+    ],
+    DOC_GAMMA: [
+        ClientTriple(Uri(DOC_GAMMA), Uri(TG_MIME_TYPE),
+                     Literal("text/plain")),
+    ],
+}
+
+EXPECTED_SOURCES = [
+    {"uri": DOC_ALPHA, "title": "Quantum Mechanics Primer"},
+    {"uri": DOC_BETA, "title": "Physics Notes"},
+    {"uri": DOC_GAMMA, "title": ""},
+]
+
+# Total triples_client.query calls query() makes against the graph above:
+# 6 label lookups + 3 tg:contains + 9 wasDerivedFrom + 3 doc metadata.
+# Sources are built from the same fetches, so this total must not grow.
+EXPECTED_TRIPLES_QUERY_CALLS = 21
+
+
+def build_source_tracing_clients(fail_tracing=False):
+    """Like build_mock_clients, but the triples client also serves the
+    tg:contains + prov:wasDerivedFrom chains and document metadata."""
+    (prompt_client, embeddings_client, graph_embeddings_client,
+     triples_client, reranker_client) = build_mock_clients()
+
+    def subgraph_for(quoted):
+        t = quoted.triple
+        if t.p.iri == "http://schema.org/relatedTo":
+            return SUBGRAPH_A
+        return SUBGRAPH_A if t.s.iri == ENTITY_A else SUBGRAPH_B
+
+    async def mock_query(s=None, p=None, o=None, limit=1,
+                         user=None, collection=None, g=None):
+        if p == TG_CONTAINS and o is not None:
+            if fail_tracing:
+                raise RuntimeError("triple store unavailable")
+            sg = subgraph_for(o)
+            return [ClientTriple(Uri(sg), Uri(TG_CONTAINS), o)]
+        if p == PROV_WAS_DERIVED_FROM:
+            return [
+                ClientTriple(Uri(str(s)), Uri(PROV_WAS_DERIVED_FROM),
+                             Uri(target))
+                for target in DERIVATIONS.get(str(s), [])
+            ]
+        if p is None and str(s) in DOC_METADATA:
+            return DOC_METADATA[str(s)]
+        return []  # Label lookups: fall back to URI
+
+    triples_client.query.side_effect = mock_query
+
+    return (prompt_client, embeddings_client, graph_embeddings_client,
+            triples_client, reranker_client)
+
+
+class TestGraphRagSourceTracing:
+    """query() should return structured source references built from the
+    provenance walk it already performs."""
+
+    @pytest.mark.asyncio
+    async def test_query_returns_sources(self):
+        """Sources are deduplicated, uri-sorted, titled where possible."""
+        clients = build_source_tracing_clients()
+        rag = GraphRag(*clients)
+
+        resp, usage, sources = await rag.query(
+            query="What is quantum computing?",
+        )
+
+        assert resp == (
+            "Quantum computing applies physics principles to computation."
+        )
+        assert sources == EXPECTED_SOURCES
+
+    @pytest.mark.asyncio
+    async def test_sources_add_zero_triple_queries(self):
+        """Building sources must not add any triple-store queries."""
+        clients = build_source_tracing_clients()
+        triples_client = clients[3]
+        rag = GraphRag(*clients)
+
+        resp, usage, sources = await rag.query(
+            query="What is quantum computing?",
+        )
+
+        assert sources == EXPECTED_SOURCES
+        assert triples_client.query.call_count == (
+            EXPECTED_TRIPLES_QUERY_CALLS
+        )
+
+    @pytest.mark.asyncio
+    async def test_doc_metadata_still_reaches_synthesis_prompt(self):
+        """The kg-synthesis prompt context keeps the document edges."""
+        clients = build_source_tracing_clients()
+        prompt_client = clients[0]
+        rag = GraphRag(*clients)
+
+        await rag.query(query="What is quantum computing?")
+
+        synthesis_calls = [
+            c for c in prompt_client.prompt.call_args_list
+            if c.args[0] == "kg-synthesis"
+        ]
+        assert len(synthesis_calls) == 1
+        knowledge = synthesis_calls[0].kwargs["variables"]["knowledge"]
+        assert {
+            "s": DOC_ALPHA, "p": DC_TITLE,
+            "o": "Quantum Mechanics Primer",
+        } in knowledge
+
+    @pytest.mark.asyncio
+    async def test_tracing_failure_degrades_to_empty_sources(self):
+        """A failing walk yields empty sources, answer unaffected."""
+        clients = build_source_tracing_clients(fail_tracing=True)
+        rag = GraphRag(*clients)
+
+        resp, usage, sources = await rag.query(
+            query="What is quantum computing?",
+        )
+
+        assert resp == (
+            "Quantum computing applies physics principles to computation."
+        )
+        assert sources == []

tests/unit/test_retrieval/test_graph_rag_service.py

@@ -8,7 +8,7 @@ import pytest
 from unittest.mock import MagicMock, AsyncMock, patch
 
 from trustgraph.retrieval.graph_rag.rag import Processor
-from trustgraph.schema import GraphRagQuery, GraphRagResponse
+from trustgraph.schema import GraphRagQuery, GraphRagResponse, Source
 
 
 class TestGraphRagService:
@@ -44,7 +44,7 @@ class TestGraphRagService:
                 await explain_callback([], "urn:trustgraph:prov:retrieval:test")
                 await explain_callback([], "urn:trustgraph:prov:selection:test")
                 await explain_callback([], "urn:trustgraph:prov:answer:test")
-            return "A small domesticated mammal.", {"in_token": None, "out_token": None, "model": None}
+            return "A small domesticated mammal.", {"in_token": None, "out_token": None, "model": None}, []
 
         mock_rag_instance.query.side_effect = mock_query
 
@@ -93,6 +93,7 @@ class TestGraphRagService:
         assert chunk_msg.response == "A small domesticated mammal."
         assert chunk_msg.end_of_stream is True
         assert chunk_msg.end_of_session is True
+        assert chunk_msg.sources == []
 
         # Verify provenance triples were sent to provenance queue
         assert mock_provenance_producer.send.call_count == 4
@@ -180,7 +181,7 @@ class TestGraphRagService:
 
         async def mock_query(**kwargs):
             # Don't call explain_callback
-            return "Response text", {"in_token": None, "out_token": None, "model": None}
+            return "Response text", {"in_token": None, "out_token": None, "model": None}, []
 
         mock_rag_instance.query.side_effect = mock_query
 
@@ -219,3 +220,112 @@ class TestGraphRagService:
         assert chunk_msg.response == "Response text"
         assert chunk_msg.end_of_stream is True
         assert chunk_msg.end_of_session is True
+
+    @patch('trustgraph.retrieval.graph_rag.rag.GraphRag')
+    @pytest.mark.asyncio
+    async def test_non_streaming_final_message_carries_sources(
+            self, mock_graph_rag_class):
+        """
+        Test that the non-streaming response carries the source references
+        returned by the query.
+        """
+        processor = Processor(
+            taskgroup=MagicMock(),
+            id="test-processor",
+        )
+
+        mock_rag_instance = AsyncMock()
+        mock_graph_rag_class.return_value = mock_rag_instance
+
+        async def mock_query(**kwargs):
+            return "Answer.", \
+                {"in_token": None, "out_token": None, "model": None}, \
+                [
+                    {"uri": "urn:document:alpha",
+                     "title": "Quantum Mechanics Primer"},
+                    {"uri": "urn:document:beta", "title": ""},
+                ]
+
+        mock_rag_instance.query.side_effect = mock_query
+
+        msg = MagicMock()
+        msg.value.return_value = GraphRagQuery(
+            query="Test query",
+            collection="default",
+            streaming=False
+        )
+        msg.properties.return_value = {"id": "test-id"}
+
+        consumer = MagicMock()
+        flow = MagicMock()
+
+        mock_response_producer = AsyncMock()
+        flow.side_effect = lambda service_name: mock_response_producer
+
+        # Execute
+        await processor.on_request(msg, consumer, flow)
+
+        # Final (only) message carries the sources
+        chunk_msg = mock_response_producer.send.call_args_list[0][0][0]
+        assert chunk_msg.end_of_session is True
+        assert chunk_msg.sources == [
+            Source(uri="urn:document:alpha",
+                   title="Quantum Mechanics Primer"),
+            Source(uri="urn:document:beta", title=""),
+        ]
+
+    @patch('trustgraph.retrieval.graph_rag.rag.GraphRag')
+    @pytest.mark.asyncio
+    async def test_streaming_final_message_carries_sources(
+            self, mock_graph_rag_class):
+        """
+        Test that in streaming mode only the final end_of_session message
+        carries the source references.
+        """
+        processor = Processor(
+            taskgroup=MagicMock(),
+            id="test-processor",
+        )
+
+        mock_rag_instance = AsyncMock()
+        mock_graph_rag_class.return_value = mock_rag_instance
+
+        async def mock_query(**kwargs):
+            chunk_callback = kwargs.get('chunk_callback')
+            await chunk_callback("Streamed answer.", True)
+            return "Streamed answer.", \
+                {"in_token": None, "out_token": None, "model": None}, \
+                [{"uri": "urn:document:alpha", "title": "Primer"}]
+
+        mock_rag_instance.query.side_effect = mock_query
+
+        msg = MagicMock()
+        msg.value.return_value = GraphRagQuery(
+            query="Test query",
+            collection="default",
+            streaming=True
+        )
+        msg.properties.return_value = {"id": "test-id"}
+
+        consumer = MagicMock()
+        flow = MagicMock()
+
+        mock_response_producer = AsyncMock()
+        flow.side_effect = lambda service_name: mock_response_producer
+
+        # Execute
+        await processor.on_request(msg, consumer, flow)
+
+        # 2 messages: streamed chunk, then end_of_session close
+        assert mock_response_producer.send.call_count == 2
+
+        chunk_msg = mock_response_producer.send.call_args_list[0][0][0]
+        assert chunk_msg.end_of_session is False
+        assert chunk_msg.sources == []
+
+        final_msg = mock_response_producer.send.call_args_list[1][0][0]
+        assert final_msg.end_of_session is True
+        assert final_msg.sources == [
+            Source(uri="urn:document:alpha", title="Primer"),
+        ]
+

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
 
 ############################################################################
 

trustgraph-cli/trustgraph/cli/invoke_graph_rag.py

@@ -29,6 +29,20 @@ default_edge_score_limit = 30
 default_edge_limit = 25
 default_max_reranker_input = 350
 
+def _print_sources(sources):
+    """Print the source document references from the final message."""
+    if not sources:
+        return
+    print("Sources:", file=sys.stderr)
+    for src in sources:
+        uri = src.get("uri", "")
+        title = src.get("title", "")
+        if title:
+            print(f"  - {title} ({uri})", file=sys.stderr)
+        else:
+            print(f"  - {uri}", file=sys.stderr)
+
+
 def _question_explainable_api(
         url, flow_id, question_text, collection, entity_limit, triple_limit,
         max_subgraph_size, max_path_length, edge_score_limit=30,
@@ -42,6 +56,8 @@ def _question_explainable_api(
     explain_client = ExplainabilityClient(flow, retry_delay=0.2, max_retries=10)
 
     try:
+        final_sources = []
+
         # Stream GraphRAG with explainability - process events as they arrive
         for item in flow.graph_rag_explain(
             query=question_text,
@@ -57,6 +73,8 @@ def _question_explainable_api(
             if isinstance(item, RAGChunk):
                 # Print response content
                 print(item.content, end="", flush=True)
+                if item.sources:
+                    final_sources = item.sources
 
             elif isinstance(item, ProvenanceEvent):
                 # Use inline entity if available, otherwise fetch from graph
@@ -134,6 +152,8 @@ def _question_explainable_api(
 
         print()  # Final newline
 
+        _print_sources(final_sources)
+
     finally:
         socket.close()
 
@@ -195,6 +215,9 @@ def question(
                 last_chunk = chunk
             print()  # Final newline
 
+            if last_chunk:
+                _print_sources(last_chunk.sources)
+
             if show_usage and last_chunk:
                 print(
                     f"Input tokens: {last_chunk.in_token}  "
@@ -221,6 +244,8 @@ def question(
         )
         print(result.text)
 
+        _print_sources(result.sources)
+
         if show_usage:
             print(
                 f"Input tokens: {result.in_token}  "

trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py

@@ -27,6 +27,7 @@ from trustgraph.provenance import (
     set_graph,
     GRAPH_RETRIEVAL, GRAPH_SOURCE,
     TG_CONTAINS, PROV_WAS_DERIVED_FROM,
+    DC_TITLE, RDFS_LABEL,
 )
 
 # Module logger
@@ -500,7 +501,9 @@ class Query:
             edge_uris: List of (s, p, o) URI string tuples
 
         Returns:
-            List of unique document titles
+            (doc_edges, sources): document metadata edges as (s, p, o)
+            string tuples, and per-document source references as
+            {"uri", "title"} dicts sorted by uri
         """
         # Step 1: Find subgraphs containing these edges via tg:contains
         subgraph_tasks = []
@@ -535,7 +538,7 @@ class Query:
                 subgraph_uris.add(str(triple.s))
 
         if not subgraph_uris:
-            return []
+            return [], []
 
         # Step 2: Walk prov:wasDerivedFrom chain to find documents
         current_uris = subgraph_uris
@@ -569,7 +572,7 @@ class Query:
             current_uris = next_uris - doc_uris
 
         if not doc_uris:
-            return []
+            return [], []
 
         # Step 3: Get all document metadata properties
         SKIP_PREDICATES = {
@@ -577,12 +580,14 @@ class Query:
             "http://www.w3.org/1999/02/22-rdf-syntax-ns#type",
         }
 
+        sorted_doc_uris = sorted(doc_uris)
+
         metadata_tasks = [
             self.rag.triples_client.query(
                 s=uri, p=None, o=None, limit=50,
                 collection=self.collection,
             )
-            for uri in doc_uris
+            for uri in sorted_doc_uris
         ]
 
         metadata_results = await asyncio.gather(
@@ -590,18 +595,22 @@ class Query:
         )
 
         doc_edges = []
-        for result in metadata_results:
+        sources = []
-            if isinstance(result, Exception) or not result:
+        for uri, result in zip(sorted_doc_uris, metadata_results):
-                continue
+            title = ""
-            for triple in result:
+            if not isinstance(result, Exception) and result:
-                p = str(triple.p)
+                for triple in result:
-                if p in SKIP_PREDICATES:
+                    p = str(triple.p)
-                    continue
+                    if p in SKIP_PREDICATES:
-                doc_edges.append((
+                        continue
-                    str(triple.s), p, str(triple.o)
+                    doc_edges.append((
-                ))
+                        str(triple.s), p, str(triple.o)
+                    ))
+                    if p == DC_TITLE or (p == RDFS_LABEL and not title):
+                        title = str(triple.o)
+            sources.append({"uri": uri, "title": title})
 
-        return doc_edges
+        return doc_edges, sources
 
 class GraphRag:
     """
@@ -740,15 +749,13 @@ class GraphRag:
             if edge_id(s, p, o) in uri_map
         ]
 
-        source_documents = await q.trace_source_documents(
+        try:
-            selected_edge_uris,
+            source_documents, sources = await q.trace_source_documents(
-        )
+                selected_edge_uris,
-
-        if isinstance(source_documents, Exception):
-            logger.warning(
-                f"Document tracing failed: {source_documents}"
             )
-            source_documents = []
+        except Exception as e:
+            logger.warning(f"Document tracing failed: {e}")
+            source_documents, sources = [], []
 
         # Build focus explainability data with cross-encoder metadata
         selected_edges_with_reasoning = []
@@ -866,4 +873,4 @@ class GraphRag:
             "model": last_model,
         }
 
-        return resp, usage
+        return resp, usage, sources

trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py

@@ -6,7 +6,7 @@ Input is query, output is response.
 
 import logging
 
-from ... schema import GraphRagQuery, GraphRagResponse, Error
+from ... schema import GraphRagQuery, GraphRagResponse, Error, Source
 from ... schema import Triples, Metadata
 from ... provenance import GRAPH_RETRIEVAL
 from . graph_rag import GraphRag
@@ -22,6 +22,9 @@ logger = logging.getLogger(__name__)
 default_ident = "graph-rag"
 default_concurrency = 1
 
+def source_refs(sources):
+    return [Source(uri=s["uri"], title=s["title"]) for s in sources]
+
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
@@ -229,7 +232,7 @@ class Processor(FlowProcessor):
                     )
 
                 # Query with streaming and real-time explain
-                response, usage = await rag.query(
+                response, usage, sources = await rag.query(
                     query = v.query, collection = v.collection,
                     entity_limit = entity_limit, triple_limit = triple_limit,
                     max_subgraph_size = max_subgraph_size,
@@ -245,7 +248,7 @@ class Processor(FlowProcessor):
 
             else:
                 # Non-streaming path with real-time explain
-                response, usage = await rag.query(
+                response, usage, sources = await rag.query(
                     query = v.query, collection = v.collection,
                     entity_limit = entity_limit, triple_limit = triple_limit,
                     max_subgraph_size = max_subgraph_size,
@@ -267,6 +270,7 @@ class Processor(FlowProcessor):
                         in_token=usage.get("in_token"),
                         out_token=usage.get("out_token"),
                         model=usage.get("model"),
+                        sources=source_refs(sources),
                     ),
                     properties={"id": id}
                 )
@@ -281,6 +285,7 @@ class Processor(FlowProcessor):
                     in_token=usage.get("in_token"),
                     out_token=usage.get("out_token"),
                     model=usage.get("model"),
+                    sources=source_refs(sources),
                 ),
                 properties={"id": id}
             )