113 lines
4.8 KiB
Python
113 lines
4.8 KiB
Python
import logging
|
||
import re
|
||
import threading
|
||
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__)
|
||
|
||
_MD_SPECIAL = re.compile(r"([_*\[\]()~`>#+\-=|{}.!\\])")
|
||
|
||
def _escape_md(text: str) -> str:
|
||
return _MD_SPECIAL.sub(r"\\\1", text)
|
||
|
||
|
||
def _mail_label(mail: Email) -> str:
|
||
sender_short = mail.sender.split("<")[0].strip() if "<" in mail.sender else mail.sender.split("@")[0]
|
||
subject_short = mail.subject[:30] + ("…" if len(mail.subject) > 30 else "")
|
||
return f"{_escape_md(sender_short)} · {_escape_md(subject_short)}"
|
||
|
||
|
||
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}")
|
||
label = _mail_label(mail)
|
||
# 阶段1：显示等待处理
|
||
if tg_msg_id:
|
||
edit_message(cfg.telegram, cfg.telegram.chat_id, tg_msg_id, f"⏳ *等待处理…*\n{label}")
|
||
else:
|
||
tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, f"⏳ *等待处理…*\n{label}")
|
||
|
||
# 阶段2：AI 处理，更新为思考中
|
||
edit_message(cfg.telegram, cfg.telegram.chat_id, tg_msg_id, f"🤖 *AI思考中…*\n{label}")
|
||
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
|
||
from src.database import save_email_context
|
||
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))
|
||
# 保存上下文用于回调
|
||
save_email_context(thinking_msg_id, {
|
||
"summary_text": info["text"],
|
||
"summary_data": info["data"],
|
||
"original_body": info["original_body"],
|
||
"plain_text": info["data"].get("plain_text", ""),
|
||
"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", ""),
|
||
"account_idx": info["acct_idx"],
|
||
"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", ""))
|
||
# 标记已读放到后台线程，不阻塞 TG 发送
|
||
def _bg_mark():
|
||
try:
|
||
mark_as_seen(acct, [(info["uid"], info.get("folder", "INBOX"))])
|
||
except Exception:
|
||
pass
|
||
threading.Thread(target=_bg_mark, daemon=True).start()
|
||
logger.info(f" TG 发送成功，标记已读已提交后台")
|