mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
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:
parent
3c675b8cfc
commit
f0c95a4c5e
5 changed files with 29 additions and 42 deletions
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue