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

Fix streaming

Browse source
This commit is contained in:
Cyber MacGeddon 2025-12-04 20:37:23 +00:00
parent 593caa08c2
commit f511b8db50
3 changed files with 19 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

14
trustgraph-base/trustgraph/api/async_socket_client.py
View file

@ -130,18 +130,26 @@ class AsyncSocketClient:
content=resp.get("content", ""), content=resp.get("content", ""),
end_of_message=resp.get("end_of_message", False) end_of_message=resp.get("end_of_message", False)
) )
elif chunk_type == "final-answer": elif chunk_type == "answer" or chunk_type == "final-answer":
return AgentAnswer( return AgentAnswer(
content=resp.get("content", ""), content=resp.get("content", ""),
end_of_message=resp.get("end_of_message", False), end_of_message=resp.get("end_of_message", False),
end_of_dialog=resp.get("end_of_dialog", False) end_of_dialog=resp.get("end_of_dialog", False)
) )
elif chunk_type == "action":
# Agent action chunks - treat as thoughts for display purposes
return AgentThought(
content=resp.get("content", ""),
end_of_message=resp.get("end_of_message", False)
)
else: else:
# RAG-style chunk (or generic chunk) # RAG-style chunk (or generic chunk)
# Text-completion uses "response" field, RAG uses "chunk" field, Prompt uses "text" field
content = resp.get("response", resp.get("chunk", resp.get("text", "")))
return RAGChunk( return RAGChunk(
content=resp.get("chunk", ""), content=content,
end_of_stream=resp.get("end_of_stream", False), end_of_stream=resp.get("end_of_stream", False),
error=resp.get("error") error=None # Errors are always thrown, never stored
) )
async def aclose(self): async def aclose(self):

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

@ -202,12 +202,18 @@ class SocketClient:
content=resp.get("content", ""), content=resp.get("content", ""),
end_of_message=resp.get("end_of_message", False) end_of_message=resp.get("end_of_message", False)
) )
elif chunk_type == "final-answer": elif chunk_type == "answer" or chunk_type == "final-answer":
return AgentAnswer( return AgentAnswer(
content=resp.get("content", ""), content=resp.get("content", ""),
end_of_message=resp.get("end_of_message", False), end_of_message=resp.get("end_of_message", False),
end_of_dialog=resp.get("end_of_dialog", False) end_of_dialog=resp.get("end_of_dialog", False)
) )
elif chunk_type == "action":
# Agent action chunks - treat as thoughts for display purposes
return AgentThought(
content=resp.get("content", ""),
end_of_message=resp.get("end_of_message", False)
)
else: else:
# RAG-style chunk (or generic chunk) # RAG-style chunk (or generic chunk)
# Text-completion uses "response" field, RAG uses "chunk" field, Prompt uses "text" field # Text-completion uses "response" field, RAG uses "chunk" field, Prompt uses "text" field

1
trustgraph-base/trustgraph/api/types.py
View file

@ -79,5 +79,6 @@ class AgentAnswer(StreamingChunk):
@dataclasses.dataclass @dataclasses.dataclass
class RAGChunk(StreamingChunk): class RAGChunk(StreamingChunk):
"""RAG streaming chunk""" """RAG streaming chunk"""
chunk_type: str = "rag"
end_of_stream: bool = False end_of_stream: bool = False
error: Optional[Dict[str, str]] = None error: Optional[Dict[str, str]] = None