diff --git a/docs/tech-specs/query-time-explainability.md b/docs/tech-specs/query-time-explainability.md index 02598563..751979bc 100644 --- a/docs/tech-specs/query-time-explainability.md +++ b/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: diff --git a/specs/api/components/schemas/rag/GraphRagResponse.yaml b/specs/api/components/schemas/rag/GraphRagResponse.yaml index 5e397003..d92c5f2b 100644 --- a/specs/api/components/schemas/rag/GraphRagResponse.yaml +++ b/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 diff --git a/specs/api/components/schemas/rag/Source.yaml b/specs/api/components/schemas/rag/Source.yaml new file mode 100644 index 00000000..87d285db --- /dev/null +++ b/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 diff --git a/specs/api/paths/flow/graph-rag.yaml b/specs/api/paths/flow/graph-rag.yaml index 7eccd78d..a598f91e 100644 --- a/specs/api/paths/flow/graph-rag.yaml +++ b/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: diff --git a/tests/contract/test_translator_completion_flags.py b/tests/contract/test_translator_completion_flags.py index 606061f9..198638cf 100644 --- a/tests/contract/test_translator_completion_flags.py +++ b/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 diff --git a/tests/integration/test_graph_rag_integration.py b/tests/integration/test_graph_rag_integration.py index 8930d159..162de372 100644 --- a/tests/integration/test_graph_rag_integration.py +++ b/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() diff --git a/tests/integration/test_graph_rag_streaming_integration.py b/tests/integration/test_graph_rag_streaming_integration.py index 8dfd9c2b..edbac03d 100644 --- a/tests/integration/test_graph_rag_streaming_integration.py +++ b/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 diff --git a/tests/unit/test_api/test_rag_chunk_sources.py b/tests/unit/test_api/test_rag_chunk_sources.py new file mode 100644 index 00000000..6aca74e6 --- /dev/null +++ b/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 == [] diff --git a/tests/unit/test_retrieval/test_graph_rag.py b/tests/unit/test_retrieval/test_graph_rag.py index 15ffdc9d..ff56b500 100644 --- a/tests/unit/test_retrieval/test_graph_rag.py +++ b/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 diff --git a/tests/unit/test_retrieval/test_graph_rag_provenance_integration.py b/tests/unit/test_retrieval/test_graph_rag_provenance_integration.py index bc2cb368..51522db7 100644 --- a/tests/unit/test_retrieval/test_graph_rag_provenance_integration.py +++ b/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 == [] diff --git a/tests/unit/test_retrieval/test_graph_rag_service.py b/tests/unit/test_retrieval/test_graph_rag_service.py index a637a350..6f364afc 100644 --- a/tests/unit/test_retrieval/test_graph_rag_service.py +++ b/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"), + ] + diff --git a/trustgraph-base/trustgraph/api/async_socket_client.py b/trustgraph-base/trustgraph/api/async_socket_client.py index 9eff3d60..62d30609 100644 --- a/trustgraph-base/trustgraph/api/async_socket_client.py +++ b/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): diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py index 95fc009b..5ebb0873 100644 --- a/trustgraph-base/trustgraph/api/flow.py +++ b/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( diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 6c51210f..518c3e5a 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/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( diff --git a/trustgraph-base/trustgraph/api/types.py b/trustgraph-base/trustgraph/api/types.py index fafb8224..94d7b1a5 100644 --- a/trustgraph-base/trustgraph/api/types.py +++ b/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: diff --git a/trustgraph-base/trustgraph/messaging/translators/retrieval.py b/trustgraph-base/trustgraph/messaging/translators/retrieval.py index 556ad758..99a265ba 100644 --- a/trustgraph-base/trustgraph/messaging/translators/retrieval.py +++ b/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) diff --git a/trustgraph-base/trustgraph/schema/services/retrieval.py b/trustgraph-base/trustgraph/schema/services/retrieval.py index 47ced73d..bd921fe0 100644 --- a/trustgraph-base/trustgraph/schema/services/retrieval.py +++ b/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 ############################################################################ diff --git a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py index 97bb8db9..cf99905f 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py +++ b/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} " diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py index c094e395..6727ea2b 100644 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py +++ b/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: - if isinstance(result, Exception) or not result: - continue - for triple in result: - p = str(triple.p) - if p in SKIP_PREDICATES: - continue - doc_edges.append(( - str(triple.s), p, str(triple.o) - )) + sources = [] + for uri, result in zip(sorted_doc_uris, metadata_results): + title = "" + if not isinstance(result, Exception) and result: + for triple in result: + p = str(triple.p) + if p in SKIP_PREDICATES: + continue + 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( - selected_edge_uris, - ) - - if isinstance(source_documents, Exception): - logger.warning( - f"Document tracing failed: {source_documents}" + try: + source_documents, sources = await q.trace_source_documents( + selected_edge_uris, ) - 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 diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 9ed802f4..143545f5 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/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} )