diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index e1b8f705..23e3dbc0 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -214,6 +214,23 @@ class SocketClient: content=resp.get("content", ""), end_of_message=resp.get("end_of_message", False) ) + # Non-streaming agent format: chunk_type is empty but has thought/observation/answer fields + elif resp.get("thought"): + return AgentThought( + content=resp.get("thought", ""), + end_of_message=resp.get("end_of_message", False) + ) + elif resp.get("observation"): + return AgentObservation( + content=resp.get("observation", ""), + end_of_message=resp.get("end_of_message", False) + ) + elif resp.get("answer"): + return AgentAnswer( + content=resp.get("answer", ""), + end_of_message=resp.get("end_of_message", False), + end_of_dialog=resp.get("end_of_dialog", False) + ) else: # RAG-style chunk (or generic chunk) # Text-completion uses "response" field, RAG uses "chunk" field, Prompt uses "text" field @@ -261,7 +278,9 @@ class SocketFlowInstance: request["history"] = history request.update(kwargs) - return self.client._send_request_sync("agent", self.flow_id, request, streaming) + # Agents always use multipart messaging (multiple complete messages) + # regardless of streaming flag, so always use the streaming code path + return self.client._send_request_sync("agent", self.flow_id, request, streaming=True) def text_completion(self, system: str, prompt: str, streaming: bool = False, **kwargs) -> Union[str, Iterator[str]]: """Text completion with optional streaming""" diff --git a/trustgraph-cli/trustgraph/cli/invoke_agent.py b/trustgraph-cli/trustgraph/cli/invoke_agent.py index de70021b..369fcdd4 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_agent.py +++ b/trustgraph-cli/trustgraph/cli/invoke_agent.py @@ -178,11 +178,22 @@ def question( print() else: - # Non-streaming response - if "answer" in response: - print(response["answer"]) - if "error" in response: - raise RuntimeError(response["error"]) + # Non-streaming response - but agents use multipart messaging + # so we iterate through the chunks (which are complete messages, not text chunks) + for chunk in response: + # Display thoughts if verbose + if chunk.chunk_type == "thought" and verbose: + output(wrap(chunk.content), "\U0001f914 ") + print() + + # Display observations if verbose + elif chunk.chunk_type == "observation" and verbose: + output(wrap(chunk.content), "\U0001f4a1 ") + print() + + # Display answer + elif chunk.chunk_type == "final-answer" or chunk.chunk_type == "answer": + print(chunk.content) finally: # Clean up socket connection