feat: v1.2.0 - 图片去重与管理、微信机器人优化、搜索设置可配置
主要功能: - 图片上传时 OCR 内容去重（3个上传端点统一使用公共函数 _check_ocr_duplicate） - 图片管理 Tab：展示所有图片、手动删除、一键去重 - 搜索结果详情弹窗增加删除按钮（带确认弹窗） - 图片管理卡片点击查看详情（复用 showOcrDetailModal） - 搜索限制和 LLM 批量判断数量可通过网站设置 - MiniMax API 调用添加 reasoning_split=True - 企业微信机器人：WebSocket 长连接、图片搜索、配置化搜索数量 - 版本号升级至 1.2.0
This commit is contained in:
EduBrain Dev
362
app/wework_bot.py
Normal file
362
app/wework_bot.py
Normal file
@@ -0,0 +1,362 @@
|
||||
"""
|
||||
企业微信智能机器人服务（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()
|
||||
Reference in New Issue
Block a user