255 lines
7.9 KiB
Python
255 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
MCP 协议模拟器 - 模拟 xiaozhi-server 向 MCP 工具发送 JSON-RPC 消息
|
||
|
||
直接用 stdio 模式跑我们的 psycho-screener MCP 工具，
|
||
手动构造 xiaozhi 风格的 JSON-RPC 消息来测试完整流程。
|
||
|
||
使用方式：
|
||
python mcp_simulator.py
|
||
"""
|
||
|
||
import json
|
||
import subprocess
|
||
import sys
|
||
import os
|
||
import re
|
||
|
||
# 加入项目路径
|
||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
|
||
|
||
from psycho_screener.screener import PsychoScreener
|
||
|
||
|
||
# ============================================================================
|
||
# 测试语料（来自 docs/test_corpus.md）
|
||
# ============================================================================
|
||
|
||
TEST_CASES = [
|
||
{
|
||
"name": "霸凌场景",
|
||
"messages": [
|
||
{"role": "system", "content": "你是一个友爱的小智AI玩偶..."},
|
||
{"role": "user", "content": "今天幼儿园有个小朋友抢了我的玩具"},
|
||
{"role": "assistant", "content": "哎呀，那真是太过分了！"},
|
||
{"role": "user", "content": "他还不让我告诉老师，他说如果我告诉老师就会打我。我好害怕啊。"},
|
||
],
|
||
"expected_detected": True,
|
||
"expected_category": "bullying",
|
||
},
|
||
{
|
||
"name": "抑郁情绪",
|
||
"messages": [
|
||
{"role": "system", "content": "你是一个友爱的小智AI玩偶..."},
|
||
{"role": "user", "content": "最近我总觉得好累，什么都不想做"},
|
||
{"role": "assistant", "content": "是不是身体不舒服呀？"},
|
||
{"role": "user", "content": "我觉得我是一个很糟糕的人，没有人喜欢我。有时候我想，要是我不在了就好了。"},
|
||
],
|
||
"expected_detected": True,
|
||
"expected_category": "depression",
|
||
},
|
||
{
|
||
"name": "正常对话",
|
||
"messages": [
|
||
{"role": "system", "content": "你是一个友爱的小智AI玩偶..."},
|
||
{"role": "user", "content": "今天我画了一幅画，是一只大恐龙！"},
|
||
{"role": "assistant", "content": "哇，好厉害！"},
|
||
{"role": "user", "content": "是绿色的！晚上妈妈还做了红烧肉，好开心呀！"},
|
||
],
|
||
"expected_detected": False,
|
||
"expected_category": "none",
|
||
},
|
||
{
|
||
"name": "家庭矛盾",
|
||
"messages": [
|
||
{"role": "system", "content": "你是一个友爱的小智AI玩偶..."},
|
||
{"role": "user", "content": "昨天晚上爸爸妈妈吵架了，妈妈哭了"},
|
||
{"role": "user", "content": "我很害怕，怕他们会离婚。我总觉得是因为我表现不好他们才吵架的。"},
|
||
],
|
||
"expected_detected": True,
|
||
"expected_category": "family_conflict",
|
||
},
|
||
]
|
||
|
||
|
||
# ============================================================================
|
||
# 模拟 xiaozhi 的 MCP JSON-RPC 格式
|
||
# ============================================================================
|
||
|
||
def build_xiaozhi_mcp_request(messages: list[dict]) -> dict:
|
||
"""
|
||
构建 xiaozhi 风格的 MCP JSON-RPC 请求
|
||
|
||
xiaozhi 调用 MCP 工具时发送的格式：
|
||
{
|
||
"jsonrpc": "2.0",
|
||
"id": 1,
|
||
"method": "tools/call",
|
||
"params": {
|
||
"name": "psycho_screen",
|
||
"arguments": {
|
||
"messages": [...],
|
||
"include_prefix": true
|
||
}
|
||
}
|
||
}
|
||
"""
|
||
return {
|
||
"jsonrpc": "2.0",
|
||
"id": 1,
|
||
"method": "tools/call",
|
||
"params": {
|
||
"name": "psycho_screen",
|
||
"arguments": {
|
||
"messages": messages,
|
||
"include_prefix": True,
|
||
}
|
||
}
|
||
}
|
||
|
||
|
||
def parse_mcp_response(raw: str) -> dict:
|
||
"""
|
||
解析 MCP 响应
|
||
FastMCP 返回格式可能是：
|
||
- 纯 JSON 对象（直接结果）
|
||
- 带有 ```json 包裹
|
||
- 带有多行日志前缀
|
||
"""
|
||
content = raw.strip()
|
||
|
||
# 去掉日志前缀（时间戳等）
|
||
lines = content.split("\n")
|
||
json_lines = []
|
||
in_json = False
|
||
for line in lines:
|
||
if re.match(r'^\d{6}\s+\d{2}:\d{2}:\d{2}', line):
|
||
continue # 跳过日志行
|
||
json_lines.append(line)
|
||
|
||
content = "\n".join(json_lines).strip()
|
||
|
||
# 提取 JSON
|
||
md_match = re.search(r"\{.*\}", content, re.DOTALL)
|
||
if md_match:
|
||
content = md_match.group()
|
||
|
||
return json.loads(content)
|
||
|
||
|
||
# ============================================================================
|
||
# 直接调用（不通过 stdio，直接 Python 调用）
|
||
# ============================================================================
|
||
|
||
def test_direct(api_key: str):
|
||
"""直接 Python 函数调用，测试核心筛查逻辑"""
|
||
print("=" * 60)
|
||
print("测试1：直接 Python 调用（模拟 xiaozhi MCP 工具调用）")
|
||
print("=" * 60)
|
||
|
||
screener = PsychoScreener(api_key=api_key)
|
||
all_passed = True
|
||
|
||
for tc in TEST_CASES:
|
||
print(f"\n▶ 测试：{tc['name']}")
|
||
|
||
# 提取孩子消息
|
||
child_msgs = [m["content"] for m in tc["messages"] if m["role"] == "user"]
|
||
context = "\n".join(child_msgs)
|
||
print(f" 孩子的话：{context[:60]}...")
|
||
|
||
# 调用筛查器
|
||
result = screener.screen(context)
|
||
prefix = screener.build_response_prefix(result)
|
||
|
||
print(f" 结果：detected={result.detected}, category={result.category}, severity={result.severity}")
|
||
if prefix:
|
||
print(f" 前缀：{prefix[:80]}...")
|
||
|
||
# 验证
|
||
passed = (result.detected == tc["expected_detected"])
|
||
if tc["expected_detected"]:
|
||
passed = passed and (result.category == tc["expected_category"])
|
||
|
||
status = "✅ PASS" if passed else "❌ FAIL"
|
||
print(f" 验证：{status}")
|
||
|
||
if not passed:
|
||
all_passed = False
|
||
|
||
print("\n" + "=" * 60)
|
||
if all_passed:
|
||
print("✅ 全部测试通过！")
|
||
else:
|
||
print("❌ 有测试失败")
|
||
print("=" * 60)
|
||
return all_passed
|
||
|
||
|
||
def test_mcp_stdio(api_key: str):
|
||
"""通过 FastMCP stdio 模式调用（模拟 xiaozhi 的 MCP 管道）"""
|
||
print("\n" + "=" * 60)
|
||
print("测试2：FastMCP stdio 模式（模拟 xiaozhi MCP 管道）")
|
||
print("=" * 60)
|
||
|
||
# 设置环境变量
|
||
env = os.environ.copy()
|
||
env["MINIMAX_API_KEY"] = api_key
|
||
|
||
# 启动 MCP 进程
|
||
proc = subprocess.Popen(
|
||
[sys.executable, "-m", "psycho_screener.mcp_tool"],
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True,
|
||
cwd=os.path.join(os.path.dirname(__file__)),
|
||
env=env,
|
||
)
|
||
|
||
try:
|
||
for tc in TEST_CASES[:2]: # 只测前两个，节省时间
|
||
print(f"\n▶ 测试：{tc['name']}")
|
||
|
||
# 构建 JSON-RPC 请求
|
||
request = build_xiaozhi_mcp_request(tc["messages"])
|
||
request_json = json.dumps(request)
|
||
|
||
print(f" 发送：{request_json[:80]}...")
|
||
|
||
# 发送请求
|
||
proc.stdin.write(request_json + "\n")
|
||
proc.stdin.flush()
|
||
|
||
# 读取响应
|
||
import select
|
||
if select.select([proc.stdout], [], [], 30)[0]:
|
||
response_line = proc.stdout.readline()
|
||
print(f" 响应：{response_line[:120]}...")
|
||
else:
|
||
print(" ❌ 超时无响应")
|
||
|
||
finally:
|
||
proc.terminate()
|
||
proc.wait(timeout=5)
|
||
|
||
|
||
# ============================================================================
|
||
# 主程序
|
||
# ============================================================================
|
||
|
||
if __name__ == "__main__":
|
||
api_key = os.environ.get("MINIMAX_API_KEY", "")
|
||
if not api_key:
|
||
print("错误：需要设置 MINIMAX_API_KEY 环境变量")
|
||
print(" export MINIMAX_API_KEY=your-key")
|
||
sys.exit(1)
|
||
|
||
# 测试1：直接调用
|
||
ok = test_direct(api_key)
|
||
|
||
# 测试2：stdio MCP 模式
|
||
test_mcp_stdio(api_key)
|
||
|
||
sys.exit(0 if ok else 1)
|