Zichao Lin
ded61e25c1
- Three independent threads: email poller, AI processor, TG worker - Two queues: _email_queue (raw emails→AI), _tg_queue (AI results→TG) - summarizer.py split: ai_process() (AI only) + tg_send_and_mark() (TG only) - AI retries no longer affect TG polling or email polling
121 lines
3.2 KiB
Python
121 lines
3.2 KiB
Python
import logging
|
||
import queue
|
||
import signal
|
||
import sys
|
||
import threading
|
||
import time
|
||
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
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||
datefmt="%Y-%m-%d %H:%M:%S",
|
||
)
|
||
logger = logging.getLogger("main")
|
||
|
||
_running = True
|
||
_pending_lock = threading.Lock()
|
||
_pending_uids: set[bytes] = set()
|
||
|
||
# Queue 1: raw emails → AI processor
|
||
_email_queue: queue.Queue = queue.Queue()
|
||
# Queue 2: AI results → TG sender
|
||
_tg_queue: queue.Queue = queue.Queue()
|
||
|
||
|
||
def _signal_handler(signum, frame):
|
||
global _running
|
||
logger.info("收到退出信号，正在停止...")
|
||
_running = False
|
||
|
||
|
||
def _email_poller(cfg):
|
||
logger.info("邮件轮询线程已启动")
|
||
while _running:
|
||
for acct_idx, mail in poll_accounts(cfg):
|
||
if not _running:
|
||
return
|
||
with _pending_lock:
|
||
if mail.uid not in _pending_uids:
|
||
_pending_uids.add(mail.uid)
|
||
_email_queue.put((acct_idx, mail))
|
||
if _running:
|
||
for _ in range(cfg.polling.interval_seconds):
|
||
if not _running:
|
||
return
|
||
time.sleep(1)
|
||
|
||
|
||
def _ai_processor(cfg):
|
||
logger.info("AI 处理线程已启动")
|
||
while _running:
|
||
try:
|
||
acct_idx, mail = _email_queue.get(timeout=1)
|
||
except queue.Empty:
|
||
continue
|
||
try:
|
||
info = ai_process(cfg, acct_idx, mail)
|
||
if info:
|
||
_tg_queue.put(info)
|
||
except Exception as e:
|
||
logger.error(f"AI 处理邮件失败: {e}", exc_info=True)
|
||
finally:
|
||
with _pending_lock:
|
||
_pending_uids.discard(mail.uid)
|
||
|
||
|
||
def _tg_worker(cfg):
|
||
logger.info("TG 线程已启动")
|
||
last_update_id = 0
|
||
while _running:
|
||
try:
|
||
last_update_id = poll_telegram(cfg.telegram, cfg, last_update_id)
|
||
except Exception as e:
|
||
logger.error(f"TG 轮询错误: {e}", exc_info=True)
|
||
|
||
try:
|
||
info = _tg_queue.get_nowait()
|
||
tg_send_and_mark(cfg, info)
|
||
except queue.Empty:
|
||
pass
|
||
except Exception as e:
|
||
logger.error(f"TG 发送失败: {e}", exc_info=True)
|
||
|
||
if _running:
|
||
time.sleep(1)
|
||
|
||
|
||
def main():
|
||
global _running
|
||
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 = load_config(cfg_path)
|
||
logger.info(f"AI邮件摘要机器人已启动，轮询间隔: {cfg.polling.interval_seconds}s")
|
||
|
||
threads = [
|
||
threading.Thread(target=_email_poller, args=(cfg,), daemon=True),
|
||
threading.Thread(target=_ai_processor, args=(cfg,), daemon=True),
|
||
threading.Thread(target=_tg_worker, args=(cfg,), daemon=True),
|
||
]
|
||
|
||
for t in threads:
|
||
t.start()
|
||
|
||
try:
|
||
while _running:
|
||
time.sleep(1)
|
||
except KeyboardInterrupt:
|
||
_running = False
|
||
|
||
for t in threads:
|
||
t.join(timeout=5)
|
||
logger.info("机器人已停止")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|