From 5cacbc1872069df93e706155270e71765e1f52a6 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 17 Dec 2025 21:00:51 +0000 Subject: [PATCH] Make GraphRAG not send everything twice --- .../trustgraph/base/prompt_client.py | 31 +++++++------ .../messaging/translators/retrieval.py | 44 ++++++------------- .../trustgraph/schema/services/retrieval.py | 2 - .../trustgraph/retrieval/document_rag/rag.py | 11 +++-- .../trustgraph/retrieval/graph_rag/rag.py | 12 ++--- 5 files changed, 40 insertions(+), 60 deletions(-) diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py index 307a118a..55c54cfc 100644 --- a/trustgraph-base/trustgraph/base/prompt_client.py +++ b/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 - full_text = "" - full_object = None + # Streaming path - just forward chunks, don't accumulate + last_text = "" + last_object = None - async def collect_chunks(resp): - nonlocal full_text, full_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)}") + async def forward_chunks(resp): + nonlocal last_text, last_object + 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 - logger.info(f"DEBUG prompt_client: Accumulated {len(full_text)} chars") + last_text = resp.text # 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: - logger.info("DEBUG prompt_client: Returning full_text") - return full_text + if last_text: + logger.info("DEBUG prompt_client: Returning last_text") + return last_text - logger.info("DEBUG prompt_client: Returning parsed full_object") - return json.loads(full_object) + logger.info("DEBUG prompt_client: Returning parsed last_object") + return json.loads(last_object) if last_object else None async def extract_definitions(self, text, timeout=600): return await self.prompt( diff --git a/trustgraph-base/trustgraph/messaging/translators/retrieval.py b/trustgraph-base/trustgraph/messaging/translators/retrieval.py index 441a9d18..d2161cff 100644 --- a/trustgraph-base/trustgraph/messaging/translators/retrieval.py +++ b/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) - if hasattr(obj, 'chunk') and obj.chunk: - result["chunk"] = obj.chunk - result["end_of_stream"] = getattr(obj, "end_of_stream", False) - else: - # Non-streaming response - if obj.response: - result["response"] = obj.response + # Include response content (chunk or complete) + if obj.response: + result["response"] = obj.response + + # Include end_of_stream flag + result["end_of_stream"] = getattr(obj, "end_of_stream", False) # 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 - 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 - + is_final = getattr(obj, 'end_of_stream', False) 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) - if hasattr(obj, 'chunk') and obj.chunk: - result["chunk"] = obj.chunk - result["end_of_stream"] = getattr(obj, "end_of_stream", False) - else: - # Non-streaming response - if obj.response: - result["response"] = obj.response + # Include response content (chunk or complete) + if obj.response: + result["response"] = obj.response + + # Include end_of_stream flag + result["end_of_stream"] = getattr(obj, "end_of_stream", False) # 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 - 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 - + is_final = getattr(obj, 'end_of_stream', False) return self.from_pulsar(obj), is_final \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/retrieval.py b/trustgraph-base/trustgraph/schema/services/retrieval.py index 813fbaa0..72085ae8 100644 --- a/trustgraph-base/trustgraph/schema/services/retrieval.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py index 670d71a1..ec67a072 100755 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py +++ b/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} diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 565921a3..de1f0e24 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/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}