Files
ai-mail-bot/src/tg_bot.py
T

637 lines
24 KiB
Python

import json
import logging
import threading
from concurrent.futures import ThreadPoolExecutor
from email.utils import parseaddr
import requests
from src.config import TelegramConfig, Config
from src.ai_client import generate_more_replies, expand_promo_detail
from src.smtp_client import send_reply
from src.retry import retry
from src.database import save_email_context, load_email_context, save_conversation, load_conversation, delete_conversation
logger = logging.getLogger(__name__)
_ai_pool = ThreadPoolExecutor(max_workers=3, thread_name_prefix="ai_regen")
# ── Public API ──────────────────────────────────────────
@retry()
def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
summary_data: dict, original_body: str, account_idx: int,
original_sender: str = "", original_recipient: str = "",
original_reply_to: str = "", account_email: str = "") -> int:
can_reply = summary_data.get("can_reply", True)
is_promotion = summary_data.get("category") == "promotion"
links = summary_data.get("links", [])
result = _tg_req(tg_cfg, "sendMessage", {
"chat_id": chat_id,
"text": summary_text,
"parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply, is_promotion, links),
})
msg_id = result["result"]["message_id"]
save_email_context(msg_id, {
"summary_text": summary_text,
"summary_data": summary_data,
"original_body": original_body,
"original_sender": original_sender or summary_data.get("sender", ""),
"original_recipient": original_recipient or summary_data.get("recipient", ""),
"original_reply_to": original_reply_to,
"account_idx": account_idx,
"account_email": account_email,
"sender": summary_data.get("sender", ""),
"subject": summary_data.get("subject", ""),
"can_reply": can_reply,
})
return msg_id
@retry()
def send_text(tg_cfg: TelegramConfig, chat_id: str, text: str,
reply_markup: dict = None) -> int:
payload = {"chat_id": chat_id, "text": text}
if reply_markup:
payload["reply_markup"] = reply_markup
result = _tg_req(tg_cfg, "sendMessage", payload)
return result["result"]["message_id"]
def send_thinking(tg_cfg: TelegramConfig, chat_id: str, text: str = "💭 思考中…") -> int:
result = _tg_req(tg_cfg, "sendMessage", {
"chat_id": chat_id,
"text": text,
})
return result["result"]["message_id"]
def edit_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int,
text: str, reply_markup: dict = None):
payload = {
"chat_id": chat_id, "message_id": msg_id,
"text": text, "parse_mode": "MarkdownV2",
}
if reply_markup:
payload["reply_markup"] = reply_markup
_tg_req(tg_cfg, "editMessageText", payload)
def delete_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int):
try:
_tg_req(tg_cfg, "deleteMessage", {
"chat_id": chat_id, "message_id": msg_id,
})
except Exception:
pass
def poll_telegram(tg_cfg: TelegramConfig, cfg: Config, last_update_id: int) -> int:
try:
resp = requests.get(
f"https://api.telegram.org/bot{tg_cfg.bot_token}/getUpdates",
params={"offset": last_update_id, "timeout": 1},
timeout=10,
)
resp.raise_for_status()
data = resp.json()
if not data.get("ok"):
return last_update_id
for update in data.get("result", []):
uid = update["update_id"]
try:
_handle_update(tg_cfg, cfg, update)
except Exception as e:
logger.warning("处理 update %d 失败: %s", uid, e)
last_update_id = uid + 1
except Exception as e:
logger.warning("Telegram polling error: %s", e)
return last_update_id
# ── Formatting ──────────────────────────────────────────
def format_summary(data: dict, account_email: str = "") -> str:
category = data.get("category", "normal")
if category == "promotion":
return _format_promotion(data, account_email)
if category == "notification" and data.get("verification_code"):
return _format_verification(data, account_email)
return _format_normal(data, account_email)
def _format_promotion(data: dict, account_email: str = "") -> str:
summary = data.get("summary", "")
lines = [
f"📢 *推广*",
f"━━━━━━━━━━━━━━━━━━",
f"*收件账户:* {_escape(account_email)}",
"",
f"{_escape(summary)}",
]
return "\n".join(lines)
def _format_verification(data: dict, account_email: str = "") -> str:
code = data.get("verification_code", "")
app = data.get("app_name", "")
lines = [
f"🔑 *验证码*",
f"━━━━━━━━━━━━━━━━━━",
f"*收件账户:* {_escape(account_email)}",
f"*应用:* {_escape(app)}",
f"*验证码:* `{_escape(code)}`",
]
return "\n".join(lines)
def _format_normal(data: dict, account_email: str = "") -> str:
lines = [
f"📧 *{_escape(data.get('subject', '邮件摘要'))}*",
f"━━━━━━━━━━━━━━━━━━",
f"*收件账户:* {_escape(account_email)}",
f"*收件人:* {_escape(data.get('recipient', '未知'))}",
f"*发件人:* {_escape(data.get('sender', '未知'))}",
"",
_escape(data.get("summary", "")),
]
return "\n".join(lines)
def _escape(text: str) -> str:
for ch in "_*[]()~`>#+-=|{}.!":
text = text.replace(ch, f"\\{ch}")
return text
# ── Inline keyboards ────────────────────────────────────
CALLBACK_VIEW_ORIG = "v"
CALLBACK_VIEW_SUMM = "s"
CALLBACK_REPLY = "r"
CALLBACK_AI_REPLY = "a"
CALLBACK_SELECT_REPLY = "sr"
CALLBACK_REGEN = "rg"
CALLBACK_HINT = "h"
CALLBACK_CANCEL = "c"
CALLBACK_PROMO_DETAIL = "pd"
def _summary_keyboard(can_reply: bool = True, is_promotion: bool = False, links: list = None) -> dict:
if is_promotion:
return {"inline_keyboard": [
[{"text": "告诉我更多", "callback_data": CALLBACK_PROMO_DETAIL}],
[{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG}],
]}
kb = [
[
{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG},
{"text": "回复", "callback_data": CALLBACK_REPLY},
],
]
if can_reply:
kb.append([{"text": "AI回复", "callback_data": CALLBACK_AI_REPLY}])
for link in (links or []):
url = link.get("url", "")
text = link.get("text", "打开链接")[:30]
# 校验 URL:合法 http/https,无多重协议头,长度合理
if url.startswith("http") and url.count("://") == 1 and len(url) < 200:
kb.append([{"text": f"🔗 {text}", "url": url}])
return {"inline_keyboard": kb}
def _orig_keyboard(can_reply: bool = True) -> dict:
kb = [
[
{"text": "返回摘要", "callback_data": CALLBACK_VIEW_SUMM},
{"text": "回复", "callback_data": CALLBACK_REPLY},
],
]
if can_reply:
kb.append([{"text": "AI回复", "callback_data": CALLBACK_AI_REPLY}])
return {"inline_keyboard": kb}
def _cancel_keyboard() -> dict:
return {
"inline_keyboard": [
[{"text": "取消回复", "callback_data": CALLBACK_CANCEL}],
]
}
def _ai_reply_keyboard(msg_id: int, suggestions: list) -> dict:
kb = []
for i, s in enumerate(suggestions):
data = f"{CALLBACK_SELECT_REPLY}|{msg_id}|{i}"
kb.append([{"text": s["text"][:30], "callback_data": data}])
kb.append([
{"text": "换一批", "callback_data": f"{CALLBACK_REGEN}|{msg_id}"},
{"text": "我想说:", "callback_data": f"{CALLBACK_HINT}|{msg_id}"},
])
kb.append([{"text": "取消", "callback_data": CALLBACK_CANCEL}])
return {"inline_keyboard": kb}
# ── Telegram API ────────────────────────────────────────
def _tg_req(tg_cfg: TelegramConfig, method: str, payload: dict) -> dict:
url = f"https://api.telegram.org/bot{tg_cfg.bot_token}/{method}"
logger.info("Telegram API %s: %s", method, {k: (v if k != "reply_markup" else "...") for k, v in payload.items()})
r = requests.post(url, json=payload, timeout=30)
if not r.ok:
logger.error("Telegram API %s 失败: %s %s", method, r.status_code, r.text[:500])
r.raise_for_status()
result = r.json()
if not result.get("ok"):
raise RuntimeError(f"Telegram API error: {result}")
return result
# ── Update handling ─────────────────────────────────────
def _do_regen(cfg: Config, chat_id: int, msg_id: int, ctx: dict, conv: dict):
"""在 AI 线程池中执行换一批,不阻塞按钮轮询"""
try:
new_suggestions = generate_more_replies(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []),
)
except Exception as e:
send_text(cfg.telegram, str(chat_id), f"生成失败: {e}")
return
if new_suggestions:
conv["ai_suggestions"] = new_suggestions
conv["all_suggestions"].extend(new_suggestions)
save_conversation(chat_id, conv)
logger.info("[ai_regen] 新生成 %d 条回复", len(new_suggestions))
_show_ai_suggestions(cfg.telegram, chat_id, msg_id, new_suggestions)
def _do_ai_hint(cfg: Config, chat_id: int, conv: dict, ctx: dict, hint: str):
"""在 AI 线程池中根据用户提示生成回复"""
try:
new_suggestions = generate_more_replies(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []), user_hint=hint,
)
except Exception as e:
send_text(cfg.telegram, str(chat_id), f"生成失败: {e}")
delete_conversation(chat_id)
return
if new_suggestions:
conv["ai_suggestions"] = new_suggestions
conv["all_suggestions"].extend(new_suggestions)
conv["state"] = "ai_reply_selection"
conv.pop("prompt_msg_id", None)
save_conversation(chat_id, conv)
logger.info("[ai_hint] 根据提示生成 %d 条新回复", len(new_suggestions))
_show_ai_suggestions(cfg.telegram, chat_id, conv["summary_msg_id"], new_suggestions)
def _handle_update(tg_cfg: TelegramConfig, cfg: Config, update: dict):
if "callback_query" in update:
_handle_callback(tg_cfg, cfg, update["callback_query"])
elif "message" in update:
_handle_message(tg_cfg, cfg, update["message"])
def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
chat_id = cb["message"]["chat"]["id"]
msg_id = cb["message"]["message_id"]
data = cb["data"]
cb_id = cb["id"]
try:
_tg_req(tg_cfg, "answerCallbackQuery", {"callback_query_id": cb_id})
except Exception:
pass
parts = data.split("|")
action = parts[0]
ctx = load_email_context(msg_id)
subject = ctx["subject"] if ctx else "?"
logger.info("TG 回调: action=%s msg_id=%d subject=%s", action, msg_id, subject)
if action == CALLBACK_VIEW_ORIG and ctx:
can_reply = ctx.get("can_reply", True)
text = f"*📄 原文 \\- {_escape(ctx['subject'])}*\n\n{_escape(ctx.get('plain_text', ctx['original_body']))}"
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": text, "parse_mode": "MarkdownV2",
"reply_markup": _orig_keyboard(can_reply),
})
logger.info(" 切换为原文视图")
elif action == CALLBACK_VIEW_SUMM and ctx:
can_reply = ctx.get("can_reply", True)
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply),
})
logger.info(" 切换回摘要视图")
elif action == CALLBACK_REPLY and ctx:
prompt_msg_id = send_text(
tg_cfg, str(chat_id), "请输入你的回复内容:",
reply_markup=_cancel_keyboard(),
)
save_conversation(chat_id, {
"state": "awaiting_reply", "summary_msg_id": msg_id,
"prompt_msg_id": prompt_msg_id,
})
logger.info(" 进入回复流程, prompt_msg_id=%d", prompt_msg_id)
elif action == CALLBACK_AI_REPLY and ctx:
suggestions = ctx["summary_data"].get("reply_suggestions", [])
if not suggestions:
send_text(tg_cfg, str(chat_id), "此邮件无需回复。")
logger.info(" AI 判定无需回复,隐藏")
return
save_conversation(chat_id, {
"state": "ai_reply_selection",
"summary_msg_id": msg_id,
"ai_suggestions": suggestions,
"all_suggestions": list(suggestions),
})
logger.info(" 显示 AI 回复建议 (%d 条)", len(suggestions))
_show_ai_suggestions(tg_cfg, chat_id, msg_id, suggestions)
elif action == CALLBACK_SELECT_REPLY and ctx:
idx = int(parts[2])
conv = load_conversation(chat_id)
if not conv or conv.get("summary_msg_id") != msg_id:
return
suggestions = conv.get("ai_suggestions", [])
if idx >= len(suggestions):
return
sel = suggestions[idx]
logger.info(" 选择回复 #%d: %s", idx, sel["text"][:60])
_do_send_reply(tg_cfg, cfg, chat_id, msg_id, ctx, sel["text"])
send_text(tg_cfg, str(chat_id), f"✅ 已发送: {sel['text']}")
delete_conversation(chat_id)
can_reply = ctx.get("can_reply", True)
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply),
})
elif action == CALLBACK_REGEN and ctx:
conv = load_conversation(chat_id)
if not conv or conv.get("summary_msg_id") != msg_id:
return
logger.info(" 换一批 (已有 %d 条历史)", len(conv.get("all_suggestions", [])))
# AI 处理中
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
# 丢到 AI 线程池处理,不阻塞按钮轮询
_ai_pool.submit(_do_regen, cfg, chat_id, msg_id, ctx, conv)
elif action == CALLBACK_HINT and ctx:
prompt_msg_id = send_text(
tg_cfg, str(chat_id), "请输入你对回复的提示/要求:",
reply_markup=_cancel_keyboard(),
)
conv = load_conversation(chat_id)
if conv and conv.get("summary_msg_id") == msg_id:
conv["state"] = "awaiting_ai_hint"
conv["prompt_msg_id"] = prompt_msg_id
save_conversation(chat_id, conv)
logger.info(" 等待用户输入 AI 提示")
elif action == CALLBACK_CANCEL:
delete_conversation(chat_id)
if ctx:
can_reply = ctx.get("can_reply", True)
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply),
})
logger.info(" 取消, 恢复摘要视图")
else:
delete_message(tg_cfg, str(chat_id), msg_id)
logger.info(" 取消, 删除提示消息")
elif action == CALLBACK_PROMO_DETAIL and ctx:
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
try:
details = expand_promo_detail(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
)
except Exception as e:
send_text(tg_cfg, str(chat_id), f"获取详情失败: {e}")
return
summary = ctx["summary_data"].get("summary", "")
detail_lines = "\n".join(f"\\- {_escape(d)}" for d in details) if details else "暂无更多详情"
expanded = (
f"📢 *推广*\n"
f"━━━━━━━━━━━━━━━━━━\n"
f"*收件账户:* {_escape(ctx.get('account_email', ''))}\n"
f"\n{_escape(summary)}\n"
f"\n*📋 详情:*\n{detail_lines}"
)
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": expanded, "parse_mode": "MarkdownV2",
"reply_markup": {"inline_keyboard": [
[{"text": "收起", "callback_data": CALLBACK_VIEW_SUMM}],
[{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG}],
]},
})
logger.info(" 显示推广详情")
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
chat_id = msg["chat"]["id"]
text = msg.get("text", "").strip()
conv = load_conversation(chat_id)
if not conv:
return
user_msg_id = msg["message_id"]
chat_id_str = str(chat_id)
prompt_msg_id = conv.get("prompt_msg_id")
state = conv.get("state")
media = _extract_media(msg)
logger.info("TG 消息: state=%s text=%s media=%d", state, text[:60] if text else "(无)", len(media))
if state == "awaiting_reply":
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
delete_conversation(chat_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
logger.info(" 上下文已丢失,清理")
return
# 有附件但没文字:累积附件,等用户输入文字
if media and not text:
pending = conv.get("pending_attachments", [])
for att in media:
pending.append(att)
conv["pending_attachments"] = pending
save_conversation(chat_id, conv)
delete_message(tg_cfg, chat_id_str, user_msg_id)
send_text(tg_cfg, chat_id_str,
f"📎 已添加 {len(pending)} 个附件,请输入回复内容发送",
reply_markup=_cancel_keyboard())
return
# 有文字(可能也有附件):发送邮件
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
# 合并累积的附件
all_attachments = conv.get("pending_attachments", [])
all_attachments.extend(media)
_do_send_reply(tg_cfg, cfg, chat_id, conv["summary_msg_id"], ctx, text,
attachments=all_attachments or None)
attach_info = f" (+{len(all_attachments)}个附件)" if all_attachments else ""
send_text(tg_cfg, chat_id_str, f"✅ 回复已发送: {text}{attach_info}")
delete_conversation(chat_id)
logger.info(" 回复发送完成")
elif state == "awaiting_ai_hint":
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
delete_conversation(chat_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
return
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
# 更新消息为 AI 思考中
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": conv["summary_msg_id"],
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
# 丢到 AI 线程池处理
_ai_pool.submit(_do_ai_hint, cfg, chat_id, conv, ctx, text)
def _show_ai_suggestions(tg_cfg: TelegramConfig, chat_id: int, msg_id: int,
suggestions: list):
text = "*AI 建议回复,请选择:*\n\n"
for i, s in enumerate(suggestions, 1):
text += f"{i}\\. {_escape(s['text'])}\n"
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": text, "parse_mode": "MarkdownV2",
"reply_markup": _ai_reply_keyboard(msg_id, suggestions),
})
def _download_tg_file(tg_cfg: TelegramConfig, file_id: str, filename: str) -> str:
"""从 Telegram 下载文件到本地 media/ 目录,返回本地路径"""
import os
from pathlib import Path
media_dir = Path("media")
media_dir.mkdir(exist_ok=True)
resp = requests.get(
f"https://api.telegram.org/bot{tg_cfg.bot_token}/getFile",
params={"file_id": file_id}, timeout=30,
)
resp.raise_for_status()
file_path = resp.json()["result"]["file_path"]
dl = requests.get(
f"https://api.telegram.org/file/bot{tg_cfg.bot_token}/{file_path}",
timeout=60,
)
dl.raise_for_status()
dest = media_dir / filename
dest.write_bytes(dl.content)
logger.info(" 下载附件: %s (%d bytes)", filename, len(dl.content))
return str(dest)
def _extract_media(msg: dict) -> list[dict]:
"""从 Telegram 消息中提取附件信息列表 [{file_id, filename, mime_type}]"""
attachments = []
# 照片取最大尺寸
if "photo" in msg:
photo = max(msg["photo"], key=lambda p: p.get("file_size", 0))
ext = ".jpg"
attachments.append({
"file_id": photo["file_id"],
"filename": f"photo_{photo['file_id'][:8]}{ext}",
"mime_type": "image/jpeg",
})
if "video" in msg:
v = msg["video"]
attachments.append({
"file_id": v["file_id"],
"filename": v.get("file_name", f"video_{v['file_id'][:8]}.mp4"),
"mime_type": v.get("mime_type", "video/mp4"),
})
if "document" in msg:
d = msg["document"]
attachments.append({
"file_id": d["file_id"],
"filename": d.get("file_name", f"doc_{d['file_id'][:8]}"),
"mime_type": d.get("mime_type", "application/octet-stream"),
})
if "voice" in msg:
vc = msg["voice"]
attachments.append({
"file_id": vc["file_id"],
"filename": f"voice_{vc['file_id'][:8]}.ogg",
"mime_type": vc.get("mime_type", "audio/ogg"),
})
if "sticker" in msg:
s = msg["sticker"]
if not s.get("is_animated"):
attachments.append({
"file_id": s["file_id"],
"filename": f"sticker_{s['file_id'][:8]}.webp",
"mime_type": "image/webp",
})
return attachments
def _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
chat_id: int, msg_id: int, ctx: dict, reply_text: str,
attachments: list[dict] | None = None):
acct = cfg.email_accounts[ctx["account_idx"]]
if not acct.smtp:
send_text(tg_cfg, str(chat_id), "❌ 该邮箱未配置 SMTP,无法发送回复。")
return
to_addr = parseaddr(ctx.get("original_reply_to", ""))[1]
if not to_addr:
to_addr = parseaddr(ctx["original_sender"])[1]
if not to_addr:
to_addr = parseaddr(ctx["sender"])[1]
if not to_addr:
to_addr = ctx["sender"]
subject = ctx["subject"]
# 下载 Telegram 附件到本地
local_files = []
if attachments:
for att in attachments:
try:
path = _download_tg_file(tg_cfg, att["file_id"], att["filename"])
local_files.append((path, att["mime_type"]))
except Exception as e:
logger.error(" 附件下载失败: %s - %s", att["filename"], e)
send_reply(acct, to_addr, subject, reply_text, attachments=local_files or None)