589 lines
20 KiB
Python
Executable File
589 lines
20 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
WeChat Integration for bbav20311
|
|
基于 mcp-wechat-server 的微信集成模块
|
|
|
|
功能：
|
|
- 任务完成通知
|
|
- 接收用户指令
|
|
- 状态查询
|
|
|
|
作者: xialab的幸运星
|
|
日期: 2026-04-10
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import time
|
|
import asyncio
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any, Callable
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
|
|
|
|
@dataclass
|
|
class WeChatConfig:
|
|
"""微信配置"""
|
|
mcp_server_command: str = "bunx"
|
|
mcp_server_args: str = "mcp-wechat-server"
|
|
qrcode_path: Path = Path.home() / ".mcp-wechat-server" / "qrcode.txt"
|
|
qrcode_png_path: Path = Path.home() / ".mcp-wechat-server" / "qrcode.png"
|
|
state_file: Path = Path.home() / ".mcp-wechat-server" / "state.json"
|
|
|
|
# 任务通知模板
|
|
task_complete_template: str = """🎉 任务完成通知
|
|
|
|
📋 任务ID: {task_id}
|
|
🔬 分析类型: {task_type}
|
|
⏱️ 耗时: {duration}
|
|
📊 状态: {status}
|
|
|
|
📁 输出文件:
|
|
{output_files}
|
|
|
|
💡 提示: 回复"详情 {task_id}"查看完整结果"""
|
|
|
|
task_failed_template: str = """❌ 任务失败通知
|
|
|
|
📋 任务ID: {task_id}
|
|
🔬 分析类型: {task_type}
|
|
⏱️ 耗时: {duration}
|
|
❗ 错误: {error}
|
|
|
|
请检查输入数据或联系管理员"""
|
|
|
|
|
|
class WeChatIntegration:
|
|
"""
|
|
微信集成类
|
|
|
|
使用方法:
|
|
1. 初始化: wechat = WeChatIntegration()
|
|
2. 登录: await wechat.login()
|
|
3. 发送通知: await wechat.notify_task_complete(user_id, task_info)
|
|
4. 监听消息: await wechat.start_message_listener(callback)
|
|
"""
|
|
|
|
def __init__(self, config: Optional[WeChatConfig] = None):
|
|
self.config = config or WeChatConfig()
|
|
self.is_logged_in = False
|
|
self.message_cursor = 0
|
|
self.use_mock = False # 是否使用模拟模式
|
|
|
|
async def login(self, timeout: int = 300) -> bool:
|
|
"""
|
|
登录微信
|
|
|
|
流程:
|
|
1. 调用 login_qrcode 获取二维码
|
|
2. 显示二维码供用户扫描
|
|
3. 轮询 check_qrcode_status 等待登录
|
|
|
|
Args:
|
|
timeout: 超时时间（秒）
|
|
|
|
Returns:
|
|
bool: 是否登录成功
|
|
"""
|
|
print("🔐 正在启动微信登录...")
|
|
|
|
# 步骤1: 获取二维码
|
|
print("📱 正在生成登录二维码...")
|
|
qrcode_result = await self._call_mcp_tool("login_qrcode", {})
|
|
|
|
if not qrcode_result:
|
|
print("❌ 获取二维码失败")
|
|
return False
|
|
|
|
# 显示二维码
|
|
await self._display_qrcode()
|
|
|
|
# 步骤2: 等待扫码
|
|
print("⏳ 请使用微信扫描二维码...")
|
|
print(f" 二维码文本: {self.config.qrcode_path}")
|
|
print(f" 二维码图片: {self.config.qrcode_png_path}")
|
|
|
|
# 如果是模拟模式，直接返回成功
|
|
if self.use_mock:
|
|
print("✅ 模拟登录成功！")
|
|
self.is_logged_in = True
|
|
return True
|
|
|
|
start_time = time.time()
|
|
while time.time() - start_time < timeout:
|
|
status = await self._call_mcp_tool("check_qrcode_status", {})
|
|
|
|
if status.get("status") == "logged_in":
|
|
print("✅ 登录成功！")
|
|
self.is_logged_in = True
|
|
return True
|
|
elif status.get("status") == "expired":
|
|
print("⚠️ 二维码已过期，正在重新生成...")
|
|
await self._call_mcp_tool("login_qrcode", {})
|
|
await self._display_qrcode()
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
print("⏰ 登录超时")
|
|
return False
|
|
|
|
async def _display_qrcode(self):
|
|
"""显示二维码"""
|
|
# 尝试显示文本二维码
|
|
if self.config.qrcode_path.exists():
|
|
try:
|
|
qrcode_text = self.config.qrcode_path.read_text()
|
|
print("\n" + "="*40)
|
|
print("微信扫码登录:")
|
|
print(qrcode_text)
|
|
print("="*40 + "\n")
|
|
except Exception as e:
|
|
print(f" 无法读取文本二维码: {e}")
|
|
|
|
# 提示图片路径
|
|
if self.config.qrcode_png_path.exists():
|
|
print(f" 图片二维码: {self.config.qrcode_png_path}")
|
|
|
|
async def send_text_message(self, user_id: str, message: str) -> bool:
|
|
"""
|
|
发送文本消息
|
|
|
|
Args:
|
|
user_id: 用户微信ID
|
|
message: 消息内容
|
|
|
|
Returns:
|
|
bool: 是否发送成功
|
|
"""
|
|
if not self.is_logged_in:
|
|
print("❌ 未登录，请先调用 login()")
|
|
return False
|
|
|
|
result = await self._call_mcp_tool("send_text_message", {
|
|
"user_id": user_id,
|
|
"message": message
|
|
})
|
|
|
|
return result.get("success", False)
|
|
|
|
async def send_typing(self, user_id: str) -> bool:
|
|
"""
|
|
发送"正在输入..."状态
|
|
|
|
Args:
|
|
user_id: 用户微信ID
|
|
|
|
Returns:
|
|
bool: 是否发送成功
|
|
"""
|
|
result = await self._call_mcp_tool("send_typing", {
|
|
"user_id": user_id
|
|
})
|
|
return result.get("success", False)
|
|
|
|
async def notify_task_complete(self, user_id: str, task_info: Dict[str, Any]) -> bool:
|
|
"""
|
|
发送任务完成通知
|
|
|
|
Args:
|
|
user_id: 用户微信ID
|
|
task_info: 任务信息字典
|
|
- task_id: 任务ID
|
|
- task_type: 任务类型
|
|
- duration: 耗时
|
|
- status: 状态 (success/failed)
|
|
- output_files: 输出文件列表
|
|
- error: 错误信息（如果失败）
|
|
|
|
Returns:
|
|
bool: 是否发送成功
|
|
"""
|
|
if task_info.get("status") == "success":
|
|
template = self.config.task_complete_template
|
|
else:
|
|
template = self.config.task_failed_template
|
|
|
|
# 格式化输出文件列表
|
|
output_files = task_info.get("output_files", [])
|
|
if output_files:
|
|
files_str = "\n".join([f" • {f}" for f in output_files[:5]]) # 最多显示5个
|
|
if len(output_files) > 5:
|
|
files_str += f"\n ...还有 {len(output_files) - 5} 个文件"
|
|
else:
|
|
files_str = " (无输出文件)"
|
|
|
|
message = template.format(
|
|
task_id=task_info.get("task_id", "未知"),
|
|
task_type=task_info.get("task_type", "未知"),
|
|
duration=task_info.get("duration", "未知"),
|
|
status=task_info.get("status", "未知"),
|
|
output_files=files_str,
|
|
error=task_info.get("error", "")
|
|
)
|
|
|
|
return await self.send_text_message(user_id, message)
|
|
|
|
async def get_messages(self, timeout: int = 30) -> list:
|
|
"""
|
|
获取新消息（长轮询）
|
|
|
|
Args:
|
|
timeout: 轮询超时时间（秒）
|
|
|
|
Returns:
|
|
list: 新消息列表
|
|
"""
|
|
if not self.is_logged_in:
|
|
print("❌ 未登录，请先调用 login()")
|
|
return []
|
|
|
|
result = await self._call_mcp_tool("get_messages", {
|
|
"cursor": self.message_cursor,
|
|
"timeout": timeout
|
|
})
|
|
|
|
messages = result.get("messages", [])
|
|
if messages:
|
|
# 更新游标
|
|
self.message_cursor = result.get("next_cursor", self.message_cursor)
|
|
|
|
return messages
|
|
|
|
async def start_message_listener(self, callback: Callable[[Dict], None], poll_interval: int = 5):
|
|
"""
|
|
启动消息监听循环
|
|
|
|
Args:
|
|
callback: 消息处理回调函数
|
|
poll_interval: 轮询间隔（秒）
|
|
"""
|
|
print("🎧 启动消息监听...")
|
|
print(f" 轮询间隔: {poll_interval}秒")
|
|
print(" 按 Ctrl+C 停止\n")
|
|
|
|
try:
|
|
while True:
|
|
messages = await self.get_messages(timeout=poll_interval)
|
|
|
|
for msg in messages:
|
|
print(f"📨 收到消息: {msg}")
|
|
await callback(msg)
|
|
|
|
await asyncio.sleep(poll_interval)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n🛑 消息监听已停止")
|
|
|
|
async def _call_mcp_tool(self, tool_name: str, params: Dict) -> Dict:
|
|
"""
|
|
调用 MCP 工具
|
|
|
|
首先尝试使用真实的 mcp-wechat-server，如果失败则回退到模拟模式
|
|
"""
|
|
try:
|
|
# 尝试真实调用
|
|
return await self._real_mcp_call(tool_name, params)
|
|
except Exception as e:
|
|
# 真实调用失败，使用模拟模式
|
|
if not self.use_mock:
|
|
print(f" ⚠️ MCP 服务器不可用: {e}")
|
|
print(" 切换到模拟模式")
|
|
self.use_mock = True
|
|
return self._mock_call(tool_name, params)
|
|
|
|
async def _real_mcp_call(self, tool_name: str, params: Dict) -> Dict:
|
|
"""真实的 MCP 调用"""
|
|
import subprocess
|
|
import select
|
|
|
|
# 构建 MCP 请求 (JSON-RPC 格式)
|
|
mcp_request = {
|
|
"jsonrpc": "2.0",
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": tool_name,
|
|
"arguments": params
|
|
},
|
|
"id": 1
|
|
}
|
|
|
|
# 调用 mcp-wechat-server
|
|
proc = subprocess.Popen(
|
|
["bunx", "mcp-wechat-server"],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True
|
|
)
|
|
|
|
# 发送请求
|
|
request_json = json.dumps(mcp_request) + "\n"
|
|
proc.stdin.write(request_json)
|
|
proc.stdin.flush()
|
|
proc.stdin.close() # 关闭 stdin 表示请求发送完毕
|
|
|
|
# 读取响应（增加超时到 30 秒）
|
|
ready, _, _ = select.select([proc.stdout], [], [], 30)
|
|
|
|
if ready:
|
|
response_line = proc.stdout.readline()
|
|
proc.terminate()
|
|
|
|
try:
|
|
response = json.loads(response_line)
|
|
if "result" in response:
|
|
# 解析 MCP 响应格式
|
|
result = response["result"]
|
|
if "content" in result and len(result["content"]) > 0:
|
|
text_content = result["content"][0].get("text", "")
|
|
# 根据工具类型返回不同格式
|
|
if tool_name == "login_qrcode":
|
|
return {"success": True, "message": text_content}
|
|
elif tool_name == "check_qrcode_status":
|
|
# 检查是否已登录
|
|
if "logged in" in text_content.lower() or "success" in text_content.lower():
|
|
return {"status": "logged_in", "message": text_content}
|
|
else:
|
|
return {"status": "waiting", "message": text_content}
|
|
elif tool_name == "send_text_message":
|
|
return {"success": "success" in text_content.lower(), "message": text_content}
|
|
else:
|
|
return {"success": True, "message": text_content}
|
|
return result
|
|
elif "error" in response:
|
|
return {"success": False, "error": response["error"]}
|
|
except json.JSONDecodeError:
|
|
return {"success": False, "error": "Invalid JSON response"}
|
|
else:
|
|
proc.terminate()
|
|
raise TimeoutError("Timeout waiting for MCP response")
|
|
|
|
def _mock_call(self, tool_name: str, params: Dict) -> Dict:
|
|
"""模拟 MCP 调用（用于测试）"""
|
|
if tool_name == "login_qrcode":
|
|
# 创建模拟的二维码文件
|
|
qrcode_dir = Path.home() / ".mcp-wechat-server"
|
|
qrcode_dir.mkdir(exist_ok=True)
|
|
|
|
# 模拟二维码文本
|
|
qrcode_text = "[模拟二维码] 请想象这里有一个微信登录二维码"
|
|
(qrcode_dir / "qrcode.txt").write_text(qrcode_text)
|
|
|
|
return {"success": True, "qrcode_url": "weixin://login/mock"}
|
|
|
|
elif tool_name == "check_qrcode_status":
|
|
# 模拟登录成功
|
|
return {"status": "logged_in"}
|
|
|
|
elif tool_name == "send_text_message":
|
|
user_id = params.get('user_id', 'unknown')
|
|
message = params.get('message', '')[:50]
|
|
print(f" [模拟发送消息给 {user_id}]: {message}...")
|
|
return {"success": True}
|
|
|
|
elif tool_name == "send_typing":
|
|
return {"success": True}
|
|
|
|
elif tool_name == "get_messages":
|
|
# 模拟返回消息
|
|
return {"messages": [], "next_cursor": self.message_cursor}
|
|
|
|
return {"success": False, "error": "Unknown tool"}
|
|
|
|
|
|
class WeChatTaskNotifier:
|
|
"""
|
|
任务通知器 - 与 bbav20311 任务系统集成
|
|
|
|
使用示例:
|
|
notifier = WeChatTaskNotifier(user_wechat_id="wxid_xxx")
|
|
await notifier.notify_task_started(task_id="task_001", task_type="identify")
|
|
# ... 任务执行 ...
|
|
await notifier.notify_task_completed(task_id="task_001", result=...)
|
|
"""
|
|
|
|
def __init__(self, user_wechat_id: str, wechat: Optional[WeChatIntegration] = None):
|
|
self.user_wechat_id = user_wechat_id
|
|
self.wechat = wechat or WeChatIntegration()
|
|
self.task_start_times: Dict[str, datetime] = {}
|
|
|
|
async def notify_task_started(self, task_id: str, task_type: str, **kwargs):
|
|
"""通知任务开始"""
|
|
self.task_start_times[task_id] = datetime.now()
|
|
|
|
message = f"""🚀 任务已启动
|
|
|
|
📋 任务ID: {task_id}
|
|
🔬 分析类型: {task_type}
|
|
⏱️ 开始时间: {self.task_start_times[task_id].strftime('%Y-%m-%d %H:%M:%S')}
|
|
|
|
正在处理中，完成后将通知您..."""
|
|
|
|
await self.wechat.send_text_message(self.user_wechat_id, message)
|
|
|
|
async def notify_task_completed(self, task_id: str, output_files: list, **kwargs):
|
|
"""通知任务完成"""
|
|
start_time = self.task_start_times.get(task_id)
|
|
duration = "未知"
|
|
if start_time:
|
|
delta = datetime.now() - start_time
|
|
duration = f"{delta.total_seconds() // 60:.0f}分{delta.total_seconds() % 60:.0f}秒"
|
|
|
|
task_info = {
|
|
"task_id": task_id,
|
|
"task_type": kwargs.get("task_type", "未知"),
|
|
"duration": duration,
|
|
"status": "success",
|
|
"output_files": output_files
|
|
}
|
|
|
|
await self.wechat.notify_task_complete(self.user_wechat_id, task_info)
|
|
|
|
async def notify_task_failed(self, task_id: str, error: str, **kwargs):
|
|
"""通知任务失败"""
|
|
start_time = self.task_start_times.get(task_id)
|
|
duration = "未知"
|
|
if start_time:
|
|
delta = datetime.now() - start_time
|
|
duration = f"{delta.total_seconds() // 60:.0f}分{delta.total_seconds() % 60:.0f}秒"
|
|
|
|
task_info = {
|
|
"task_id": task_id,
|
|
"task_type": kwargs.get("task_type", "未知"),
|
|
"duration": duration,
|
|
"status": "failed",
|
|
"output_files": [],
|
|
"error": error
|
|
}
|
|
|
|
await self.wechat.notify_task_complete(self.user_wechat_id, task_info)
|
|
|
|
|
|
# ==================== 测试代码 ====================
|
|
|
|
async def test_wechat_integration():
|
|
"""测试微信集成"""
|
|
print("="*60)
|
|
print("🧪 bbav20311 微信集成测试")
|
|
print("="*60)
|
|
|
|
# 1. 初始化
|
|
print("\n📦 步骤1: 初始化微信集成模块")
|
|
wechat = WeChatIntegration()
|
|
|
|
# 2. 登录
|
|
print("\n🔐 步骤2: 登录测试")
|
|
success = await wechat.login(timeout=10)
|
|
if not success:
|
|
print("❌ 登录失败，跳过后续测试")
|
|
return
|
|
|
|
# 3. 发送测试消息
|
|
print("\n💬 步骤3: 发送测试消息")
|
|
test_user_id = "test_user_001"
|
|
test_message = "你好！这是来自 bbav20311 的测试消息 🎉"
|
|
|
|
success = await wechat.send_text_message(test_user_id, test_message)
|
|
print(f" 发送结果: {'✅ 成功' if success else '❌ 失败'}")
|
|
|
|
# 4. 发送任务通知
|
|
print("\n📊 步骤4: 发送任务完成通知")
|
|
task_info = {
|
|
"task_id": "task_20260410_001",
|
|
"task_type": "基因家族鉴定 (identify)",
|
|
"duration": "5分32秒",
|
|
"status": "success",
|
|
"output_files": [
|
|
"Ft.SPL.identify.gene.list.txt",
|
|
"Ft.SPL.identify.protein.fa",
|
|
"Ft.SPL.identify.report.pdf"
|
|
]
|
|
}
|
|
|
|
success = await wechat.notify_task_complete(test_user_id, task_info)
|
|
print(f" 通知结果: {'✅ 成功' if success else '❌ 失败'}")
|
|
|
|
# 5. 发送失败通知
|
|
print("\n❌ 步骤5: 发送任务失败通知")
|
|
failed_task_info = {
|
|
"task_id": "task_20260410_002",
|
|
"task_type": "进化树分析 (phylogeny)",
|
|
"duration": "2分15秒",
|
|
"status": "failed",
|
|
"output_files": [],
|
|
"error": "输入文件格式错误: 缺少必需的序列数据"
|
|
}
|
|
|
|
success = await wechat.notify_task_complete(test_user_id, failed_task_info)
|
|
print(f" 通知结果: {'✅ 成功' if success else '❌ 失败'}")
|
|
|
|
# 6. 测试任务通知器
|
|
print("\n🤖 步骤6: 测试任务通知器集成")
|
|
notifier = WeChatTaskNotifier(user_wechat_id=test_user_id, wechat=wechat)
|
|
|
|
task_id = "task_20260410_003"
|
|
await notifier.notify_task_started(task_id=task_id, task_type="Motif分析")
|
|
|
|
# 模拟任务执行时间
|
|
print(" 模拟任务执行中...")
|
|
await asyncio.sleep(2)
|
|
|
|
await notifier.notify_task_completed(
|
|
task_id=task_id,
|
|
task_type="Motif分析",
|
|
output_files=[
|
|
"Ft.SPL.motif.meme.txt",
|
|
"Ft.SPL.motif.logo.png"
|
|
]
|
|
)
|
|
|
|
print("\n" + "="*60)
|
|
print("✅ 所有测试完成！")
|
|
print("="*60)
|
|
|
|
|
|
async def test_message_listener():
|
|
"""测试消息监听"""
|
|
print("\n🎧 测试消息监听（按 Ctrl+C 停止）")
|
|
|
|
wechat = WeChatIntegration()
|
|
await wechat.login(timeout=10)
|
|
|
|
async def message_handler(msg: Dict):
|
|
"""处理收到的消息"""
|
|
print(f" 处理消息: {msg}")
|
|
|
|
# 示例：根据消息内容执行不同操作
|
|
content = msg.get("content", "").lower()
|
|
|
|
if "状态" in content or "status" in content:
|
|
# 查询任务状态
|
|
reply = "📊 当前系统状态正常，所有服务运行中"
|
|
await wechat.send_text_message(msg.get("from"), reply)
|
|
|
|
elif "帮助" in content or "help" in content:
|
|
reply = """🤖 bbav20311 基因家族分析助手
|
|
|
|
可用命令:
|
|
• 状态 - 查询系统状态
|
|
• 帮助 - 显示此帮助信息
|
|
• 任务列表 - 查看进行中的任务
|
|
• 详情 <任务ID> - 查看任务详情
|
|
|
|
直接发送分析需求即可开始任务！"""
|
|
await wechat.send_text_message(msg.get("from"), reply)
|
|
|
|
await wechat.start_message_listener(callback=message_handler, poll_interval=5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) > 1 and sys.argv[1] == "listen":
|
|
# 运行消息监听测试
|
|
asyncio.run(test_message_listener())
|
|
else:
|
|
# 运行基础功能测试
|
|
asyncio.run(test_wechat_integration()) |