apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-23 20:21:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Fix end-of-stream anomaly with some graph-rag and document-rag

Browse source
This commit is contained in:
Cyber MacGeddon 2026-01-06 15:34:34 +00:00
parent 3c675b8cfc
commit 00292940c4
3 changed files with 16 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

11
trustgraph-base/trustgraph/base/prompt_client.py
View file

@ -49,20 +49,21 @@ class PromptClient(RequestResponse):
logger.error(f"DEBUG prompt_client: Error in response: {resp.error.message}")
raise RuntimeError(resp.error.message)
end_stream = getattr(resp, 'end_of_stream', False)
if resp.text:
last_text = resp.text
# Call chunk callback if provided
# Call chunk callback if provided with both chunk and end_of_stream flag
if chunk_callback:
logger.info(f"DEBUG prompt_client: Calling chunk_callback")
logger.info(f"DEBUG prompt_client: Calling chunk_callback with end_of_stream={end_stream}")
if asyncio.iscoroutinefunction(chunk_callback):
await chunk_callback(resp.text)
await chunk_callback(resp.text, end_stream)
else:
chunk_callback(resp.text)
chunk_callback(resp.text, end_stream)
elif resp.object:
logger.info(f"DEBUG prompt_client: Got object response")
last_object = resp.object
end_stream = getattr(resp, 'end_of_stream', False)
logger.info(f"DEBUG prompt_client: Returning end_of_stream={end_stream}")
return end_stream

19
trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
View file

@ -95,19 +95,20 @@ class Processor(FlowProcessor):
# Check if streaming is requested
if v.streaming:
# Define async callback for streaming chunks
async def send_chunk(chunk):
# Receives chunk text and end_of_stream flag from prompt client
async def send_chunk(chunk, end_of_stream):
await flow("response").send(
DocumentRagResponse(
response=chunk,
end_of_stream=False,
end_of_stream=end_of_stream,
error=None
),
properties={"id": id}
)
# Query with streaming enabled
# The query returns the last chunk (not accumulated text)
final_response = await self.rag.query(
# All chunks (including final one with end_of_stream=True) are sent via callback
await self.rag.query(
v.query,
user=v.user,
collection=v.collection,
@ -115,16 +116,6 @@ class Processor(FlowProcessor):
streaming=True,
chunk_callback=send_chunk,
)
# Send final message with last chunk
await flow("response").send(
DocumentRagResponse(
response=final_response if final_response else "",
end_of_stream=True,
error=None
),
properties={"id": id}
)
else:
# Non-streaming path (existing behavior)
response = await self.rag.query(

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

@ -138,19 +138,20 @@ class Processor(FlowProcessor):
# Check if streaming is requested
if v.streaming:
# Define async callback for streaming chunks
async def send_chunk(chunk):
# Receives chunk text and end_of_stream flag from prompt client
async def send_chunk(chunk, end_of_stream):
await flow("response").send(
GraphRagResponse(
response=chunk,
end_of_stream=False,
end_of_stream=end_of_stream,
error=None
),
properties={"id": id}
)
# Query with streaming enabled
# The query will send chunks via callback AND return the complete text
final_response = await rag.query(
# All chunks (including final one with end_of_stream=True) are sent via callback
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,
@ -158,17 +159,6 @@ class Processor(FlowProcessor):
streaming = True,
chunk_callback = send_chunk,
)
# Send final message - may have last chunk of content with end_of_stream=True
# (prompt service may send final chunk with text, so we pass through whatever we got)
await flow("response").send(
GraphRagResponse(
response=final_response if final_response else "",
end_of_stream=True,
error=None
),
properties={"id": id}
)
else:
# Non-streaming path (existing behavior)
response = await rag.query(