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:
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),
|
||||
|
||||
Reference in New Issue
Block a user