forked from xiaohei/taiji-AI-PAD
283 lines
8.3 KiB
Python
283 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
MCP Server WebSocket 客户端示例
|
|
演示如何通过 WebSocket 与 Agent 进行实时通信
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
import websocket
|
|
from typing import Optional
|
|
import threading
|
|
|
|
# 配置
|
|
WEBSOCKET_URL = "ws://localhost:8002/ws/{agent_id}"
|
|
|
|
|
|
class MCPWebSocketClient:
|
|
"""MCP WebSocket 客户端"""
|
|
|
|
def __init__(self, agent_id: str):
|
|
self.agent_id = agent_id
|
|
self.url = WEBSOCKET_URL.format(agent_id=agent_id)
|
|
self.ws: Optional[websocket.WebSocketApp] = None
|
|
self.connected = False
|
|
self.messages = []
|
|
|
|
def on_message(self, ws, message):
|
|
"""收到消息时的回调"""
|
|
try:
|
|
data = json.loads(message)
|
|
msg_type = data.get('type')
|
|
|
|
print(f"\n📨 收到消息 [{msg_type}]:")
|
|
|
|
if msg_type == 'welcome':
|
|
print(f" 欢迎连接到 Agent: {data.get('agent_name')}")
|
|
print(f" Agent ID: {data.get('agent_id')}")
|
|
print(f" 时间: {data.get('timestamp')}")
|
|
|
|
elif msg_type == 'heartbeat':
|
|
print(f" 心跳: {data.get('timestamp')}")
|
|
|
|
elif msg_type == 'pong':
|
|
print(f" Pong 响应: {data.get('timestamp')}")
|
|
|
|
elif msg_type == 'mcp_response':
|
|
payload = data.get('payload', {})
|
|
print(f" 执行结果:")
|
|
print(f" 成功: {payload.get('success')}")
|
|
print(f" 结果: {payload.get('result')}")
|
|
print(f" 执行时间: {payload.get('execution_time', 0):.3f}s")
|
|
|
|
elif msg_type == 'error':
|
|
print(f" ❌ 错误: {data.get('error')}")
|
|
if 'request_id' in data:
|
|
print(f" 请求 ID: {data.get('request_id')}")
|
|
|
|
else:
|
|
print(f" 未知消息类型: {json.dumps(data, indent=2, ensure_ascii=False)}")
|
|
|
|
self.messages.append(data)
|
|
|
|
except json.JSONDecodeError:
|
|
print(f" 无法解析消息: {message}")
|
|
|
|
def on_error(self, ws, error):
|
|
"""错误回调"""
|
|
print(f"\n❌ WebSocket 错误: {error}")
|
|
|
|
def on_close(self, ws, close_status_code, close_msg):
|
|
"""关闭回调"""
|
|
print(f"\n🔌 WebSocket 连接已关闭")
|
|
print(f" 状态码: {close_status_code}")
|
|
print(f" 消息: {close_msg}")
|
|
self.connected = False
|
|
|
|
def on_open(self, ws):
|
|
"""连接建立回调"""
|
|
print(f"\n✓ WebSocket 连接已建立")
|
|
print(f" URL: {self.url}")
|
|
self.connected = True
|
|
|
|
def connect(self):
|
|
"""连接到 WebSocket"""
|
|
print(f"\n连接到 WebSocket: {self.url}")
|
|
|
|
self.ws = websocket.WebSocketApp(
|
|
self.url,
|
|
on_open=self.on_open,
|
|
on_message=self.on_message,
|
|
on_error=self.on_error,
|
|
on_close=self.on_close
|
|
)
|
|
|
|
# 在后台线程运行
|
|
wst = threading.Thread(target=self.ws.run_forever)
|
|
wst.daemon = True
|
|
wst.start()
|
|
|
|
# 等待连接建立
|
|
timeout = 5
|
|
start_time = time.time()
|
|
while not self.connected and (time.time() - start_time) < timeout:
|
|
time.sleep(0.1)
|
|
|
|
if not self.connected:
|
|
raise TimeoutError("WebSocket 连接超时")
|
|
|
|
def send_ping(self):
|
|
"""发送 ping 消息"""
|
|
if not self.connected:
|
|
print("未连接，无法发送消息")
|
|
return
|
|
|
|
message = {
|
|
"type": "ping"
|
|
}
|
|
|
|
print(f"\n📤 发送 Ping")
|
|
self.ws.send(json.dumps(message))
|
|
|
|
def execute_tool(self, tool_name: str, arguments: dict):
|
|
"""执行工具"""
|
|
if not self.connected:
|
|
print("未连接，无法发送消息")
|
|
return
|
|
|
|
message = {
|
|
"type": "mcp_request",
|
|
"payload": {
|
|
"jsonrpc": "2.0",
|
|
"id": f"req-{int(time.time() * 1000)}",
|
|
"method": "tools/call",
|
|
"params": {
|
|
"tool": {
|
|
"name": tool_name,
|
|
"function_name": tool_name
|
|
},
|
|
"arguments": arguments,
|
|
"context": {}
|
|
}
|
|
}
|
|
}
|
|
|
|
print(f"\n📤 发送工具执行请求:")
|
|
print(f" 工具: {tool_name}")
|
|
print(f" 参数: {json.dumps(arguments, ensure_ascii=False)}")
|
|
|
|
self.ws.send(json.dumps(message))
|
|
|
|
def close(self):
|
|
"""关闭连接"""
|
|
if self.ws:
|
|
self.ws.close()
|
|
|
|
|
|
def create_test_agent() -> Optional[str]:
|
|
"""创建测试 Agent"""
|
|
import requests
|
|
|
|
print("\n创建测试 Agent...")
|
|
|
|
agent_data = {
|
|
"name": f"websocket-test-agent-{int(time.time())}",
|
|
"description": "Agent for WebSocket testing",
|
|
"role": "assistant",
|
|
"goal": "Test WebSocket communication",
|
|
"tools": ["math_add", "string_upper", "datetime_now"],
|
|
"config": {},
|
|
"capabilities": ["websocket_communication"]
|
|
}
|
|
|
|
response = requests.post(
|
|
"http://localhost:8002/agents",
|
|
json=agent_data,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-API-Key": "test-api-key"
|
|
}
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
agent = response.json()
|
|
agent_id = agent['id']
|
|
print(f"✓ Agent 创建成功: {agent_id}")
|
|
return agent_id
|
|
else:
|
|
print(f"✗ Agent 创建失败: {response.text}")
|
|
return None
|
|
|
|
|
|
def main():
|
|
"""主函数"""
|
|
print("=" * 60)
|
|
print("MCP Server WebSocket 客户端示例")
|
|
print("=" * 60)
|
|
|
|
try:
|
|
# 创建测试 Agent
|
|
agent_id = create_test_agent()
|
|
|
|
if not agent_id:
|
|
print("\n无法创建测试 Agent，退出")
|
|
return
|
|
|
|
# 创建 WebSocket 客户端
|
|
client = MCPWebSocketClient(agent_id)
|
|
|
|
# 连接
|
|
client.connect()
|
|
|
|
print("\n开始测试...")
|
|
|
|
# 等待欢迎消息
|
|
time.sleep(1)
|
|
|
|
# 测试 1: 发送 ping
|
|
print("\n" + "=" * 60)
|
|
print("测试 1: Ping/Pong")
|
|
print("=" * 60)
|
|
client.send_ping()
|
|
time.sleep(1)
|
|
|
|
# 测试 2: 执行数学函数
|
|
print("\n" + "=" * 60)
|
|
print("测试 2: 执行数学函数 (math_add)")
|
|
print("=" * 60)
|
|
client.execute_tool("math_add", {"a": 15, "b": 25})
|
|
time.sleep(2)
|
|
|
|
# 测试 3: 执行字符串函数
|
|
print("\n" + "=" * 60)
|
|
print("测试 3: 执行字符串函数 (string_upper)")
|
|
print("=" * 60)
|
|
client.execute_tool("string_upper", {"s": "hello websocket"})
|
|
time.sleep(2)
|
|
|
|
# 测试 4: 获取当前时间
|
|
print("\n" + "=" * 60)
|
|
print("测试 4: 获取当前时间 (datetime_now)")
|
|
print("=" * 60)
|
|
client.execute_tool("datetime_now", {})
|
|
time.sleep(2)
|
|
|
|
# 测试 5: 连续发送多个请求
|
|
print("\n" + "=" * 60)
|
|
print("测试 5: 连续发送多个请求")
|
|
print("=" * 60)
|
|
|
|
for i in range(3):
|
|
client.execute_tool("math_add", {"a": i * 10, "b": i * 5})
|
|
time.sleep(0.5)
|
|
|
|
time.sleep(2)
|
|
|
|
print("\n" + "=" * 60)
|
|
print("WebSocket 测试完成！")
|
|
print("=" * 60)
|
|
print(f"\n总共收到 {len(client.messages)} 条消息")
|
|
print(f"测试 Agent ID: {agent_id}")
|
|
|
|
# 关闭连接
|
|
print("\n关闭连接...")
|
|
client.close()
|
|
time.sleep(1)
|
|
|
|
except TimeoutError as e:
|
|
print(f"\n❌ 超时错误: {e}")
|
|
except ConnectionRefusedError:
|
|
print("\n❌ 错误: 无法连接到 MCP Server")
|
|
print("请确保服务正在运行: docker-compose up -d mcp-server")
|
|
except Exception as e:
|
|
print(f"\n❌ 发生错误: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 需要安装 websocket-client
|
|
# pip install websocket-client
|
|
main()
|