apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-25 08:26:21 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Streaming rag responses (#568)

Browse source 
* Tech spec for streaming RAG

* Support for streaming Graph/Doc RAG
This commit is contained in:
cybermaggedon 2025-11-26 19:47:39 +00:00 • committed by GitHub
parent b1cc724f7d
commit 1948edaa50
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 3087 additions and 94 deletions
Download patch file Download diff file Expand all files Collapse all files

83
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
View file

@ -135,20 +135,56 @@ class Processor(FlowProcessor):
else:
max_path_length = self.default_max_path_length
response = await rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
)
# Check if streaming is requested
if v.streaming:
# Define async callback for streaming chunks
async def send_chunk(chunk):
await flow("response").send(
GraphRagResponse(
chunk=chunk,
end_of_stream=False,
response=None,
error=None
),
properties={"id": id}
)
await flow("response").send(
GraphRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
# Query with streaming enabled
full_response = await rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
streaming = True,
chunk_callback = send_chunk,
)
# Send final message with complete response
await flow("response").send(
GraphRagResponse(
chunk=None,
end_of_stream=True,
response=full_response,
error=None
),
properties={"id": id}
)
else:
# Non-streaming path (existing behavior)
response = await rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
)
await flow("response").send(
GraphRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
logger.info("Request processing complete")
@ -158,14 +194,21 @@ class Processor(FlowProcessor):
logger.debug("Sending error response...")
await flow("response").send(
GraphRagResponse(
response = None,
error = Error(
type = "graph-rag-error",
message = str(e),
),
# Send error response with end_of_stream flag if streaming was requested
error_response = GraphRagResponse(
response = None,
error = Error(
type = "graph-rag-error",
message = str(e),
),
)
# If streaming was requested, indicate stream end
if v.streaming:
error_response.end_of_stream = True
await flow("response").send(
error_response,
properties = {"id": id}
)