earthjasonlin/ai-mail-bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: replace all in-memory state with SQLite database

Browse Source
This commit is contained in:
Zichao Lin
2026-07-14 14:49:23 +08:00
parent c3132f746a
commit 54c4bf00fb
4 changed files with 218 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
src/tg_bot.py
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,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: