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
This commit is contained in:
32
config.example.json
Normal file
32
config.example.json
Normal file
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"bot_token": "your_bot_token_here",
|
||||
"polling_interval": 30,
|
||||
"users": [
|
||||
{
|
||||
"chat_id": "your_telegram_chat_id",
|
||||
"ai": {
|
||||
"api_key": "your_deepseek_api_key",
|
||||
"model": "deepseek-v4-flash",
|
||||
"base_url": "https://api.deepseek.com"
|
||||
},
|
||||
"email_accounts": [
|
||||
{
|
||||
"username": "you@126.com",
|
||||
"password": "your_imap_password",
|
||||
"imap": {
|
||||
"server": "imap.126.com",
|
||||
"port": 993,
|
||||
"ssl": true,
|
||||
"starttls": false
|
||||
},
|
||||
"smtp": {
|
||||
"server": "smtp.126.com",
|
||||
"port": 465,
|
||||
"ssl": true,
|
||||
"starttls": false
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
70
main.py
70
main.py
@@ -6,12 +6,12 @@ import sys
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from src.config import load_config
|
||||
from src.config import load_config, Config, UserConfig
|
||||
from src.summarizer import poll_accounts, ai_process, tg_send_and_mark, _mail_label
|
||||
from src.tg_bot import poll_telegram, send_thinking
|
||||
from src.database import init_db, is_uid_processed, mark_fetched, mark_processing, mark_sent, mark_failed
|
||||
|
||||
# Thread-safe logging: all log records go through a single QueueListener thread
|
||||
# Thread-safe logging
|
||||
_log_queue = queue_module.Queue(-1)
|
||||
_queue_handler = logging.handlers.QueueHandler(_log_queue)
|
||||
_console_handler = logging.StreamHandler()
|
||||
@@ -31,9 +31,9 @@ logger = logging.getLogger("main")
|
||||
_running = True
|
||||
_shutdown_event = threading.Event()
|
||||
|
||||
# Queue 1: raw emails → AI processor
|
||||
# Queue 1: raw emails → AI processor (user, acct_idx, mail, tg_msg_id)
|
||||
_email_queue: queue_module.Queue = queue_module.Queue()
|
||||
# Queue 2: AI summaries → TG sender
|
||||
# Queue 2: AI summaries → TG sender (user, info_dict)
|
||||
_tg_queue: queue_module.Queue = queue_module.Queue()
|
||||
|
||||
|
||||
@@ -44,33 +44,36 @@ def _signal_handler(signum, frame):
|
||||
_shutdown_event.set()
|
||||
|
||||
|
||||
def _email_poller(cfg):
|
||||
def _email_poller(cfg: Config):
|
||||
logger.info("邮件轮询线程已启动")
|
||||
while _running:
|
||||
try:
|
||||
for acct_idx, mail in poll_accounts(cfg):
|
||||
for user in cfg.users:
|
||||
if not _running:
|
||||
return
|
||||
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,
|
||||
f"📥 *获取邮件中…*\n{_mail_label(mail)}")
|
||||
logger.info(" 获取邮件中消息已发送: msg_id=%d", tg_msg_id)
|
||||
except Exception as e:
|
||||
logger.error(" 获取邮件中消息发送失败: %s", e)
|
||||
tg_msg_id = 0
|
||||
_email_queue.put((acct_idx, mail, tg_msg_id))
|
||||
for acct_idx, mail in poll_accounts(user):
|
||||
if not _running:
|
||||
return
|
||||
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.bot_token, user.chat_id,
|
||||
f"📥 *获取邮件中…*\n{_mail_label(mail)}")
|
||||
logger.info(" [%s] 获取邮件中消息已发送: msg_id=%d", user.chat_id, tg_msg_id)
|
||||
except Exception as e:
|
||||
logger.error(" [%s] 获取邮件中消息发送失败: %s", user.chat_id, e)
|
||||
tg_msg_id = 0
|
||||
_email_queue.put((user, acct_idx, mail, tg_msg_id))
|
||||
except Exception as e:
|
||||
logger.error(f"邮件轮询线程异常: {e}", exc_info=True)
|
||||
time.sleep(5)
|
||||
if _running:
|
||||
_shutdown_event.wait(cfg.polling.interval_seconds)
|
||||
_shutdown_event.wait(cfg.polling_interval)
|
||||
|
||||
|
||||
def _ai_processor(cfg):
|
||||
def _ai_processor(cfg: Config):
|
||||
logger.info("AI 处理线程已启动 (并行 workers=10)")
|
||||
max_workers = 10
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
@@ -78,21 +81,21 @@ def _ai_processor(cfg):
|
||||
while _running:
|
||||
while len(pending_futures) < max_workers:
|
||||
try:
|
||||
acct_idx, mail, tg_msg_id = _email_queue.get_nowait()
|
||||
user, 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)
|
||||
future = executor.submit(ai_process, cfg, user, acct_idx, mail, tg_msg_id)
|
||||
pending_futures[future] = (user, acct_idx, mail)
|
||||
|
||||
completed = [f for f in pending_futures if f.done()]
|
||||
for future in completed:
|
||||
acct_idx, mail = pending_futures.pop(future)
|
||||
user, acct_idx, mail = pending_futures.pop(future)
|
||||
try:
|
||||
info = future.result()
|
||||
if info:
|
||||
_tg_queue.put(info)
|
||||
_tg_queue.put((user, info))
|
||||
except Exception as e:
|
||||
logger.error(f"AI 处理邮件失败: {e}", exc_info=True)
|
||||
uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
|
||||
@@ -101,24 +104,24 @@ def _ai_processor(cfg):
|
||||
time.sleep(0.1 if pending_futures else 0.5)
|
||||
|
||||
|
||||
def _tg_poller(cfg):
|
||||
def _tg_poller(cfg: Config):
|
||||
logger.info("TG 轮询线程已启动(按钮回调专用)")
|
||||
last_update_id = 0
|
||||
while _running:
|
||||
try:
|
||||
last_update_id = poll_telegram(cfg.telegram, cfg, last_update_id)
|
||||
last_update_id = poll_telegram(cfg.bot_token, cfg, last_update_id)
|
||||
except Exception as e:
|
||||
logger.error(f"TG 轮询错误: {e}", exc_info=True)
|
||||
if _running:
|
||||
_shutdown_event.wait(0.05)
|
||||
|
||||
|
||||
def _tg_sender(cfg):
|
||||
def _tg_sender(cfg: Config):
|
||||
logger.info("TG 发送线程已启动(邮件摘要专用)")
|
||||
while _running:
|
||||
try:
|
||||
info = _tg_queue.get_nowait()
|
||||
tg_send_and_mark(cfg, info)
|
||||
user, info = _tg_queue.get_nowait()
|
||||
tg_send_and_mark(cfg, user, info)
|
||||
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:
|
||||
@@ -139,10 +142,13 @@ def main():
|
||||
signal.signal(signal.SIGINT, _signal_handler)
|
||||
signal.signal(signal.SIGTERM, _signal_handler)
|
||||
|
||||
cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
|
||||
cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.json"
|
||||
cfg = load_config(cfg_path)
|
||||
init_db()
|
||||
logger.info(f"AI邮件摘要机器人已启动,轮询间隔: {cfg.polling.interval_seconds}s")
|
||||
|
||||
user_count = len(cfg.users)
|
||||
acct_count = sum(len(u.email_accounts) for u in cfg.users)
|
||||
logger.info(f"AI邮件摘要机器人已启动: {user_count} 用户, {acct_count} 邮箱, 轮询间隔: {cfg.polling_interval}s")
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=_email_poller, args=(cfg,), daemon=True),
|
||||
|
||||
@@ -1,25 +1,30 @@
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
import yaml
|
||||
import json
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImapConfig:
|
||||
server: str
|
||||
port: int = 993
|
||||
ssl: bool = True
|
||||
starttls: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class SmtpConfig:
|
||||
server: str
|
||||
port: int
|
||||
use_ssl: bool = True
|
||||
use_starttls: bool = False
|
||||
port: int = 465
|
||||
ssl: bool = True
|
||||
starttls: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmailAccount:
|
||||
imap_server: str
|
||||
imap_port: int
|
||||
username: str
|
||||
password: str
|
||||
use_ssl: bool = True
|
||||
use_starttls: bool = False
|
||||
smtp: Optional[SmtpConfig] = None
|
||||
imap: ImapConfig
|
||||
smtp: SmtpConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -30,40 +35,50 @@ class AIConfig:
|
||||
|
||||
|
||||
@dataclass
|
||||
class TelegramConfig:
|
||||
bot_token: str
|
||||
class UserConfig:
|
||||
chat_id: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class PollingConfig:
|
||||
interval_seconds: int
|
||||
ai: AIConfig
|
||||
email_accounts: List[EmailAccount] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
email_accounts: List[EmailAccount]
|
||||
ai: AIConfig
|
||||
telegram: TelegramConfig
|
||||
polling: PollingConfig
|
||||
bot_token: str
|
||||
polling_interval: int
|
||||
users: List[UserConfig] = field(default_factory=list)
|
||||
|
||||
|
||||
def load_config(path: str) -> Config:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
raw = yaml.safe_load(f)
|
||||
raw = json.load(f)
|
||||
|
||||
accounts = []
|
||||
for a in raw["email_accounts"]:
|
||||
if "smtp" in a and a["smtp"] is not None:
|
||||
a["smtp"] = SmtpConfig(**a["smtp"])
|
||||
accounts.append(EmailAccount(**a))
|
||||
ai = AIConfig(**raw["ai"])
|
||||
tg = TelegramConfig(**raw["telegram"])
|
||||
polling = PollingConfig(**raw["polling"])
|
||||
users = []
|
||||
for u in raw.get("users", []):
|
||||
email_accounts = []
|
||||
for a in u.get("email_accounts", []):
|
||||
email_accounts.append(EmailAccount(
|
||||
username=a["username"],
|
||||
password=a["password"],
|
||||
imap=ImapConfig(**a.get("imap", {})),
|
||||
smtp=SmtpConfig(**a.get("smtp", {})),
|
||||
))
|
||||
users.append(UserConfig(
|
||||
chat_id=str(u["chat_id"]),
|
||||
ai=AIConfig(**u["ai"]),
|
||||
email_accounts=email_accounts,
|
||||
))
|
||||
|
||||
return Config(
|
||||
email_accounts=accounts,
|
||||
ai=ai,
|
||||
telegram=tg,
|
||||
polling=polling,
|
||||
bot_token=raw["bot_token"],
|
||||
polling_interval=raw.get("polling_interval", 30),
|
||||
users=users,
|
||||
)
|
||||
|
||||
|
||||
def find_user_by_chat_id(cfg: Config, chat_id: int | str) -> Optional[UserConfig]:
|
||||
"""根据 chat_id 查找用户"""
|
||||
cid = str(chat_id)
|
||||
for user in cfg.users:
|
||||
if user.chat_id == cid:
|
||||
return user
|
||||
return None
|
||||
|
||||
@@ -2,7 +2,7 @@ import logging
|
||||
import re
|
||||
import threading
|
||||
from typing import Generator, Optional
|
||||
from src.config import Config
|
||||
from src.config import Config, UserConfig
|
||||
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
|
||||
@@ -21,31 +21,32 @@ def _mail_label(mail: Email) -> str:
|
||||
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):
|
||||
def poll_accounts(user: UserConfig, skip_uids: set[bytes] | None = None) -> Generator[tuple[int, Email], None, None]:
|
||||
"""轮询指定用户的所有邮箱"""
|
||||
for idx, acct in enumerate(user.email_accounts):
|
||||
try:
|
||||
logger.info(f"检查邮箱: {acct.username}")
|
||||
logger.info("检查邮箱: %s (user=%s)", acct.username, user.chat_id)
|
||||
emails = fetch_unseen_emails(acct, skip_uids=skip_uids)
|
||||
if emails:
|
||||
logger.info(f" 发现 {len(emails)} 封新邮件")
|
||||
logger.info(" 发现 %d 封新邮件", len(emails))
|
||||
for mail in emails:
|
||||
yield idx, mail
|
||||
except Exception as e:
|
||||
logger.error(f"轮询 {acct.username} 失败: {e}", exc_info=True)
|
||||
logger.error("轮询 %s 失败: %s", 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}")
|
||||
def ai_process(cfg: Config, user: UserConfig, acct_idx: int, mail: Email, tg_msg_id: int = 0) -> Optional[dict]:
|
||||
logger.info(" [%s] 正在摘要: %s", user.chat_id, 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}")
|
||||
edit_message(cfg.bot_token, user.chat_id, tg_msg_id, f"⏳ *等待处理…*\n{label}")
|
||||
else:
|
||||
tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, f"⏳ *等待处理…*\n{label}")
|
||||
tg_msg_id = send_thinking(cfg.bot_token, user.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)
|
||||
edit_message(cfg.bot_token, user.chat_id, tg_msg_id, f"🤖 *AI思考中…*\n{label}")
|
||||
summary = summarize_email(user.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
|
||||
@@ -66,19 +67,18 @@ def ai_process(cfg: Config, acct_idx: int, mail: Email, tg_msg_id: int = 0) -> O
|
||||
}
|
||||
|
||||
|
||||
def tg_send_and_mark(cfg: Config, info: dict):
|
||||
acct = cfg.email_accounts[info["acct_idx"]]
|
||||
def tg_send_and_mark(cfg: Config, user: UserConfig, info: dict):
|
||||
# 找到用户对应的邮箱账户
|
||||
acct = user.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,
|
||||
edit_message(cfg.bot_token, user.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"],
|
||||
@@ -94,19 +94,18 @@ def tg_send_and_mark(cfg: Config, info: dict):
|
||||
"can_reply": can_reply,
|
||||
})
|
||||
else:
|
||||
# 兜底:没有思考中消息就正常发送
|
||||
send_summary(cfg.telegram, cfg.telegram.chat_id,
|
||||
send_summary(cfg.bot_token, user.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 发送成功,标记已读已提交后台")
|
||||
logger.info(" [%s] TG 发送成功,标记已读已提交后台", user.chat_id)
|
||||
|
||||
142
src/tg_bot.py
142
src/tg_bot.py
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user