Compare commits

..

3 Commits

5 changed files with 224 additions and 54 deletions

47
main.py
View File

@@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor
from src.config import load_config
from src.summarizer import poll_accounts, ai_process, tg_send_and_mark
from src.tg_bot import poll_telegram, send_thinking
from src.database import init_db, is_uids_processed, mark_fetched, mark_processing, mark_sent, mark_failed
# Thread-safe logging: all log records go through a single QueueListener thread
_log_queue = queue_module.Queue(-1)
@@ -29,9 +30,6 @@ logger = logging.getLogger("main")
_running = True
_shutdown_event = threading.Event()
_pending_lock = threading.Lock()
_pending_uids: set[bytes] = set()
_processing_uids: set[bytes] = set() # 正在处理含等待TG发送的邮件UID
# Queue 1: raw emails → AI processor
_email_queue: queue_module.Queue = queue_module.Queue()
@@ -50,19 +48,18 @@ def _email_poller(cfg):
logger.info("邮件轮询线程已启动")
while _running:
try:
for acct_idx, mail in poll_accounts(cfg, skip_uids=_processing_uids):
for acct_idx, mail in poll_accounts(cfg):
if not _running:
return
with _pending_lock:
if mail.uid not in _processing_uids:
_processing_uids.add(mail.uid)
_pending_uids.add(mail.uid)
# 发送"获取邮件中"提示
try:
tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, "📥 获取邮件中...")
except Exception:
tg_msg_id = 0
_email_queue.put((acct_idx, mail, tg_msg_id))
uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
if not is_uid_processed(uid_str):
mark_fetched(uid_str, mail.account_email, mail.folder,
mail.subject, mail.sender)
try:
tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, "📥 获取邮件中…")
except Exception:
tg_msg_id = 0
_email_queue.put((acct_idx, mail, tg_msg_id))
except Exception as e:
logger.error(f"邮件轮询线程异常: {e}", exc_info=True)
time.sleep(5)
@@ -81,6 +78,8 @@ def _ai_processor(cfg):
acct_idx, mail, tg_msg_id = _email_queue.get_nowait()
except queue_module.Empty:
break
uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
mark_processing(uid_str)
future = executor.submit(ai_process, cfg, acct_idx, mail, tg_msg_id)
pending_futures[future] = (acct_idx, mail)
@@ -93,10 +92,8 @@ def _ai_processor(cfg):
_tg_queue.put(info)
except Exception as e:
logger.error(f"AI 处理邮件失败: {e}", exc_info=True)
# AI 失败也要移除 pending否则永远卡住
with _pending_lock:
_pending_uids.discard(mail.uid)
_processing_uids.discard(mail.uid)
uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
mark_failed(uid_str)
time.sleep(0.1 if pending_futures else 0.5)
@@ -114,16 +111,15 @@ def _tg_worker(cfg):
try:
info = _tg_queue.get_nowait()
tg_send_and_mark(cfg, info)
# 发送并标记已读成功后移除 pending_processing_uids 保留直到重启)
with _pending_lock:
_pending_uids.discard(info["uid"])
uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"])
mark_sent(uid_str)
except queue_module.Empty:
break
except Exception as e:
logger.error(f"TG 发送失败: {e}", exc_info=True)
try:
with _pending_lock:
_pending_uids.discard(info["uid"])
uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"])
mark_failed(uid_str)
except Exception:
pass
@@ -138,6 +134,7 @@ def main():
cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
cfg = load_config(cfg_path)
init_db()
logger.info(f"AI邮件摘要机器人已启动轮询间隔: {cfg.polling.interval_seconds}s")
threads = [
@@ -155,8 +152,8 @@ def main():
while _running:
tick += 1
if tick % 30 == 0:
logger.info("队列状态: email_queue=%d tg_queue=%d pending_uids=%d",
_email_queue.qsize(), _tg_queue.qsize(), len(_pending_uids))
logger.info("队列状态: email_queue=%d tg_queue=%d",
_email_queue.qsize(), _tg_queue.qsize())
_shutdown_event.wait(1)
except KeyboardInterrupt:
_running = False

View File

@@ -16,6 +16,7 @@ SYSTEM_PROMPT = """你是一个邮件摘要助手。请分析邮件内容并以
"recipient": "收件邮箱地址",
"sender": "发件人",
"category": "normal/promotion/notification",
"plain_text": "邮件全文的纯文本版本去除HTML标签保留所有文字内容保持原始段落结构不删减不修改任何内容",
"summary": "50字以内的核心摘要",
"verification_code": "",
"app_name": "",
@@ -35,6 +36,7 @@ SYSTEM_PROMPT = """你是一个邮件摘要助手。请分析邮件内容并以
- app_name: 仅当 category 为 "notification" 且邮件包含验证码时填写,填写发送验证码的应用/平台名称如「GitHub」「微信」「Kraken」。非验证码邮件留空字符串
- links: 仅提取以下两类链接1邮件明确要求用户点击才能完成某操作的链接如确认邮箱、重置密码、查看账单、审批请求、查看通知详情等2退订/取消订阅链接。不要提取推广性质的链接如下载App、注册领优惠、查看更多商品等营销链接、追踪像素、邮件签名中的社交链接。每条包含简短描述文字和完整 URL。最多 3 条。没有符合条件的链接时为空数组
- 退订链接提取规则推广邮件必须提取退订链接。退订链接通常在邮件最底部包含以下关键词之一unsubscribe、退订、取消订阅、opt out、manage preferences、email preferences。即使链接被包裹在 HTML 标签中,也要提取 href 属性里的完整 URL。如果邮件底部同时有多个退订相关链接只提取最主要的那个
- plain_text: 将邮件全文转换为可读的纯文本版本。去除所有HTML标签、CSS样式、脚本代码保留完整的文字内容和原始段落结构。不删减、不修改、不总结任何内容完整保留原文所有信息
- summary 要简洁,把需要知道的信息浓缩成一句话,不需要分条列出
- 当 category 为 "promotion"summary 直接写成一句话:「[发送方名称] 在推广 [推广的产品/服务/活动]」,例如「腾讯云在推广双十一云服务器折扣」。此时 can_reply 为 falsereply_suggestions 为空数组
- can_reply 为 false 表示无需回复或不适合建议回复(此时忽略 reply_suggestions

173
src/database.py Normal file
View File

@@ -0,0 +1,173 @@
import json
import logging
import sqlite3
import threading
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path
logger = logging.getLogger(__name__)
DB_PATH = Path(__file__).parent.parent / "data" / "mail_bot.db"
_local = threading.local()
def init_db():
"""初始化数据库表"""
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
conn = _get_conn()
conn.executescript("""
CREATE TABLE IF NOT EXISTS processed_emails (
uid TEXT PRIMARY KEY,
account TEXT NOT NULL,
folder TEXT NOT NULL DEFAULT 'INBOX',
subject TEXT,
sender TEXT,
status TEXT NOT NULL DEFAULT 'fetched',
created_at TEXT NOT NULL,
sent_at TEXT
);
CREATE TABLE IF NOT EXISTS email_contexts (
msg_id INTEGER PRIMARY KEY,
data TEXT NOT NULL,
created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS conversations (
chat_id INTEGER PRIMARY KEY,
data TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_processed_status ON processed_emails(status);
CREATE INDEX IF NOT EXISTS idx_processed_account ON processed_emails(account);
""")
conn.commit()
logger.info("数据库初始化完成: %s", DB_PATH)
def _get_conn() -> sqlite3.Connection:
if not hasattr(_local, "conn") or _local.conn is None:
_local.conn = sqlite3.connect(str(DB_PATH), timeout=10)
_local.conn.execute("PRAGMA journal_mode=WAL")
_local.conn.execute("PRAGMA busy_timeout=5000")
_local.conn.row_factory = sqlite3.Row
return _local.conn
@contextmanager
def get_db():
conn = _get_conn()
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
# ── processed_emails ────────────────────────────────────
def is_uid_processed(uid: str) -> bool:
with get_db() as conn:
row = conn.execute(
"SELECT 1 FROM processed_emails WHERE uid=? AND status != 'failed'",
(uid,)
).fetchone()
return row is not None
def is_uids_processed(uids: list[str]) -> set[str]:
"""批量检查,返回已处理的 uid 集合"""
if not uids:
return set()
with get_db() as conn:
placeholders = ",".join("?" * len(uids))
rows = conn.execute(
f"SELECT uid FROM processed_emails WHERE uid IN ({placeholders}) AND status != 'failed'",
uids
).fetchall()
return {row["uid"] for row in rows}
def mark_fetched(uid: str, account: str, folder: str = "INBOX",
subject: str = "", sender: str = ""):
with get_db() as conn:
conn.execute(
"""INSERT OR IGNORE INTO processed_emails
(uid, account, folder, subject, sender, status, created_at)
VALUES (?, ?, ?, ?, ?, 'fetched', ?)""",
(uid, account, folder, subject, sender, datetime.now().isoformat())
)
def mark_processing(uid: str):
with get_db() as conn:
conn.execute(
"UPDATE processed_emails SET status='processing' WHERE uid=?",
(uid,)
)
def mark_sent(uid: str):
with get_db() as conn:
conn.execute(
"UPDATE processed_emails SET status='sent', sent_at=? WHERE uid=?",
(datetime.now().isoformat(), uid)
)
def mark_failed(uid: str):
with get_db() as conn:
conn.execute(
"UPDATE processed_emails SET status='failed' WHERE uid=?",
(uid,)
)
# ── email_contexts ──────────────────────────────────────
def save_email_context(msg_id: int, data: dict):
with get_db() as conn:
conn.execute(
"""INSERT OR REPLACE INTO email_contexts (msg_id, data, created_at)
VALUES (?, ?, ?)""",
(msg_id, json.dumps(data, ensure_ascii=False), datetime.now().isoformat())
)
def load_email_context(msg_id: int) -> dict | None:
with get_db() as conn:
row = conn.execute(
"SELECT data FROM email_contexts WHERE msg_id=?", (msg_id,)
).fetchone()
if row:
return json.loads(row["data"])
return None
# ── conversations ───────────────────────────────────────
def save_conversation(chat_id: int, data: dict):
with get_db() as conn:
conn.execute(
"""INSERT OR REPLACE INTO conversations (chat_id, data, updated_at)
VALUES (?, ?, ?)""",
(chat_id, json.dumps(data, ensure_ascii=False), datetime.now().isoformat())
)
def load_conversation(chat_id: int) -> dict | None:
with get_db() as conn:
row = conn.execute(
"SELECT data FROM conversations WHERE chat_id=?", (chat_id,)
).fetchone()
if row:
return json.loads(row["data"])
return None
def delete_conversation(chat_id: int):
with get_db() as conn:
conn.execute("DELETE FROM conversations WHERE chat_id=?", (chat_id,))

View File

@@ -58,17 +58,18 @@ def tg_send_and_mark(cfg: Config, info: dict):
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))
# 保存上下文用于回调
from src.tg_bot import _email_contexts
_email_contexts[thinking_msg_id] = {
save_email_context(thinking_msg_id, {
"summary_text": info["text"],
"summary_data": info["data"],
"original_body": info["original_body"][:10000],
"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", ""),
@@ -76,7 +77,7 @@ def tg_send_and_mark(cfg: Config, info: dict):
"sender": info["data"].get("sender", ""),
"subject": info["data"].get("subject", ""),
"can_reply": can_reply,
}
})
else:
# 兜底:没有思考中消息就正常发送
send_summary(cfg.telegram, cfg.telegram.chat_id,

View File

@@ -6,13 +6,10 @@ 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__)
_email_contexts: dict[int, dict] = {}
# chat_id -> conversation state
_conversations: dict[int, dict] = {}
# ── Public API ──────────────────────────────────────────
@@ -31,10 +28,10 @@ def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
"reply_markup": _summary_keyboard(can_reply, is_promotion, links),
})
msg_id = result["result"]["message_id"]
_email_contexts[msg_id] = {
save_email_context(msg_id, {
"summary_text": summary_text,
"summary_data": summary_data,
"original_body": original_body[:10000],
"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,
@@ -43,7 +40,7 @@ def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
"sender": summary_data.get("sender", ""),
"subject": summary_data.get("subject", ""),
"can_reply": can_reply,
}
})
return msg_id
@@ -270,13 +267,13 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
parts = data.split("|")
action = parts[0]
ctx = _email_contexts.get(msg_id)
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['original_body'][:3500])}"
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",
@@ -298,10 +295,10 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
tg_cfg, str(chat_id), "请输入你的回复内容:",
reply_markup=_cancel_keyboard(),
)
_conversations[chat_id] = {
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:
@@ -311,18 +308,18 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
logger.info(" AI 判定无需回复,隐藏")
return
_conversations[chat_id] = {
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 = _conversations.get(chat_id)
conv = load_conversation(chat_id)
if not conv or conv.get("summary_msg_id") != msg_id:
return
suggestions = conv.get("ai_suggestions", [])
@@ -332,7 +329,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
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)
delete_conversation(chat_id)
can_reply = ctx.get("can_reply", True)
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
@@ -341,7 +338,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
})
elif action == CALLBACK_REGEN and ctx:
conv = _conversations.get(chat_id)
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", [])))
@@ -369,14 +366,14 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
tg_cfg, str(chat_id), "请输入你对回复的提示/要求:",
reply_markup=_cancel_keyboard(),
)
conv = _conversations.get(chat_id)
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
logger.info(" 等待用户输入 AI 提示")
elif action == CALLBACK_CANCEL:
_conversations.pop(chat_id, None)
delete_conversation(chat_id)
if ctx:
can_reply = ctx.get("can_reply", True)
_tg_req(tg_cfg, "editMessageText", {
@@ -424,7 +421,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
chat_id = msg["chat"]["id"]
text = msg["text"].strip()
conv = _conversations.get(chat_id)
conv = load_conversation(chat_id)
if not conv:
return
@@ -437,9 +434,9 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
logger.info("TG 消息: state=%s text=%s", state, text[:60])
if state == "awaiting_reply":
ctx = _email_contexts.get(conv["summary_msg_id"])
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
_conversations.pop(chat_id, None)
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)
@@ -451,13 +448,13 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
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)
delete_conversation(chat_id)
logger.info(" 回复发送完成")
elif state == "awaiting_ai_hint":
ctx = _email_contexts.get(conv["summary_msg_id"])
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
_conversations.pop(chat_id, None)
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)
@@ -473,7 +470,7 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
)
except Exception as e:
send_text(tg_cfg, chat_id_str, f"生成失败: {e}")
_conversations.pop(chat_id, None)
delete_conversation(chat_id)
return
if new_suggestions: