Fix streaming API niggles (#599)

* Fix end-of-stream anomally with some graph-rag and document-rag

* Fix gateway translators dropping responses
This commit is contained in:
cybermaggedon 2026-01-06 16:41:35 +00:00 committed by GitHub
parent 3c675b8cfc
commit f0c95a4c5e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 29 additions and 42 deletions

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(