4c6a20e5fc
全局替换 EduBrain -> HuiBrain, edu-brain -> huibrain, edu_brain -> hui_brain, EDUBRAIN -> HUIBRAIN 涉及文件：README.md, pyproject.toml, docker-compose.yml, .env, .env.example, app/config.py, app/main.py, app/wework_bot.py, app/__init__.py, app/mcp/server.py, static/index.html, docs/IMAGE_API_GUIDE.md
494 lines
18 KiB
Python
494 lines
18 KiB
Python
"""
|
||
企业微信智能机器人服务（WebSocket 长连接模式）
|
||
|
||
使用官方 SDK (wecom-aibot-python-sdk) 的 WSClient 类，
|
||
封装了连接管理、心跳、断线重连等能力。
|
||
|
||
支持两种运行方式：
|
||
1. 集成模式：后端启动时自动拉起（BotManager）
|
||
2. 独立进程：python -m app.wework_bot
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import math
|
||
import os
|
||
import re
|
||
import tempfile
|
||
import time
|
||
import asyncio
|
||
from datetime import datetime
|
||
|
||
import httpx
|
||
from aibot import WSClient, WSClientOptions, generate_req_id
|
||
from dotenv import load_dotenv
|
||
|
||
# 加载 .env 文件（独立运行时需要）
|
||
load_dotenv()
|
||
|
||
from app.config import settings
|
||
|
||
logger = logging.getLogger("wework_bot")
|
||
|
||
# 分片上传参数（参考官方文档）
|
||
_CHUNK_SIZE = 512 * 1024 # 单个分片最大 512KB（Base64 编码前）
|
||
|
||
|
||
class WeComBotService:
|
||
"""企业微信智能机器人服务（WebSocket 长连接模式）"""
|
||
|
||
def __init__(self):
|
||
self.bot_id = settings.WEWORK_BOT_ID or ""
|
||
self.bot_secret = settings.WEWORK_BOT_SECRET or ""
|
||
self.huibrain_base_url = settings.HUIBRAIN_BASE_URL
|
||
self.enabled = settings.WEWORK_BOT_ENABLED
|
||
self.ws_client: WSClient | None = None
|
||
|
||
def start(self):
|
||
"""启动机器人服务（阻塞模式，用于独立进程运行）"""
|
||
if not self.enabled:
|
||
logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)")
|
||
return
|
||
|
||
if not self.bot_id or not self.bot_secret:
|
||
logger.error("WEWORK_BOT_ID 或 WEWORK_BOT_SECRET 未配置")
|
||
return
|
||
|
||
self.ws_client = WSClient(
|
||
WSClientOptions(
|
||
bot_id=self.bot_id,
|
||
secret=self.bot_secret,
|
||
)
|
||
)
|
||
|
||
self._register_events()
|
||
logger.info("启动企业微信智能机器人...")
|
||
self.ws_client.run()
|
||
|
||
async def start_async(self):
|
||
"""启动机器人服务（异步模式，用于嵌入到其他事件循环）"""
|
||
if not self.enabled:
|
||
logger.info("企业微信机器人未启用 (WEWORK_BOT_ENABLED=false)")
|
||
return
|
||
|
||
if not self.bot_id or not self.bot_secret:
|
||
logger.error("WEWORK_BOT_ID 或 WEWORK_BOT_SECRET 未配置")
|
||
return
|
||
|
||
self.ws_client = WSClient(
|
||
WSClientOptions(
|
||
bot_id=self.bot_id,
|
||
secret=self.bot_secret,
|
||
)
|
||
)
|
||
|
||
self._register_events()
|
||
logger.info("启动企业微信智能机器人（异步模式）...")
|
||
await self.ws_client.connect()
|
||
# 永远等待，保持连接
|
||
await asyncio.Future()
|
||
|
||
def _register_events(self):
|
||
"""注册所有事件处理器"""
|
||
self.ws_client.on("authenticated", self._on_authenticated)
|
||
self.ws_client.on("message.text", self._on_text_message)
|
||
self.ws_client.on("message.image", self._on_image_message)
|
||
self.ws_client.on("event.enter_chat", self._on_enter_chat)
|
||
self.ws_client.on("disconnected", self._on_disconnected)
|
||
self.ws_client.on("error", self._on_error)
|
||
|
||
# ──────────────────────────── 事件处理 ────────────────────────────
|
||
|
||
def _on_authenticated(self):
|
||
"""认证成功回调"""
|
||
logger.info("企业微信机器人认证成功")
|
||
|
||
async def _on_enter_chat(self, frame: dict):
|
||
"""进入会话事件 -> 回复欢迎语"""
|
||
await self.ws_client.reply_welcome(
|
||
frame,
|
||
{
|
||
"msgtype": "text",
|
||
"text": {
|
||
"content": "你好！我是知识库助手，支持以下功能：\n\n1. 发送关键词 → 搜索相关图片\n2. 发送图片 → 自动识别并录入知识库"
|
||
},
|
||
},
|
||
)
|
||
|
||
async def _on_text_message(self, frame: dict):
|
||
"""收到文本消息 -> 搜索图片 -> 流式回复文本 + 逐张发送图片"""
|
||
body = frame.get("body", {})
|
||
text_content = body.get("text", {}).get("content", "")
|
||
|
||
# 去掉 @机器人 的前缀
|
||
text_content = re.sub(r"@.*?\s*", "", text_content).strip()
|
||
|
||
if not text_content:
|
||
stream_id = generate_req_id("stream")
|
||
await self.ws_client.reply_stream(
|
||
frame, stream_id, "请输入搜索关键词", True
|
||
)
|
||
return
|
||
|
||
# 清理口语化前缀，提取关键词
|
||
keyword = re.sub(
|
||
r"(帮我|请|找|搜|搜索|查|查一下|来|给我|要|想要|的?图|图片|相关|关于|一下)",
|
||
"", text_content,
|
||
).strip()
|
||
if not keyword:
|
||
keyword = text_content # 兜底
|
||
|
||
stream_id = generate_req_id("stream")
|
||
|
||
# 立即回复：开始搜索
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"\U0001f50d 正在搜索「{keyword}」...",
|
||
False,
|
||
)
|
||
|
||
# 调用 HuiBrain 搜索
|
||
try:
|
||
results = await self._search_images(keyword)
|
||
|
||
if not results:
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"未找到与「{text_content}」相关的图片",
|
||
True,
|
||
)
|
||
return
|
||
|
||
# 构建文本回复内容
|
||
content_parts = [f"找到 {len(results)} 张相关图片："]
|
||
for i, r in enumerate(results, 1):
|
||
preview = r.get("ocr_text_preview", "")[:80]
|
||
content_parts.append(f"\n{i}. {preview}...")
|
||
|
||
content = "\n".join(content_parts)
|
||
|
||
# 完成流式消息（仅文本，不使用 msg_item）
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
content,
|
||
True,
|
||
)
|
||
|
||
# 逐张上传临时素材并发送图片消息
|
||
for i, r in enumerate(results, 1):
|
||
b64 = r.get("image_base64", "")
|
||
if not b64:
|
||
continue
|
||
|
||
try:
|
||
# 解码 base64 获取原始图片数据
|
||
image_data = base64.b64decode(b64)
|
||
file_name = r.get("image_url", f"image_{i}.jpg")
|
||
|
||
# 上传临时素材（分片上传）
|
||
media_id = await self._upload_temp_media(
|
||
image_data, file_name
|
||
)
|
||
if not media_id:
|
||
logger.warning("第 %d 张图片上传失败，跳过", i)
|
||
continue
|
||
|
||
# 发送图片消息（通过 aibot_respond_msg）
|
||
await self.ws_client.reply(
|
||
frame,
|
||
{
|
||
"msgtype": "image",
|
||
"image": {"media_id": media_id},
|
||
},
|
||
)
|
||
logger.info("已发送第 %d 张图片 (media_id: %s)", i, media_id)
|
||
except Exception as e:
|
||
logger.error("发送第 %d 张图片失败: %s", i, e, exc_info=True)
|
||
|
||
except Exception as e:
|
||
logger.error("搜索失败: %s", e, exc_info=True)
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"搜索出错：{e}",
|
||
True,
|
||
)
|
||
|
||
def _on_disconnected(self, reason: str):
|
||
"""连接断开回调"""
|
||
logger.warning("企业微信连接断开: %s", reason)
|
||
|
||
async def _on_image_message(self, frame: dict):
|
||
"""收到图片消息 -> 下载解密 -> 调用OCR识别 -> 回复结果"""
|
||
body = frame.get("body", {})
|
||
image_info = body.get("image", {})
|
||
image_url = image_info.get("url", "")
|
||
aes_key = image_info.get("aeskey", "")
|
||
|
||
if not image_url:
|
||
logger.warning("收到图片消息但无下载地址")
|
||
return
|
||
|
||
stream_id = generate_req_id("stream")
|
||
await self.ws_client.reply_stream(
|
||
frame, stream_id, "收到图片，正在识别...", False
|
||
)
|
||
|
||
try:
|
||
# 1. 下载并解密图片
|
||
image_data, original_filename = await self.ws_client.download_file(
|
||
image_url, aes_key
|
||
)
|
||
if not image_data:
|
||
await self.ws_client.reply_stream(
|
||
frame, stream_id, "图片下载失败，请重试", True
|
||
)
|
||
return
|
||
|
||
logger.info(
|
||
"图片下载成功: filename=%s, size=%d bytes",
|
||
original_filename, len(image_data),
|
||
)
|
||
|
||
# 2. 生成保存文件名（避免重名）
|
||
ext = os.path.splitext(original_filename or "image.jpg")[1] or ".jpg"
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
save_filename = f"wecom_{timestamp}{ext}"
|
||
|
||
# 3. 调用后端 recognize-direct 接口
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
resp = await client.post(
|
||
f"{self.huibrain_base_url}/api/v1/images/recognize-direct",
|
||
files={"file": (save_filename, image_data, "image/jpeg")},
|
||
)
|
||
|
||
if resp.status_code != 200:
|
||
logger.error("OCR 识别失败: status=%d, body=%s", resp.status_code, resp.text)
|
||
await self.ws_client.reply_stream(
|
||
frame, stream_id, "图片识别失败，请稍后重试", True
|
||
)
|
||
return
|
||
|
||
result = resp.json()
|
||
|
||
# 4. 根据识别结果回复
|
||
status = result.get("status", "")
|
||
|
||
if status == "duplicate":
|
||
dup_id = result.get("duplicate_of", "?")
|
||
ocr_text = result.get("ocr_text", "")
|
||
preview = ocr_text[:200].replace("\n", " ") if ocr_text else "无识别内容"
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"🔄 该图片内容已存在于知识库中（ID: {dup_id}），无需重复录入。\n\n识别内容：\n{preview}{'...' if len(ocr_text) > 200 else ''}",
|
||
True,
|
||
)
|
||
elif status == "completed":
|
||
ocr_text = result.get("ocr_text", "")
|
||
preview = ocr_text[:200].replace("\n", " ") if ocr_text else "无识别内容"
|
||
record_id = result.get("id", "?")
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"✅ 图片已录入知识库（ID: {record_id}）\n\n识别内容预览：\n{preview}{'...' if len(ocr_text) > 200 else ''}",
|
||
True,
|
||
)
|
||
elif status == "failed":
|
||
error_msg = result.get("message", result.get("detail", "未知错误"))
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"图片识别失败：{error_msg}",
|
||
True,
|
||
)
|
||
else:
|
||
await self.ws_client.reply_stream(
|
||
frame,
|
||
stream_id,
|
||
f"图片处理结果：{json.dumps(result, ensure_ascii=False)}",
|
||
True,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error("处理图片消息失败: %s", e, exc_info=True)
|
||
await self.ws_client.reply_stream(
|
||
frame, stream_id, f"处理图片时出错：{e}", True
|
||
)
|
||
|
||
def _on_error(self, error):
|
||
"""错误回调"""
|
||
logger.error("企业微信机器人错误: %s", error)
|
||
|
||
# ──────────────────────── 上传临时素材 ────────────────────────
|
||
|
||
async def _upload_temp_media(
|
||
self, file_data: bytes, filename: str
|
||
) -> str | None:
|
||
"""
|
||
通过 WebSocket 通道上传临时素材（三步分片上传）。
|
||
|
||
流程：aibot_upload_media_init → aibot_upload_media_chunk × N → aibot_upload_media_finish
|
||
|
||
Args:
|
||
file_data: 文件原始二进制数据
|
||
filename: 文件名
|
||
|
||
Returns:
|
||
成功返回 media_id，失败返回 None
|
||
"""
|
||
total_size = len(file_data)
|
||
total_chunks = math.ceil(total_size / _CHUNK_SIZE)
|
||
md5 = hashlib.md5(file_data).hexdigest()
|
||
ws = self.ws_client._ws_manager
|
||
|
||
# 第一步：上传初始化
|
||
init_req_id = generate_req_id("upload")
|
||
try:
|
||
init_result = await ws.send_reply(
|
||
init_req_id,
|
||
{
|
||
"type": "image",
|
||
"filename": filename,
|
||
"total_size": total_size,
|
||
"total_chunks": total_chunks,
|
||
"md5": md5,
|
||
},
|
||
cmd="aibot_upload_media_init",
|
||
)
|
||
except Exception as e:
|
||
logger.error("上传初始化失败: %s", e)
|
||
return None
|
||
|
||
upload_id = init_result.get("body", {}).get("upload_id", "")
|
||
if not upload_id:
|
||
logger.error("上传初始化未返回 upload_id: %s", init_result)
|
||
return None
|
||
|
||
logger.info(
|
||
"上传初始化成功: upload_id=%s, total_size=%d, chunks=%d",
|
||
upload_id, total_size, total_chunks,
|
||
)
|
||
|
||
# 第二步：分片上传
|
||
for chunk_idx in range(total_chunks):
|
||
start = chunk_idx * _CHUNK_SIZE
|
||
end = min(start + _CHUNK_SIZE, total_size)
|
||
chunk_data = file_data[start:end]
|
||
chunk_b64 = base64.b64encode(chunk_data).decode("utf-8")
|
||
|
||
chunk_req_id = generate_req_id("upload_chunk")
|
||
try:
|
||
await ws.send_reply(
|
||
chunk_req_id,
|
||
{
|
||
"upload_id": upload_id,
|
||
"chunk_index": chunk_idx, # 分片索引从 0 开始
|
||
"base64_data": chunk_b64,
|
||
},
|
||
cmd="aibot_upload_media_chunk",
|
||
)
|
||
except Exception as e:
|
||
logger.error(
|
||
"分片上传失败 (chunk %d/%d): %s",
|
||
chunk_idx + 1, total_chunks, e,
|
||
)
|
||
return None
|
||
|
||
logger.info("所有分片上传完成: upload_id=%s", upload_id)
|
||
|
||
# 第三步：完成上传
|
||
finish_req_id = generate_req_id("upload_finish")
|
||
try:
|
||
finish_result = await ws.send_reply(
|
||
finish_req_id,
|
||
{"upload_id": upload_id},
|
||
cmd="aibot_upload_media_finish",
|
||
)
|
||
except Exception as e:
|
||
logger.error("完成上传失败: %s", e)
|
||
return None
|
||
|
||
media_id = finish_result.get("body", {}).get("media_id", "")
|
||
if not media_id:
|
||
logger.error("完成上传未返回 media_id: %s", finish_result)
|
||
return None
|
||
|
||
logger.info("上传临时素材成功: media_id=%s", media_id)
|
||
return media_id
|
||
|
||
# ──────────────────────────── 搜索调用 ────────────────────────────
|
||
|
||
async def _get_search_limit(self) -> int:
|
||
"""从后端 API 实时获取当前 search_limit 配置"""
|
||
try:
|
||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||
resp = await client.get(f"{self.huibrain_base_url}/api/v1/settings")
|
||
if resp.status_code == 200:
|
||
data = resp.json()
|
||
limit = data.get("search_limit")
|
||
if limit:
|
||
return limit
|
||
except Exception as e:
|
||
logger.warning("获取 search_limit 失败，使用默认值: %s", e)
|
||
return settings.SEARCH_LIMIT
|
||
|
||
async def _search_images(self, keyword: str) -> list:
|
||
"""
|
||
调用 HuiBrain 搜索接口（SSE 流式），收集所有匹配结果。
|
||
|
||
Args:
|
||
keyword: 搜索关键词
|
||
|
||
Returns:
|
||
匹配结果列表
|
||
"""
|
||
results = []
|
||
url = f"{self.huibrain_base_url}/api/v1/images/search"
|
||
limit = await self._get_search_limit()
|
||
|
||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||
async with client.stream(
|
||
"GET",
|
||
url,
|
||
params={
|
||
"keyword": keyword,
|
||
"limit": limit,
|
||
"sort": "time_desc",
|
||
},
|
||
) as resp:
|
||
async for line in resp.aiter_lines():
|
||
if not line.startswith("data: "):
|
||
continue
|
||
try:
|
||
data = json.loads(line[6:])
|
||
if data.get("type") == "result":
|
||
results.append(data)
|
||
elif data.get("type") == "done":
|
||
break
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
return results
|
||
|
||
|
||
def run_wework_bot():
|
||
"""入口函数"""
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
|
||
)
|
||
service = WeComBotService()
|
||
service.start()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_wework_bot()
|