feat: 企业微信机器人支持接收图片消息自动OCR识别入库

- 注册 message.image 事件监听
- 使用 SDK download_file() 下载并解密图片
- 调用后端 recognize-direct 接口进行 OCR 识别(自动去重)
- 根据识别结果回复:成功录入/重复内容/识别失败
- 更新欢迎语提示支持图片发送功能
This commit is contained in:
EduBrain Dev
2026-04-14 13:28:36 +08:00
parent b17786b57b
commit 8ebf9a4587

View File

@@ -16,6 +16,9 @@ import logging
import math
import os
import re
import tempfile
import time
from datetime import datetime
import httpx
from aibot import WSClient, WSClientOptions, generate_req_id
@@ -63,6 +66,7 @@ class WeComBotService:
# 注册事件
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)
@@ -83,7 +87,7 @@ class WeComBotService:
{
"msgtype": "text",
"text": {
"content": "你好!我是知识库图片搜索助手,发送关键词即可搜索相关图片"
"content": "你好!我是知识库助手,支持以下功能:\n\n1. 发送关键词搜索相关图片\n2. 发送图片 → 自动识别并录入知识库"
},
},
)
@@ -194,6 +198,102 @@ class WeComBotService:
"""连接断开回调"""
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.edubrain_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", "?")
await self.ws_client.reply_stream(
frame,
stream_id,
f"该图片内容已存在于知识库中ID: {dup_id}),无需重复录入。",
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)