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

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
-        streaming_text, _ = streaming_response
+        non_streaming_text, _, _ = non_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"),
+        ]
+