MaiBot/websocket_client_example.py

59 lines
2.1 KiB
Python

import asyncio
import websockets
import json
from datetime import datetime
async def connect_to_output_manager():
"""连接到OutputManager的WebSocket服务器并接收消息"""
uri = "ws://localhost:8765"
print(f"正在连接到 {uri}...")
try:
async with websockets.connect(uri) as websocket:
print("已连接到服务器")
while True:
# 接收消息
message = await websocket.recv()
data = json.loads(message)
# 格式化时间戳
timestamp = datetime.fromisoformat(data["timestamp"]).strftime(
"%H:%M:%S"
)
# 根据消息类型和来源显示不同颜色
if data["type"] == "output":
if data["source"] == "stdout":
# 标准输出 - 白色
print(f"\033[37m[{timestamp}] {data['content']}\033[0m", end="")
elif data["source"] == "stderr":
# 错误输出 - 红色
print(f"\033[31m[{timestamp}] {data['content']}\033[0m", end="")
elif data["type"] == "connection":
# 连接消息 - 绿色
print(f"\033[32m[{timestamp}] {data['message']}\033[0m")
elif data["type"] == "history":
# 历史记录 - 蓝色
print(
f"\033[34m[{timestamp}] 收到历史记录: {len(data['messages'])} 条消息\033[0m"
)
else:
# 其他消息 - 黄色
content = data.get("content", "")
if isinstance(content, dict):
content = json.dumps(content, ensure_ascii=False)
print(f"\033[33m[{timestamp}] [{data['type']}] {content}\033[0m")
except websockets.exceptions.ConnectionClosed:
print("与服务器的连接已断开")
except Exception as e:
print(f"发生错误: {e}")
if __name__ == "__main__":
# 运行客户端
asyncio.run(connect_to_output_manager())