From 54c4bf00fb0bd54c2315b2a3e4cf5070c8533418 Mon Sep 17 00:00:00 2001 From: Zichao Lin Date: Tue, 14 Jul 2026 14:49:23 +0800 Subject: [PATCH] refactor: replace all in-memory state with SQLite database --- main.py | 47 ++++++------- src/database.py | 173 ++++++++++++++++++++++++++++++++++++++++++++++ src/summarizer.py | 6 +- src/tg_bot.py | 43 ++++++------ 4 files changed, 218 insertions(+), 51 deletions(-) create mode 100644 src/database.py diff --git a/main.py b/main.py index f6286b2..f501cac 100644 --- a/main.py +++ b/main.py @@ -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 diff --git a/src/database.py b/src/database.py new file mode 100644 index 0000000..ffe55c8 --- /dev/null +++ b/src/database.py @@ -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,)) diff --git a/src/summarizer.py b/src/summarizer.py index 69068c3..f2561a8 100644 --- a/src/summarizer.py +++ b/src/summarizer.py @@ -58,14 +58,14 @@ 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"], @@ -77,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, diff --git a/src/tg_bot.py b/src/tg_bot.py index 8682955..947ac90 100644 --- a/src/tg_bot.py +++ b/src/tg_bot.py @@ -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,7 +28,7 @@ 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, @@ -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,7 +267,7 @@ 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) @@ -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: