feat: 企业微信机器人支持接收图片消息自动OCR识别入库
- 注册 message.image 事件监听 - 使用 SDK download_file() 下载并解密图片 - 调用后端 recognize-direct 接口进行 OCR 识别(自动去重) - 根据识别结果回复:成功录入/重复内容/识别失败 - 更新欢迎语提示支持图片发送功能
This commit is contained in:
@@ -16,6 +16,9 @@ import logging
|
|||||||
import math
|
import math
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
from aibot import WSClient, WSClientOptions, generate_req_id
|
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("authenticated", self._on_authenticated)
|
||||||
self.ws_client.on("message.text", self._on_text_message)
|
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("event.enter_chat", self._on_enter_chat)
|
||||||
self.ws_client.on("disconnected", self._on_disconnected)
|
self.ws_client.on("disconnected", self._on_disconnected)
|
||||||
self.ws_client.on("error", self._on_error)
|
self.ws_client.on("error", self._on_error)
|
||||||
@@ -83,7 +87,7 @@ class WeComBotService:
|
|||||||
{
|
{
|
||||||
"msgtype": "text",
|
"msgtype": "text",
|
||||||
"text": {
|
"text": {
|
||||||
"content": "你好!我是知识库图片搜索助手,发送关键词即可搜索相关图片。"
|
"content": "你好!我是知识库助手,支持以下功能:\n\n1. 发送关键词 → 搜索相关图片\n2. 发送图片 → 自动识别并录入知识库"
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
@@ -194,6 +198,102 @@ class WeComBotService:
|
|||||||
"""连接断开回调"""
|
"""连接断开回调"""
|
||||||
logger.warning("企业微信连接断开: %s", reason)
|
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):
|
def _on_error(self, error):
|
||||||
"""错误回调"""
|
"""错误回调"""
|
||||||
logger.error("企业微信机器人错误: %s", error)
|
logger.error("企业微信机器人错误: %s", error)
|
||||||
|
|||||||
Reference in New Issue
Block a user