157 lines
5.3 KiB
Python
157 lines
5.3 KiB
Python
"""企业微信智能机器人 · Skills 助手（skills 快查 + 截图预览 + Cursor 执行任务）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import sys
|
||
|
||
import env_config
|
||
from bridge_manager import shutdown_cursor_bridge, warm_cursor_bridge
|
||
from router import route_message, routing_mode
|
||
from skills_service import handle_command, warm_feed_cache
|
||
from wecom_media import reply_image, upload_image
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
)
|
||
logger = logging.getLogger("skills-bot")
|
||
|
||
BOT_ID = env_config.env("WECOM_BOT_ID") or env_config.env("WECHAT_BOT_ID")
|
||
BOT_SECRET = env_config.env("WECOM_BOT_SECRET") or env_config.env("WECHAT_BOT_SECRET")
|
||
|
||
|
||
def _require_credentials() -> None:
|
||
if not BOT_ID or not BOT_SECRET:
|
||
print(
|
||
"请设置环境变量 WECOM_BOT_ID 和 WECOM_BOT_SECRET\n"
|
||
"（企业微信 → 智能机器人 → API 模式 → 长连接）",
|
||
file=sys.stderr,
|
||
)
|
||
sys.exit(1)
|
||
|
||
|
||
def create_client():
|
||
from aibot import WSClient, WSClientOptions, generate_req_id
|
||
|
||
ws_client = WSClient(
|
||
WSClientOptions(
|
||
bot_id=BOT_ID,
|
||
secret=BOT_SECRET,
|
||
)
|
||
)
|
||
|
||
@ws_client.on("authenticated")
|
||
def on_authenticated():
|
||
logger.info("企业微信长连接认证成功，路由模式=%s", routing_mode())
|
||
cursor_key = env_config.env("CURSOR_API_KEY")
|
||
if cursor_key:
|
||
logger.info("CURSOR_API_KEY 已加载（%s…）", cursor_key[:8])
|
||
try:
|
||
warm_cursor_bridge()
|
||
logger.info("Cursor bridge 预启动完成")
|
||
except Exception as exc:
|
||
logger.warning("Cursor bridge 预启动失败（Cursor 任务时会重试）: %s", exc)
|
||
else:
|
||
logger.warning("CURSOR_API_KEY 未配置，Cursor 任务将失败")
|
||
try:
|
||
warm_feed_cache()
|
||
logger.info("skills 数据预加载完成")
|
||
except Exception as exc:
|
||
logger.warning("skills 数据预加载失败: %s", exc)
|
||
|
||
@ws_client.on("event.enter_chat")
|
||
async def on_enter_chat(frame):
|
||
help_text = handle_command("help")
|
||
extra = (
|
||
"\n\n---\n"
|
||
"**单页截图**：`preview` / `截图` / `预览 [路径或URL]`\n"
|
||
"**网页操作**：自然语言多步操作，或 `browser 场景名`\n"
|
||
"例：`访问登录页，输入账号密码，点击登录，点击智能体管理，截图`\n"
|
||
"场景文件：`bot/scenarios/*.yaml`（可用 `browser xiaobao-agent-manage`）"
|
||
)
|
||
await ws_client.reply_welcome(
|
||
frame,
|
||
{
|
||
"msgtype": "markdown",
|
||
"markdown": {"content": help_text + extra},
|
||
},
|
||
)
|
||
|
||
@ws_client.on("message.text")
|
||
async def on_text(frame):
|
||
body = frame.get("body", {})
|
||
content = body.get("text", {}).get("content", "")
|
||
logger.info("收到消息: %s", content)
|
||
|
||
stream_id = generate_req_id("stream")
|
||
last_progress = ""
|
||
|
||
async def on_progress(message: str) -> None:
|
||
nonlocal last_progress
|
||
if message != last_progress:
|
||
last_progress = message
|
||
await ws_client.reply_stream(frame, stream_id, message, False)
|
||
|
||
await ws_client.reply_stream(frame, stream_id, "收到，正在处理…", False)
|
||
|
||
try:
|
||
result = await route_message(content, on_progress=on_progress)
|
||
reply = result.text
|
||
logger.info(
|
||
"回复来源: %s, 文本长度=%d, 图片=%s",
|
||
result.source,
|
||
len(reply),
|
||
result.image_path or "-",
|
||
)
|
||
except Exception as exc:
|
||
logger.exception("处理失败")
|
||
reply = f"处理失败：{exc}"
|
||
result = None
|
||
|
||
if len(reply) > 3800:
|
||
reply = reply[:3800] + "\n\n> …内容已截断"
|
||
|
||
await ws_client.reply_stream(frame, stream_id, reply, True)
|
||
|
||
if result and result.image_path:
|
||
try:
|
||
media_id = await upload_image(ws_client, result.image_path)
|
||
await reply_image(ws_client, frame, media_id)
|
||
logger.info("图片已发送到企微: %s", result.image_path)
|
||
except Exception as exc:
|
||
logger.exception("发送图片失败")
|
||
await ws_client.reply(
|
||
frame,
|
||
{
|
||
"msgtype": "markdown",
|
||
"markdown": {
|
||
"content": f"截图文件：`{result.image_path}`\n发图失败：{exc}\n\n请确认 bot 已重启，或发送 `截图` 重试。",
|
||
},
|
||
},
|
||
)
|
||
|
||
@ws_client.on("error")
|
||
def on_error(error):
|
||
logger.error("连接错误: %s", error)
|
||
|
||
@ws_client.on("disconnected")
|
||
def on_disconnected(reason):
|
||
logger.warning("连接断开: %s", reason)
|
||
|
||
return ws_client
|
||
|
||
|
||
def main() -> None:
|
||
import atexit
|
||
|
||
atexit.register(shutdown_cursor_bridge)
|
||
_require_credentials()
|
||
client = create_client()
|
||
logger.info("启动 Skills 助手，Bot ID=%s…", BOT_ID[:8] if BOT_ID else "?")
|
||
client.run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|