mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-14 07:42:11 +02:00
feat: structured source document references in graph-rag responses (#1035)
This commit is contained in:
parent
0374098ee9
commit
e9c6a850ad
20 changed files with 568 additions and 39 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
15
specs/api/components/schemas/rag/Source.yaml
Normal file
15
specs/api/components/schemas/rag/Source.yaml
Normal 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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
100
tests/unit/test_api/test_rag_chunk_sources.py
Normal file
100
tests/unit/test_api/test_rag_chunk_sources.py
Normal 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 == []
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
############################################################################
|
||||
|
||||
|
|
|
|||
|
|
@ -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} "
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue