Files
ai-mail-bot/src/summarizer.py

92 lines
4.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import logging
from typing import Generator, Optional
from src.config import Config
from src.email_client import fetch_unseen_emails, mark_as_seen, Email
from src.ai_client import summarize_email
from src.tg_bot import send_summary, send_thinking, edit_message, format_summary
logger = logging.getLogger(__name__)
def poll_accounts(cfg: Config, skip_uids: set[bytes] | None = None) -> Generator[tuple[int, Email], None, None]:
for idx, acct in enumerate(cfg.email_accounts):
try:
logger.info(f"检查邮箱: {acct.username}")
emails = fetch_unseen_emails(acct, skip_uids=skip_uids)
if emails:
logger.info(f" 发现 {len(emails)} 封新邮件")
for mail in emails:
yield idx, mail
except Exception as e:
logger.error(f"轮询 {acct.username} 失败: {e}", exc_info=True)
def ai_process(cfg: Config, acct_idx: int, mail: Email, tg_msg_id: int = 0) -> Optional[dict]:
logger.info(f" 正在摘要: {mail.subject}")
# 阶段1显示等待处理
if tg_msg_id:
edit_message(cfg.telegram, cfg.telegram.chat_id, tg_msg_id, "⏳ 等待处理…")
else:
tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, "⏳ 等待处理…")
# 阶段2AI 处理,更新为思考中
edit_message(cfg.telegram, cfg.telegram.chat_id, tg_msg_id, "🤖 AI思考中…")
summary = summarize_email(cfg.ai, mail.recipient, mail.subject, mail.sender, mail.body, mail.account_email)
summary["recipient"] = mail.recipient
if "@" not in summary.get("sender", ""):
summary["sender"] = mail.sender
text = format_summary(summary, mail.account_email)
# 阶段3编辑为最终结果
return {
"text": text,
"data": summary,
"original_body": mail.body,
"original_sender": mail.sender,
"original_recipient": mail.recipient,
"original_reply_to": mail.reply_to,
"acct_idx": acct_idx,
"account_email": mail.account_email,
"uid": mail.uid,
"folder": mail.folder,
"thinking_msg_id": tg_msg_id if tg_msg_id else 0,
}
def tg_send_and_mark(cfg: Config, info: dict):
acct = cfg.email_accounts[info["acct_idx"]]
thinking_msg_id = info.get("thinking_msg_id")
if thinking_msg_id:
# 编辑思考中消息为实际摘要
from src.tg_bot import _summary_keyboard
can_reply = info["data"].get("can_reply", True)
is_promotion = info["data"].get("category") == "promotion"
links = info["data"].get("links", [])
edit_message(cfg.telegram, cfg.telegram.chat_id, thinking_msg_id,
info["text"], _summary_keyboard(can_reply, is_promotion, links))
# 保存上下文用于回调
from src.tg_bot import _email_contexts
_email_contexts[thinking_msg_id] = {
"summary_text": info["text"],
"summary_data": info["data"],
"original_body": info["original_body"][:10000],
"original_sender": info.get("original_sender", ""),
"original_recipient": info.get("original_recipient", ""),
"original_reply_to": info.get("original_reply_to", ""),
"account_email": info.get("account_email", ""),
"sender": info["data"].get("sender", ""),
"subject": info["data"].get("subject", ""),
"can_reply": can_reply,
}
else:
# 兜底:没有思考中消息就正常发送
send_summary(cfg.telegram, cfg.telegram.chat_id,
info["text"], info["data"],
info["original_body"], info["acct_idx"],
original_sender=info.get("original_sender", ""),
original_recipient=info.get("original_recipient", ""),
original_reply_to=info.get("original_reply_to", ""),
account_email=info.get("account_email", ""))
# 只有发送成功才标记已读
mark_as_seen(acct, [(info["uid"], info.get("folder", "INBOX"))])
logger.info(f" TG 发送成功,已标记已读")