apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-26 00:46:22 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix non streaming RAG problems (#607)

Browse source 
* Fix non-streaming failure in RAG services

* Fix non-streaming failure in API

* Fix agent non-streaming messaging

* Agent messaging unit & contract tests
This commit is contained in:
cybermaggedon 2026-01-12 18:45:52 +00:00 • committed by GitHub
parent 30ca1d2e8b
commit 807f6cc4e2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 677 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

30
trustgraph-base/trustgraph/api/socket_client.py
View file

@ -275,13 +275,17 @@ class SocketFlowInstance:
result = self.client._send_request_sync("text-completion", self.flow_id, request, streaming)
if streaming:
# For text completion, yield just the content
for chunk in result:
if hasattr(chunk, 'content'):
yield chunk.content
# For text completion, return generator that yields content
return self._text_completion_generator(result)
else:
return result.get("response", "")
def _text_completion_generator(self, result: Iterator[StreamingChunk]) -> Iterator[str]:
"""Generator for text completion streaming"""
for chunk in result:
if hasattr(chunk, 'content'):
yield chunk.content
def graph_rag(
self,
query: str,
@ -308,9 +312,7 @@ class SocketFlowInstance:
result = self.client._send_request_sync("graph-rag", self.flow_id, request, streaming)
if streaming:
for chunk in result:
if hasattr(chunk, 'content'):
yield chunk.content
return self._rag_generator(result)
else:
return result.get("response", "")
@ -336,12 +338,16 @@ class SocketFlowInstance:
result = self.client._send_request_sync("document-rag", self.flow_id, request, streaming)
if streaming:
for chunk in result:
if hasattr(chunk, 'content'):
yield chunk.content
return self._rag_generator(result)
else:
return result.get("response", "")
def _rag_generator(self, result: Iterator[StreamingChunk]) -> Iterator[str]:
"""Generator for RAG streaming (graph-rag and document-rag)"""
for chunk in result:
if hasattr(chunk, 'content'):
yield chunk.content
def prompt(
self,
id: str,
@ -360,9 +366,7 @@ class SocketFlowInstance:
result = self.client._send_request_sync("prompt", self.flow_id, request, streaming)
if streaming:
for chunk in result:
if hasattr(chunk, 'content'):
yield chunk.content
return self._rag_generator(result)
else:
return result.get("response", "")