455 lines
17 KiB
Python
455 lines
17 KiB
Python
import json
|
||
import logging
|
||
from email.utils import parseaddr
|
||
import requests
|
||
from src.config import TelegramConfig, Config
|
||
from src.ai_client import generate_more_replies
|
||
from src.smtp_client import send_reply
|
||
from src.retry import retry
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_email_contexts: dict[int, dict] = {}
|
||
# chat_id -> conversation state
|
||
_conversations: dict[int, dict] = {}
|
||
|
||
|
||
# ── 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"]
|
||
_email_contexts[msg_id] = {
|
||
"summary_text": summary_text,
|
||
"summary_data": summary_data,
|
||
"original_body": original_body[:10000],
|
||
"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 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": 10},
|
||
timeout=15,
|
||
)
|
||
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"
|
||
|
||
|
||
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_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]
|
||
if url.startswith("http"):
|
||
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 _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 and "text" in update["message"]:
|
||
_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 = _email_contexts.get(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['original_body'][:3500])}"
|
||
_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(),
|
||
)
|
||
_conversations[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
|
||
|
||
_conversations[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 = _conversations.get(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']}")
|
||
_conversations.pop(chat_id, None)
|
||
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 = _conversations.get(chat_id)
|
||
if not conv or conv.get("summary_msg_id") != msg_id:
|
||
return
|
||
logger.info(" 换一批 (已有 %d 条历史)", len(conv.get("all_suggestions", [])))
|
||
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(tg_cfg, str(chat_id), f"生成失败: {e}")
|
||
return
|
||
if new_suggestions:
|
||
conv["ai_suggestions"] = new_suggestions
|
||
conv["all_suggestions"].extend(new_suggestions)
|
||
logger.info(" 新生成 %d 条回复", len(new_suggestions))
|
||
_show_ai_suggestions(tg_cfg, chat_id, msg_id, new_suggestions)
|
||
|
||
elif action == CALLBACK_HINT and ctx:
|
||
prompt_msg_id = send_text(
|
||
tg_cfg, str(chat_id), "请输入你对回复的提示/要求：",
|
||
reply_markup=_cancel_keyboard(),
|
||
)
|
||
conv = _conversations.get(chat_id)
|
||
if conv and conv.get("summary_msg_id") == msg_id:
|
||
conv["state"] = "awaiting_ai_hint"
|
||
conv["prompt_msg_id"] = prompt_msg_id
|
||
logger.info(" 等待用户输入 AI 提示")
|
||
|
||
elif action == CALLBACK_CANCEL:
|
||
_conversations.pop(chat_id, None)
|
||
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(" 取消, 删除提示消息")
|
||
|
||
|
||
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
|
||
chat_id = msg["chat"]["id"]
|
||
text = msg["text"].strip()
|
||
conv = _conversations.get(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")
|
||
|
||
logger.info("TG 消息: state=%s text=%s", state, text[:60])
|
||
|
||
if state == "awaiting_reply":
|
||
ctx = _email_contexts.get(conv["summary_msg_id"])
|
||
if not ctx:
|
||
_conversations.pop(chat_id, None)
|
||
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
|
||
|
||
delete_message(tg_cfg, chat_id_str, user_msg_id)
|
||
if prompt_msg_id:
|
||
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
|
||
_do_send_reply(tg_cfg, cfg, chat_id, conv["summary_msg_id"], ctx, text)
|
||
send_text(tg_cfg, chat_id_str, f"✅ 回复已发送: {text}")
|
||
_conversations.pop(chat_id, None)
|
||
logger.info(" 回复发送完成")
|
||
|
||
elif state == "awaiting_ai_hint":
|
||
ctx = _email_contexts.get(conv["summary_msg_id"])
|
||
if not ctx:
|
||
_conversations.pop(chat_id, None)
|
||
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)
|
||
try:
|
||
new_suggestions = generate_more_replies(
|
||
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
|
||
conv.get("all_suggestions", []), user_hint=text,
|
||
)
|
||
except Exception as e:
|
||
send_text(tg_cfg, chat_id_str, f"生成失败: {e}")
|
||
_conversations.pop(chat_id, None)
|
||
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)
|
||
logger.info(" 根据提示生成 %d 条新回复", len(new_suggestions))
|
||
_show_ai_suggestions(tg_cfg, chat_id, conv["summary_msg_id"], new_suggestions)
|
||
|
||
|
||
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 _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
|
||
chat_id: int, msg_id: int, ctx: dict, reply_text: str):
|
||
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"]
|
||
send_reply(acct, to_addr, subject, reply_text)
|