项目初始化

bot/wecom_media.py

@@ -0,0 +1,96 @@
+"""企业微信 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},
+        },
+    )