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

feat: multi-user support - Config format: JSON with bot_token (global), users array - Each user has chat_id, AI config, email_accounts - Each email account has separate IMAP/SMTP with ssl/starttls - Config module: new dataclasses (UserConfig, ImapConfig, SmtpConfig) - find_user_by_chat_id() for callback dispatch - Main: email poller iterates all users, queues carry user object - Summarizer: accepts UserConfig, uses user-specific AI and chat_id - TG bot: all functions use bot_token string instead of TelegramConfig - Callback/message handlers look up user by chat_id - AI regen/hint threads use find_user_by_chat_id for AI config - Config example: config.example.json - Tests: 3/3 passing

Browse Source
This commit is contained in:
Zichao Lin
2026-07-17 22:05:59 +08:00
parent 9906f1675f
commit f2b3793d4c
5 changed files with 213 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
src/tg_bot.py
View File

@@ -4,7 +4,7 @@ import threading
from concurrent.futures import ThreadPoolExecutor
from email.utils import parseaddr
import requests
from src.config import TelegramConfig, Config
from src.config import Config, UserConfig, find_user_by_chat_id
from src.ai_client import generate_more_replies, expand_promo_detail
from src.smtp_client import send_reply
from src.retry import retry
@@ -18,14 +18,14 @@ _ai_pool = ThreadPoolExecutor(max_workers=3, thread_name_prefix="ai_regen")
# ── Public API ──────────────────────────────────────────
@retry()
def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
def send_summary(bot_token: str, chat_id: str, summary_text: str,
summary_data: dict, original_body: str, account_idx: int,
original_sender: str = "", original_recipient: str = "",
original_reply_to: str = "", account_email: str = "") -> int:
can_reply = summary_data.get("can_reply", True)
is_promotion = summary_data.get("category") == "promotion"
links = summary_data.get("links", [])
result = _tg_req(tg_cfg, "sendMessage", {
result = _tg_req(bot_token, "sendMessage", {
"chat_id": chat_id,
"text": summary_text,
"parse_mode": "MarkdownV2",
@@ -49,24 +49,24 @@ def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
@retry()
def send_text(tg_cfg: TelegramConfig, chat_id: str, text: str,
def send_text(bot_token: str, chat_id: str, text: str,
reply_markup: dict = None) -> int:
payload = {"chat_id": chat_id, "text": text}
if reply_markup:
payload["reply_markup"] = reply_markup
result = _tg_req(tg_cfg, "sendMessage", payload)
result = _tg_req(bot_token, "sendMessage", payload)
return result["result"]["message_id"]
def send_thinking(tg_cfg: TelegramConfig, chat_id: str, text: str = "💭 思考中…") -> int:
result = _tg_req(tg_cfg, "sendMessage", {
def send_thinking(bot_token: str, chat_id: str, text: str = "💭 思考中…") -> int:
result = _tg_req(bot_token, "sendMessage", {
"chat_id": chat_id,
"text": text,
})
return result["result"]["message_id"]
def edit_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int,
def edit_message(bot_token: str, chat_id: str, msg_id: int,
text: str, reply_markup: dict = None):
payload = {
"chat_id": chat_id, "message_id": msg_id,
@@ -74,22 +74,22 @@ def edit_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int,
}
if reply_markup:
payload["reply_markup"] = reply_markup
_tg_req(tg_cfg, "editMessageText", payload)
_tg_req(bot_token, "editMessageText", payload)
def delete_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int):
def delete_message(bot_token: str, chat_id: str, msg_id: int):
try:
_tg_req(tg_cfg, "deleteMessage", {
_tg_req(bot_token, "deleteMessage", {
"chat_id": chat_id, "message_id": msg_id,
})
except Exception:
pass
def poll_telegram(tg_cfg: TelegramConfig, cfg: Config, last_update_id: int) -> int:
def poll_telegram(bot_token: str, cfg: Config, last_update_id: int) -> int:
try:
resp = requests.get(
f"https://api.telegram.org/bot{tg_cfg.bot_token}/getUpdates",
f"https://api.telegram.org/bot{bot_token}/getUpdates",
params={"offset": last_update_id, "timeout": 1},
timeout=10,
)
@@ -100,7 +100,7 @@ def poll_telegram(tg_cfg: TelegramConfig, cfg: Config, last_update_id: int) -> i
for update in data.get("result", []):
uid = update["update_id"]
try:
_handle_update(tg_cfg, cfg, update)
_handle_update(bot_token, cfg, update)
except Exception as e:
logger.warning("处理 update %d 失败: %s", uid, e)
last_update_id = uid + 1
@@ -247,8 +247,8 @@ def _ai_reply_keyboard(msg_id: int, suggestions: list) -> dict:
# ── Telegram API ────────────────────────────────────────
def _tg_req(tg_cfg: TelegramConfig, method: str, payload: dict) -> dict:
url = f"https://api.telegram.org/bot{tg_cfg.bot_token}/{method}"
def _tg_req(bot_token: str, method: str, payload: dict) -> dict:
url = f"https://api.telegram.org/bot{bot_token}/{method}"
logger.info("Telegram API %s: %s", method, {k: (v if k != "reply_markup" else "...") for k, v in payload.items()})
r = requests.post(url, json=payload, timeout=30)
if not r.ok:
@@ -262,35 +262,41 @@ def _tg_req(tg_cfg: TelegramConfig, method: str, payload: dict) -> dict:
# ── Update handling ─────────────────────────────────────
def _do_regen(cfg: Config, chat_id: int, msg_id: int, ctx: dict, conv: dict):
def _do_regen(bot_token: str, cfg: Config, chat_id: int, msg_id: int, ctx: dict, conv: dict):
"""在 AI 线程池中执行换一批，不阻塞按钮轮询"""
user = find_user_by_chat_id(cfg, chat_id)
if not user:
return
style = conv.get("reply_style", "short")
try:
new_suggestions = generate_more_replies(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
user.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []), style=style,
)
except Exception as e:
send_text(cfg.telegram, str(chat_id), f"生成失败: {e}")
send_text(bot_token, str(chat_id), f"生成失败: {e}")
return
if new_suggestions:
conv["ai_suggestions"] = new_suggestions
conv["all_suggestions"].extend(new_suggestions)
save_conversation(chat_id, conv)
logger.info("[ai_regen] 新生成 %d 条回复 (style=%s)", len(new_suggestions), style)
_show_ai_suggestions(cfg.telegram, chat_id, msg_id, new_suggestions)
_show_ai_suggestions(bot_token, chat_id, msg_id, new_suggestions)
def _do_ai_hint(cfg: Config, chat_id: int, conv: dict, ctx: dict, hint: str):
def _do_ai_hint(bot_token: str, cfg: Config, chat_id: int, conv: dict, ctx: dict, hint: str):
"""在 AI 线程池中根据用户提示生成回复"""
user = find_user_by_chat_id(cfg, chat_id)
if not user:
return
style = conv.get("reply_style", "short")
try:
new_suggestions = generate_more_replies(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
user.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []), user_hint=hint, style=style,
)
except Exception as e:
send_text(cfg.telegram, str(chat_id), f"生成失败: {e}")
send_text(bot_token, str(chat_id), f"生成失败: {e}")
delete_conversation(chat_id)
return
if new_suggestions:
@@ -300,24 +306,24 @@ def _do_ai_hint(cfg: Config, chat_id: int, conv: dict, ctx: dict, hint: str):
conv.pop("prompt_msg_id", None)
save_conversation(chat_id, conv)
logger.info("[ai_hint] 根据提示生成 %d 条新回复", len(new_suggestions))
_show_ai_suggestions(cfg.telegram, chat_id, conv["summary_msg_id"], new_suggestions)
_show_ai_suggestions(bot_token, chat_id, conv["summary_msg_id"], new_suggestions)
def _handle_update(tg_cfg: TelegramConfig, cfg: Config, update: dict):
def _handle_update(bot_token: str, cfg: Config, update: dict):
if "callback_query" in update:
_handle_callback(tg_cfg, cfg, update["callback_query"])
_handle_callback(bot_token, cfg, update["callback_query"])
elif "message" in update:
_handle_message(tg_cfg, cfg, update["message"])
_handle_message(bot_token, cfg, update["message"])
def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
def _handle_callback(bot_token: str, cfg: Config, cb: dict):
chat_id = cb["message"]["chat"]["id"]
msg_id = cb["message"]["message_id"]
data = cb["data"]
cb_id = cb["id"]
try:
_tg_req(tg_cfg, "answerCallbackQuery", {"callback_query_id": cb_id})
_tg_req(bot_token, "answerCallbackQuery", {"callback_query_id": cb_id})
except Exception:
pass
@@ -331,7 +337,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
if action == CALLBACK_VIEW_ORIG and ctx:
can_reply = ctx.get("can_reply", True)
text = f"*📄 原文 \\- {_escape(ctx['subject'])}*\n\n{_escape(ctx.get('plain_text', ctx['original_body']))}"
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": text, "parse_mode": "MarkdownV2",
"reply_markup": _orig_keyboard(can_reply),
@@ -342,7 +348,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
can_reply = ctx.get("can_reply", True)
is_promotion = ctx.get("summary_data", {}).get("category") == "promotion"
links = ctx.get("summary_data", {}).get("links", [])
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply, is_promotion, links),
@@ -351,7 +357,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
elif action == CALLBACK_REPLY and ctx:
prompt_msg_id = send_text(
tg_cfg, str(chat_id), "请输入你的回复内容：",
bot_token, str(chat_id), "请输入你的回复内容：",
reply_markup=_cancel_keyboard(),
)
save_conversation(chat_id, {
@@ -364,7 +370,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
style = "full" if action == CALLBACK_AI_REPLY_FULL else "short"
# 总是用对应风格重新生成，不使用摘要时的通用缓存
logger.info(" 重新生成 AI 回复 (style=%s)", style)
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
@@ -374,7 +380,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
conv["reply_style"] = style
conv["all_suggestions"] = []
save_conversation(chat_id, conv)
_ai_pool.submit(_do_regen, cfg, chat_id, msg_id, ctx, conv)
_ai_pool.submit(_do_regen, bot_token, cfg, chat_id, msg_id, ctx, conv)
elif action == CALLBACK_SEND_REPLY and ctx:
idx = int(parts[2])
@@ -386,13 +392,13 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
return
sel_text = suggestions[idx]["text"]
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"✅ 回复已发送")
_do_send_reply(bot_token, cfg, chat_id, msg_id, ctx, sel_text)
send_text(bot_token, str(chat_id), f"✅ 回复已发送")
delete_conversation(chat_id)
can_reply = ctx.get("can_reply", True)
is_promotion = ctx.get("summary_data", {}).get("category") == "promotion"
links = ctx.get("summary_data", {}).get("links", [])
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply, is_promotion, links),
@@ -412,7 +418,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
# 用代码块显示，点击可复制（MarkdownV2）
escaped = sel_text.replace("\\", "\\\\").replace("`", "\\`")
_tg_req(tg_cfg, "sendMessage", {
_tg_req(bot_token, "sendMessage", {
"chat_id": str(chat_id),
"text": f"*建议回复（点击代码块复制）:*\n`{escaped}`",
"parse_mode": "MarkdownV2",
@@ -420,7 +426,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
# 进入 awaiting_reply 状态，用户可编辑后发送
prompt_msg_id = send_text(
tg_cfg, str(chat_id),
bot_token, str(chat_id),
"请编辑后发送（或直接回复附件）：",
reply_markup=_cancel_keyboard(),
)
@@ -437,16 +443,16 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
return
logger.info(" 换一批 (已有 %d 条历史)", len(conv.get("all_suggestions", [])))
# AI 处理中
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
# 丢到 AI 线程池处理，不阻塞按钮轮询
_ai_pool.submit(_do_regen, cfg, chat_id, msg_id, ctx, conv)
_ai_pool.submit(_do_regen, bot_token, cfg, chat_id, msg_id, ctx, conv)
elif action == CALLBACK_HINT and ctx:
prompt_msg_id = send_text(
tg_cfg, str(chat_id), "请输入你对回复的提示/要求：",
bot_token, str(chat_id), "请输入你对回复的提示/要求：",
reply_markup=_cancel_keyboard(),
)
conv = load_conversation(chat_id)
@@ -465,18 +471,18 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
can_reply = ctx.get("can_reply", True)
is_promotion = ctx.get("summary_data", {}).get("category") == "promotion"
links = ctx.get("summary_data", {}).get("links", [])
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply, is_promotion, links),
})
logger.info(" 取消, 恢复摘要视图")
else:
delete_message(tg_cfg, str(chat_id), msg_id)
delete_message(bot_token, str(chat_id), msg_id)
logger.info(" 取消, 删除提示消息")
elif action == CALLBACK_PROMO_DETAIL and ctx:
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
@@ -485,7 +491,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
)
except Exception as e:
send_text(tg_cfg, str(chat_id), f"获取详情失败: {e}")
send_text(bot_token, str(chat_id), f"获取详情失败: {e}")
return
summary = ctx["summary_data"].get("summary", "")
detail_lines = "\n".join(f"\\- {_escape(d)}" for d in details) if details else "暂无更多详情"
@@ -496,7 +502,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
f"\n{_escape(summary)}\n"
f"\n*📋 详情:*\n{detail_lines}"
)
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": expanded, "parse_mode": "MarkdownV2",
"reply_markup": {"inline_keyboard": [
@@ -507,7 +513,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
logger.info(" 显示推广详情")
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
def _handle_message(bot_token: str, cfg: Config, msg: dict):
chat_id = msg["chat"]["id"]
text = msg.get("text") or msg.get("caption") or ""
text = text.strip()
@@ -528,24 +534,24 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
delete_conversation(chat_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
delete_message(bot_token, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
delete_message(bot_token, chat_id_str, prompt_msg_id)
logger.info(" 上下文已丢失，清理")
return
delete_message(tg_cfg, chat_id_str, user_msg_id)
delete_message(bot_token, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
delete_message(bot_token, chat_id_str, prompt_msg_id)
# 合并累积的附件
all_attachments = conv.get("pending_attachments", [])
all_attachments.extend(media)
_do_send_reply(tg_cfg, cfg, chat_id, conv["summary_msg_id"], ctx, text,
_do_send_reply(bot_token, cfg, chat_id, conv["summary_msg_id"], ctx, text,
attachments=all_attachments or None)
attach_info = f" (+{len(all_attachments)}个附件)" if all_attachments else ""
send_text(tg_cfg, chat_id_str, f"✅ 回复已发送{attach_info}")
send_text(bot_token, chat_id_str, f"✅ 回复已发送{attach_info}")
delete_conversation(chat_id)
logger.info(" 回复发送完成")
@@ -553,50 +559,50 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
delete_conversation(chat_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
delete_message(bot_token, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
delete_message(bot_token, chat_id_str, prompt_msg_id)
return
delete_message(tg_cfg, chat_id_str, user_msg_id)
delete_message(bot_token, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
delete_message(bot_token, chat_id_str, prompt_msg_id)
# 更新消息为 AI 思考中
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": conv["summary_msg_id"],
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
# 丢到 AI 线程池处理
_ai_pool.submit(_do_ai_hint, cfg, chat_id, conv, ctx, text)
_ai_pool.submit(_do_ai_hint, bot_token, cfg, chat_id, conv, ctx, text)
def _show_ai_suggestions(tg_cfg: TelegramConfig, chat_id: int, msg_id: int,
def _show_ai_suggestions(bot_token: str, chat_id: int, msg_id: int,
suggestions: list):
text = "*AI 建议回复，请选择：*\n\n"
for i, s in enumerate(suggestions, 1):
text += f"{i}\\. {_escape(s['text'])}\n"
_tg_req(tg_cfg, "editMessageText", {
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": text, "parse_mode": "MarkdownV2",
"reply_markup": _ai_reply_keyboard(msg_id, suggestions),
})
def _download_tg_file(tg_cfg: TelegramConfig, file_id: str, filename: str) -> str:
def _download_tg_file(bot_token: str, file_id: str, filename: str) -> str:
"""从 Telegram 下载文件到本地 media/ 目录，返回本地路径"""
import os
from pathlib import Path
media_dir = Path("media")
media_dir.mkdir(exist_ok=True)
resp = requests.get(
f"https://api.telegram.org/bot{tg_cfg.bot_token}/getFile",
f"https://api.telegram.org/bot{bot_token}/getFile",
params={"file_id": file_id}, timeout=30,
)
resp.raise_for_status()
file_path = resp.json()["result"]["file_path"]
dl = requests.get(
f"https://api.telegram.org/file/bot{tg_cfg.bot_token}/{file_path}",
f"https://api.telegram.org/file/bot{bot_token}/{file_path}",
timeout=60,
)
dl.raise_for_status()
@@ -650,12 +656,12 @@ def _extract_media(msg: dict) -> list[dict]:
return attachments
def _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
def _do_send_reply(bot_token: str, cfg: Config,
chat_id: int, msg_id: int, ctx: dict, reply_text: str,
attachments: list[dict] | None = None):
acct = cfg.email_accounts[ctx["account_idx"]]
if not acct.smtp:
send_text(tg_cfg, str(chat_id), "❌ 该邮箱未配置 SMTP，无法发送回复。")
send_text(bot_token, str(chat_id), "❌ 该邮箱未配置 SMTP，无法发送回复。")
return
to_addr = parseaddr(ctx.get("original_reply_to", ""))[1]
if not to_addr:
@@ -671,7 +677,7 @@ def _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
if attachments:
for att in attachments:
try:
path = _download_tg_file(tg_cfg, att["file_id"], att["filename"])
path = _download_tg_file(bot_token, att["file_id"], att["filename"])
local_files.append((path, att["mime_type"]))
except Exception as e:
logger.error(" 附件下载失败: %s - %s", att["filename"], e)