97 lines
2.8 KiB
Python
97 lines
2.8 KiB
Python
"""企业微信 API 模式：上传图片并回复。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from aibot import generate_req_id
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
CHUNK_SIZE = 512 * 1024
|
||
MAX_IMAGE_BYTES = 9 * 1024 * 1024
|
||
|
||
|
||
def _ensure_image_size(path: Path) -> bytes:
|
||
data = path.read_bytes()
|
||
if len(data) > MAX_IMAGE_BYTES:
|
||
raise RuntimeError(
|
||
f"截图过大（{len(data) // 1024}KB），请缩小页面或使用 viewport 截图（上限 9MB）"
|
||
)
|
||
return data
|
||
|
||
|
||
def _response_body(frame: dict[str, Any]) -> dict[str, Any]:
|
||
if frame.get("errcode", 0) != 0:
|
||
raise RuntimeError(
|
||
f"企微接口错误 errcode={frame.get('errcode')} errmsg={frame.get('errmsg')}"
|
||
)
|
||
body = frame.get("body")
|
||
return body if isinstance(body, dict) else {}
|
||
|
||
|
||
async def upload_image(ws_client: Any, image_path: str | Path) -> str:
|
||
path = Path(image_path)
|
||
if not path.exists():
|
||
raise RuntimeError(f"截图不存在: {path}")
|
||
|
||
data = _ensure_image_size(path)
|
||
md5 = hashlib.md5(data).hexdigest()
|
||
chunks = [data[i : i + CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
|
||
total_chunks = len(chunks)
|
||
|
||
manager = ws_client._ws_manager
|
||
|
||
init_frame = await manager.send_reply(
|
||
generate_req_id("upload_init"),
|
||
{
|
||
"type": "image",
|
||
"filename": path.name,
|
||
"total_size": len(data),
|
||
"total_chunks": total_chunks,
|
||
"md5": md5,
|
||
},
|
||
"aibot_upload_media_init",
|
||
)
|
||
upload_id = _response_body(init_frame).get("upload_id")
|
||
if not upload_id:
|
||
raise RuntimeError("上传初始化失败：未返回 upload_id")
|
||
|
||
for index, chunk in enumerate(chunks):
|
||
chunk_frame = await manager.send_reply(
|
||
generate_req_id("upload_chunk"),
|
||
{
|
||
"upload_id": upload_id,
|
||
"chunk_index": index,
|
||
"base64_data": base64.b64encode(chunk).decode("ascii"),
|
||
},
|
||
"aibot_upload_media_chunk",
|
||
)
|
||
_response_body(chunk_frame)
|
||
|
||
finish_frame = await manager.send_reply(
|
||
generate_req_id("upload_finish"),
|
||
{"upload_id": upload_id},
|
||
"aibot_upload_media_finish",
|
||
)
|
||
media_id = _response_body(finish_frame).get("media_id")
|
||
if not media_id:
|
||
raise RuntimeError("上传完成但未返回 media_id")
|
||
|
||
logger.info("图片已上传 media_id=%s…", str(media_id)[:12])
|
||
return str(media_id)
|
||
|
||
|
||
async def reply_image(ws_client: Any, frame: dict[str, Any], media_id: str) -> None:
|
||
await ws_client.reply(
|
||
frame,
|
||
{
|
||
"msgtype": "image",
|
||
"image": {"media_id": media_id},
|
||
},
|
||
)
|