| 4 | from websockets.exceptions import ConnectionClosed |
| 5 | |
| 6 | class TermNetClient: |
| 7 | def __init__(self, uri="ws://localhost:876"): |
| 8 | self.uri = uri |
| 9 | self.websocket = None |
| 10 | self.connected = False |
| 11 | |
| 12 | async def connect(self): |
| 13 | try: |
| 14 | self.websocket = await websockets.connect(self.uri) |
| 15 | welcome = await self.websocket.recv() |
| 16 | data = json.loads(welcome) |
| 17 | print(f"Connected: {data.get('message', 'Ready')}") |
| 18 | self.connected = True |
| 19 | return True |
| 20 | except Exception as e: |
| 21 | print(f"Connection failed: {e}") |
| 22 | return False |
| 23 | |
| 24 | async def send_message(self, message): |
| 25 | # Auto-reconnect silently if needed |
| 26 | if not self.connected or not self.websocket: |
| 27 | if not await self.reconnect(): |
| 28 | return False |
| 29 | |
| 30 | try: |
| 31 | msg_data = { |
| 32 | "type": "message", |
| 33 | "message": message, |
| 34 | "timestamp": asyncio.get_event_loop().time() |
| 35 | } |
| 36 | |
| 37 | await self.websocket.send(json.dumps(msg_data)) |
| 38 | print("AI: ", end="", flush=True) |
| 39 | |
| 40 | full_response = "" |
| 41 | while True: |
| 42 | try: |
| 43 | response = await asyncio.wait_for(self.websocket.recv(), timeout=60.0) |
| 44 | data = json.loads(response) |
| 45 | msg_type = data.get("type", "") |
| 46 | |
| 47 | if msg_type == "response_chunk": |
| 48 | chunk = data.get("chunk", "") |
| 49 | print(chunk, end="", flush=True) |
| 50 | full_response += chunk |
| 51 | elif msg_type == "tool_execution": |
| 52 | tool_msg = data.get("message", "") |
| 53 | print(f"\nTool: {tool_msg}\nAI: ", end="", flush=True) |
| 54 | elif msg_type == "response_end": |
| 55 | break |
| 56 | elif msg_type == "error": |
| 57 | error_msg = data.get("message", "") |
| 58 | print(f"\nError: {error_msg}") |
| 59 | break |
| 60 | |
| 61 | except asyncio.TimeoutError: |
| 62 | # Silently timeout - don't print error |
| 63 | break |