feat: 企业微信机器人支持接收图片消息自动OCR识别入库
- 注册 message.image 事件监听 - 使用 SDK download_file() 下载并解密图片 - 调用后端 recognize-direct 接口进行 OCR 识别(自动去重) - 根据识别结果回复:成功录入/重复内容/识别失败 - 更新欢迎语提示支持图片发送功能
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user