b17786b57b
主要功能: - 图片上传时 OCR 内容去重（3个上传端点统一使用公共函数 _check_ocr_duplicate） - 图片管理 Tab：展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮（带确认弹窗） - 图片管理卡片点击查看详情（复用 showOcrDetailModal） - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人：WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
363 lines
12 KiB
Python
363 lines
12 KiB
Python
"""
|
||
企业微信智能机器人服务（WebSocket 长连接模式）
|
||
|
||
使用官方 SDK (wecom-aibot-python-sdk) 的 WSClient 类，
|
||
封装了连接管理、心跳、断线重连等能力。
|
||
|
||
作为独立进程运行：python -m app.wework_bot
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import math
|
||
import os
|
||
import re
|
||
|
||
import httpx
|
||
from aibot import WSClient, WSClientOptions, generate_req_id
|
||
from dotenv import load_dotenv
|
||
|
||
# 加载 .env 文件（独立运行时需要）
|
||
load_dotenv()
|
||
|
||
logger = logging.getLogger("wework_bot")
|
||
|
||
# 分片上传参数（参考官方文档）
|
||
_CHUNK_SIZE = 512 * 1024 # 单个分片最大 512KB（Base64 编码前）
|
||
|
||
|
||
class WeComBotService:
|
||
"""企业微信智能机器人服务（WebSocket 长连接模式）"""
|
||
|
||
def __init__(self):
|
||
self.bot_id = os.getenv("WEWORK_BOT_ID", "")
|
||
self.bot_secret = os.getenv("WEWORK_BOT_SECRET", "")
|
||
self.edubrain_base_url = os.getenv(
|
||
"EDUBRAIN_BASE_URL", "http://localhost:8765"
|
||
)
|
||
self.search_limit = int(os.getenv("WEWORK_BOT_SEARCH_LIMIT", "3"))
|
||
self.enabled = os.getenv("WEWORK_BOT_ENABLED", "false").lower() == "true"
|
||
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.ws_client.on("authenticated", self._on_authenticated)
|
||
self.ws_client.on("message.text", self._on_text_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)
|
||
|
||
logger.info("启动企业微信智能机器人...")
|
||
self.ws_client.run()
|
||
|
||
# ──────────────────────────── 事件处理 ────────────────────────────
|
||
|
||
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": "你好！我是知识库图片搜索助手，发送关键词即可搜索相关图片。"
|
||
},
|
||
},
|
||
)
|
||
|
||
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,
|
||
)
|
||
|
||
# 调用 EduBrain 搜索
|
||
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)
|
||
|
||
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.edubrain_base_url}/api/v1/settings")
|
||
if resp.status_code == 200:
|
||
data = resp.json()
|
||
return data.get("search_limit", self.search_limit)
|
||
except Exception as e:
|
||
logger.warning("获取 search_limit 失败，使用默认值: %s", e)
|
||
return self.search_limit
|
||
|
||
async def _search_images(self, keyword: str) -> list:
|
||
"""
|
||
调用 EduBrain 搜索接口（SSE 流式），收集所有匹配结果。
|
||
|
||
Args:
|
||
keyword: 搜索关键词
|
||
|
||
Returns:
|
||
匹配结果列表
|
||
"""
|
||
results = []
|
||
url = f"{self.edubrain_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()
|