diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index e1b8f705..0810956d 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -118,20 +118,32 @@ class SocketClient: async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket: await websocket.send(json.dumps(message)) - # Wait for single response - raw_message = await websocket.recv() - response = json.loads(raw_message) + # Some services (like agent) send multiple messages even in non-streaming mode + # Collect all messages until complete=True + messages = [] - if response.get("id") != request_id: - raise ProtocolException(f"Response ID mismatch") + while True: + raw_message = await websocket.recv() + response = json.loads(raw_message) - if "error" in response: - raise_from_error_dict(response["error"]) + if response.get("id") != request_id: + raise ProtocolException(f"Response ID mismatch") - if "response" not in response: - raise ProtocolException(f"Missing response in message") + if "error" in response: + raise_from_error_dict(response["error"]) - return response["response"] + if "response" not in response: + raise ProtocolException(f"Missing response in message") + + messages.append(response["response"]) + + # Check if this is the final message + if response.get("complete", False): + break + + # For single-message responses, return the dict directly + # For multi-message responses (like agent), return a list + return messages[0] if len(messages) == 1 else messages async def _send_request_async_streaming( self, diff --git a/trustgraph-cli/trustgraph/cli/invoke_agent.py b/trustgraph-cli/trustgraph/cli/invoke_agent.py index de70021b..3fb876d4 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_agent.py +++ b/trustgraph-cli/trustgraph/cli/invoke_agent.py @@ -179,10 +179,27 @@ def question( else: # Non-streaming response - if "answer" in response: - print(response["answer"]) - if "error" in response: - raise RuntimeError(response["error"]) + # Response may be a list of messages (thought/observation/answer) or a single dict + messages = response if isinstance(response, list) else [response] + + for msg in messages: + # Display thoughts if verbose + if "thought" in msg and msg["thought"] and verbose: + output(wrap(msg["thought"]), "\U0001f914 ") + print() + + # Display observations if verbose + if "observation" in msg and msg["observation"] and verbose: + output(wrap(msg["observation"]), "\U0001f4a1 ") + print() + + # Display answer + if "answer" in msg and msg["answer"]: + print(msg["answer"]) + + # Handle errors + if "error" in msg: + raise RuntimeError(msg["error"]) finally: # Clean up socket connection