apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-23 20:21:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Make GraphRAG not send everything twice

Browse source
This commit is contained in:
Cyber MacGeddon 2025-12-17 21:00:51 +00:00
parent 47d83a80de
commit 5cacbc1872
5 changed files with 40 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

31
trustgraph-base/trustgraph/base/prompt_client.py
View file

@ -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(

44
trustgraph-base/trustgraph/messaging/translators/retrieval.py
View file

@ -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

2
trustgraph-base/trustgraph/schema/services/retrieval.py
View file

@ -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

11
trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
View file

@ -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}

12
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
View file

@ -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}