Make GraphRAG not send everything twice

trustgraph-base/trustgraph/base/prompt_client.py

@@ -37,21 +37,20 @@ class PromptClient(RequestResponse):
 
         else:
             logger.info("DEBUG prompt_client: Streaming path")
-            # Streaming path - collect all chunks
+            # Streaming path - just forward chunks, don't accumulate
-            full_text = ""
+            last_text = ""
-            full_object = None
+            last_object = None
 
-            async def collect_chunks(resp):
+            async def forward_chunks(resp):
-                nonlocal full_text, full_object
+                nonlocal last_text, last_object
-                logger.info(f"DEBUG prompt_client: collect_chunks called, resp.text={resp.text[:50] if resp.text else None}, end_of_stream={getattr(resp, 'end_of_stream', False)}")
+                logger.info(f"DEBUG prompt_client: forward_chunks called, resp.text={resp.text[:50] if resp.text else None}, end_of_stream={getattr(resp, 'end_of_stream', False)}")
 
                 if resp.error:
                     logger.error(f"DEBUG prompt_client: Error in response: {resp.error.message}")
                     raise RuntimeError(resp.error.message)
 
                 if resp.text:
-                    full_text += resp.text
+                    last_text = resp.text
-                    logger.info(f"DEBUG prompt_client: Accumulated {len(full_text)} chars")
                     # Call chunk callback if provided
                     if chunk_callback:
                         logger.info(f"DEBUG prompt_client: Calling chunk_callback")
@@ -61,7 +60,7 @@ class PromptClient(RequestResponse):
                             chunk_callback(resp.text)
                 elif resp.object:
                     logger.info(f"DEBUG prompt_client: Got object response")
-                    full_object = resp.object
+                    last_object = resp.object
 
                 end_stream = getattr(resp, 'end_of_stream', False)
                 logger.info(f"DEBUG prompt_client: Returning end_of_stream={end_stream}")
@@ -79,17 +78,17 @@ class PromptClient(RequestResponse):
             logger.info(f"DEBUG prompt_client: About to call self.request with recipient, timeout={timeout}")
             await self.request(
                 req,
-                recipient=collect_chunks,
+                recipient=forward_chunks,
                 timeout=timeout
             )
-            logger.info(f"DEBUG prompt_client: self.request returned, full_text has {len(full_text)} chars")
+            logger.info(f"DEBUG prompt_client: self.request returned, last_text={last_text[:50] if last_text else None}")
 
-            if full_text:
+            if last_text:
-                logger.info("DEBUG prompt_client: Returning full_text")
+                logger.info("DEBUG prompt_client: Returning last_text")
-                return full_text
+                return last_text
 
