From 00292940c45a9204e0dd0d6858434095616a16e6 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Tue, 6 Jan 2026 15:34:34 +0000
Subject: [PATCH] Fix end-of-stream anomaly with some graph-rag and
 document-rag

---
 .../trustgraph/base/prompt_client.py          | 11 +++++-----
 .../trustgraph/retrieval/document_rag/rag.py  | 19 +++++-------------
 .../trustgraph/retrieval/graph_rag/rag.py     | 20 +++++--------------
 3 files changed, 16 insertions(+), 34 deletions(-)

diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py
index 55c54cfc..370cf78a 100644
--- a/trustgraph-base/trustgraph/base/prompt_client.py
+++ b/trustgraph-base/trustgraph/base/prompt_client.py
@@ -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
 
diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
index ec67a072..14d71d97 100755
--- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
+++ b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
@@ -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(
diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
index de1f0e24..d159dbae 100755
--- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
+++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
@@ -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(