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

Fix non-streaming failure in API

Browse source
This commit is contained in:
Cyber MacGeddon 2026-01-12 17:46:07 +00:00
parent aacc96064f
commit 69f10e50b3
1 changed files with 17 additions and 13 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", "")