-            logger.info("DEBUG prompt_client: Returning parsed full_object")
+            logger.info("DEBUG prompt_client: Returning parsed last_object")
-            return json.loads(full_object)
+            return json.loads(last_object) if last_object else None
 
     async def extract_definitions(self, text, timeout=600):
         return await self.prompt(

trustgraph-base/trustgraph/messaging/translators/retrieval.py

@@ -34,14 +34,12 @@ class DocumentRagResponseTranslator(MessageTranslator):
     def from_pulsar(self, obj: DocumentRagResponse) -> Dict[str, Any]:
         result = {}
 
-        # Check if this is a streaming response (has chunk)
+        # Include response content (chunk or complete)
-        if hasattr(obj, 'chunk') and obj.chunk:
+        if obj.response:
-            result["chunk"] = obj.chunk
+            result["response"] = obj.response
-            result["end_of_stream"] = getattr(obj, "end_of_stream", False)
+
-        else:
+        # Include end_of_stream flag
-            # Non-streaming response
+        result["end_of_stream"] = getattr(obj, "end_of_stream", False)
-            if obj.response:
-                result["response"] = obj.response
 
         # Always include error if present
         if hasattr(obj, 'error') and obj.error and obj.error.message:
@@ -51,13 +49,7 @@ class DocumentRagResponseTranslator(MessageTranslator):
 
     def from_response_with_completion(self, obj: DocumentRagResponse) -> Tuple[Dict[str, Any], bool]:
         """Returns (response_dict, is_final)"""
-        # For streaming responses, check end_of_stream
+        is_final = getattr(obj, 'end_of_stream', False)
-        if hasattr(obj, 'chunk') and obj.chunk:
-            is_final = getattr(obj, 'end_of_stream', False)
-        else:
-            # For non-streaming responses, it's always final
-            is_final = True
-
         return self.from_pulsar(obj), is_final
 
 
@@ -98,14 +90,12 @@ class GraphRagResponseTranslator(MessageTranslator):
     def from_pulsar(self, obj: GraphRagResponse) -> Dict[str, Any]:
         result = {}
 
-        # Check if this is a streaming response (has chunk)
+        # Include response content (chunk or complete)
-        if hasattr(obj, 'chunk') and obj.chunk:
+        if obj.response:
-            result["chunk"] = obj.chunk
+            result["response"] = obj.response
-            result["end_of_stream"] = getattr(obj, "end_of_stream", False)
+
-        else:
+        # Include end_of_stream flag
-            # Non-streaming response
+        result["end_of_stream"] = getattr(obj, "end_of_stream", False)
-            if obj.response:
-                result["response"] = obj.response
 
         # Always include error if present
         if hasattr(obj, 'error') and obj.error and obj.error.message:
@@ -115,11 +105,5 @@ class GraphRagResponseTranslator(MessageTranslator):
 
     def from_response_with_completion(self, obj: GraphRagResponse) -> Tuple[Dict[str, Any], bool]:
         """Returns (response_dict, is_final)"""
-        # For streaming responses, check end_of_stream
+        is_final = getattr(obj, 'end_of_stream', False)
-        if hasattr(obj, 'chunk') and obj.chunk:
-            is_final = getattr(obj, 'end_of_stream', False)
-        else:
-            # For non-streaming responses, it's always final
-            is_final = True
-
         return self.from_pulsar(obj), is_final

trustgraph-base/trustgraph/schema/services/retrieval.py

@@ -21,7 +21,6 @@ class GraphRagQuery:
 class GraphRagResponse:
     error: Error | None = None
     response: str = ""
-    chunk: str = ""
     end_of_stream: bool = False
 
 ############################################################################
@@ -40,5 +39,4 @@ class DocumentRagQuery:
 class DocumentRagResponse:
     error: Error | None = None
     response: str = ""
-    chunk: str = ""
     end_of_stream: bool = False

trustgraph-flow/trustgraph/retrieval/document_rag/rag.py

@@ -98,16 +98,16 @@ class Processor(FlowProcessor):
                 async def send_chunk(chunk):
                     await flow("response").send(
                         DocumentRagResponse(
-                            chunk=chunk,
+                            response=chunk,
                             end_of_stream=False,
-                            response=None,
                             error=None
                         ),
                         properties={"id": id}
                     )
 
                 # Query with streaming enabled
-                full_response = await self.rag.query(
+                # The query returns the last chunk (not accumulated text)
+                final_response = await self.rag.query(
                     v.query,
                     user=v.user,
                     collection=v.collection,
@@ -116,12 +116,11 @@ class Processor(FlowProcessor):
                     chunk_callback=send_chunk,
                 )
 
-                # Send final message with complete response
+                # Send final message with last chunk
                 await flow("response").send(
                     DocumentRagResponse(
-                        chunk=None,
+                        response=final_response if final_response else "",
                         end_of_stream=True,
-                        response=full_response,
                         error=None
                     ),
                     properties={"id": id}

trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py

@@ -141,16 +141,16 @@ class Processor(FlowProcessor):
                 async def send_chunk(chunk):
                     await flow("response").send(
                         GraphRagResponse(
-                            chunk=chunk,
+                            response=chunk,
                             end_of_stream=False,
-                            response=None,
                             error=None
                         ),
                         properties={"id": id}
                     )
 
                 # Query with streaming enabled
-                full_response = await rag.query(
+                # The query will send chunks via callback AND return the complete text
+                final_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,
@@ -159,12 +159,12 @@ class Processor(FlowProcessor):
                     chunk_callback = send_chunk,
                 )
 
-                # Send final message with complete response
+                # 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(
-                        chunk=None,
+                        response=final_response if final_response else "",
                         end_of_stream=True,
-                        response=full_response,
                         error=None
                     ),
                     properties={"id": id}