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

This commit is contained in:
Sunny Yang 2026-07-08 16:59:56 -06:00 committed by GitHub
parent 0374098ee9
commit e9c6a850ad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 568 additions and 39 deletions

View file

@ -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:

View file

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

View file

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

View file

@ -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:

View file

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

View file

@ -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()

View file

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

View file

@ -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 == []

View file

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

View file

@ -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 == []

View file

@ -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"),
]

View file

@ -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):

View file

@ -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(

View file

@ -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(

View file

@ -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:

View file

@ -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)

View file

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

View file

@ -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} "

View file

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

View file

@ -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}
)