133 lines
4.1 KiB
Python
133 lines
4.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
轻量的 QQ Bot 客户端（用于发送私聊文本和媒体）
|
||
|
||
使用环境变量：
|
||
- QQ_BOT_URL: bot 接收地址（例如 http://localhost:30000/send_private_msg）
|
||
- QQ_BOT_TARGET_ID: 默认目标 QQ id（可被函数参数覆盖）
|
||
|
||
示例：
|
||
from qq import send_msg, send_media_msg
|
||
send_msg('hello')
|
||
# 发送图片时将自动转换为 "base64://..." 格式
|
||
send_media_msg('figures/btcusd_interval_pred_60m.png', type='image')
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import json
|
||
from typing import Optional
|
||
import base64
|
||
|
||
try:
|
||
import requests
|
||
except Exception: # pragma: no cover - runtime dependency
|
||
requests = None
|
||
|
||
|
||
def _bot_url() -> str:
|
||
return os.getenv("QQ_BOT_URL", "http://localhost:30000/send_private_msg")
|
||
|
||
|
||
def _default_target() -> Optional[str]:
|
||
return os.getenv("QQ_BOT_TARGET_ID")
|
||
|
||
|
||
def _file_to_base64_uri(path: str) -> str:
|
||
"""将本地文件读取为 base64 uri 字符串（前缀 base64://）。"""
|
||
with open(path, "rb") as f:
|
||
data = f.read()
|
||
b64 = base64.b64encode(data).decode("ascii")
|
||
return f"base64://{b64}"
|
||
|
||
|
||
def send_msg(msg: str, target_id: Optional[str] = None, timeout: float = 8.0):
|
||
"""发送文本消息到 QQ 机器人。返回 requests.Response 或 None（失败）。"""
|
||
if target_id is None:
|
||
target_id = _default_target()
|
||
if target_id is None:
|
||
raise ValueError("target_id 未提供，且环境变量 QQ_BOT_TARGET_ID 未设置")
|
||
|
||
payload = {
|
||
"user_id": str(target_id),
|
||
"message": [
|
||
{"type": "text", "data": {"text": str(msg)}}
|
||
],
|
||
}
|
||
|
||
url = _bot_url()
|
||
print(f"[QQ] 发送文本到 {target_id}: {msg}")
|
||
if requests is None:
|
||
print("requests 未安装，无法发送消息")
|
||
return None
|
||
|
||
try:
|
||
resp = requests.post(url, json=payload, timeout=timeout)
|
||
try:
|
||
# 非必要：打印返回的简短信息
|
||
print(f"[QQ] 返回: {resp.status_code} {resp.text[:200]}")
|
||
except Exception:
|
||
pass
|
||
return resp
|
||
except Exception as e:
|
||
print(f"[QQ] 发送文本消息出错: {e}")
|
||
return None
|
||
|
||
|
||
def send_media_msg(file_path: str, target_id: Optional[str] = None, type: str = "image", timeout: float = 15.0):
|
||
"""发送媒体消息（image 或 video）。
|
||
|
||
- 当 type == "image" 时：读取文件并以 "base64://..." 形式发送。
|
||
- 当 type == "video" 时：仍使用 "file://" 本地文件引用（保持原有行为）。
|
||
"""
|
||
if type not in ("image", "video"):
|
||
raise ValueError("type 必须为 'image' 或 'video'")
|
||
|
||
if target_id is None:
|
||
target_id = _default_target()
|
||
if target_id is None:
|
||
raise ValueError("target_id 未提供，且环境变量 QQ_BOT_TARGET_ID 未设置")
|
||
|
||
# 根据类型构造 file 字段
|
||
if type == "image":
|
||
try:
|
||
file_field = _file_to_base64_uri(file_path)
|
||
except Exception as e:
|
||
print(f"[QQ] 读取图片失败: {e}")
|
||
return None
|
||
else: # video
|
||
file_field = f"file://{file_path}"
|
||
|
||
payload = {
|
||
"user_id": str(target_id),
|
||
"message": [
|
||
{"type": type, "data": {"file": file_field}}
|
||
],
|
||
}
|
||
|
||
url = _bot_url()
|
||
print(f"[QQ] 发送媒体 {type} 到 {target_id}: {file_path}")
|
||
if requests is None:
|
||
print("requests 未安装，无法发送媒体消息")
|
||
return None
|
||
|
||
try:
|
||
resp = requests.post(url, json=payload, timeout=timeout)
|
||
try:
|
||
print(f"[QQ] 返回: {resp.status_code} {resp.text[:200]}")
|
||
except Exception:
|
||
pass
|
||
return resp
|
||
except Exception as e:
|
||
print(f"[QQ] 发送媒体消息出错: {e}")
|
||
return None
|
||
|
||
|
||
def send_prediction_result(summary: str, png_path: Optional[str] = None, target_id: Optional[str] = None):
|
||
"""便捷函数：先发送 summary 文本，再发送 png（若提供）。"""
|
||
r1 = send_msg(summary, target_id=target_id)
|
||
r2 = None
|
||
if png_path:
|
||
r2 = send_media_msg(png_path, target_id=target_id, type="image")
|
||
return r1, r2
|