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

Compare commits

base: earthjasonlin/ai-mail-bot:main
earthjasonlin/ai-mail-bot:main
..
compare: earthjasonlin/ai-mail-bot:bedf5b887c6439f10db6ff7a926be53fd93c0bcf
earthjasonlin/ai-mail-bot:main
0 Commits
main .. bedf5b887c
11 changed files with 279 additions and 619 deletions
Expand all files Collapse all files

No files matched your search

.gitignore
+1 -2
View File
@@ -1,6 +1,5 @@
.venv/
__pycache__/
*.pyc
config.json
config.yaml
data/
*_TODO.md
config.example.json
-32
View File
@@ -1,32 +0,0 @@
{
"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
}
}
]
}
]
}
config.yaml.example
+32
View File
@@ -0,0 +1,32 @@
email_accounts:
- imap_server: "imap.gmail.com"
imap_port: 993
username: "your_email@gmail.com"
password: "your_app_password"
use_ssl: true
smtp:
server: "smtp.gmail.com"
port: 465
use_ssl: true
- imap_server: "imap.126.com"
imap_port: 993
username: "your_email@126.com"
password: "your_auth_code"
use_ssl: true
smtp:
server: "smtp.126.com"
port: 465
use_ssl: true
ai:
api_key: "sk-your_deepseek_api_key"
model: "deepseek-chat"
base_url: "https://api.deepseek.com"
telegram:
bot_token: "1234567890:ABCdefGHIjklmNOPqrstUVwxyz"
chat_id: "123456789"
polling:
interval_seconds: 60
main.py
+32 -38
View File
@@ -6,12 +6,12 @@ import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from src.config import load_config, Config, UserConfig
from src.config import load_config
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
# Thread-safe logging: all log records go through a single QueueListener thread
_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 (user, acct_idx, mail, tg_msg_id)
# Queue 1: raw emails → AI processor
_email_queue: queue_module.Queue = queue_module.Queue()
# Queue 2: AI summaries → TG sender (user, info_dict)
# Queue 2: AI summaries → TG sender
_tg_queue: queue_module.Queue = queue_module.Queue()
@@ -44,36 +44,33 @@ def _signal_handler(signum, frame):
_shutdown_event.set()
def _email_poller(cfg: Config):
def _email_poller(cfg):
logger.info("邮件轮询线程已启动")
while _running:
try:
for user in cfg.users:
for acct_idx, mail in poll_accounts(cfg):
if not _running:
return
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))
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))
except Exception as e:
logger.error(f"邮件轮询线程异常: {e}", exc_info=True)
time.sleep(5)
if _running:
_shutdown_event.wait(cfg.polling_interval)
_shutdown_event.wait(cfg.polling.interval_seconds)
def _ai_processor(cfg: Config):
def _ai_processor(cfg):
logger.info("AI 处理线程已启动 (并行 workers=10)")
max_workers = 10
with ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -81,21 +78,21 @@ def _ai_processor(cfg: Config):
while _running:
while len(pending_futures) < max_workers:
try:
user, acct_idx, mail, tg_msg_id = _email_queue.get_nowait()
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, user, acct_idx, mail, tg_msg_id)
pending_futures[future] = (user, acct_idx, mail)
future = executor.submit(ai_process, cfg, acct_idx, mail, tg_msg_id)
pending_futures[future] = (acct_idx, mail)
completed = [f for f in pending_futures if f.done()]
for future in completed:
user, acct_idx, mail = pending_futures.pop(future)
acct_idx, mail = pending_futures.pop(future)
try:
info = future.result()
if info:
_tg_queue.put((user, info))
_tg_queue.put(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)
@@ -104,24 +101,24 @@ def _ai_processor(cfg: Config):
time.sleep(0.1 if pending_futures else 0.5)
def _tg_poller(cfg: Config):
def _tg_poller(cfg):
logger.info("TG 轮询线程已启动（按钮回调专用）")
last_update_id = 0
while _running:
try:
last_update_id = poll_telegram(cfg.bot_token, cfg, last_update_id)
last_update_id = poll_telegram(cfg.telegram, 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: Config):
def _tg_sender(cfg):
logger.info("TG 发送线程已启动（邮件摘要专用）")
while _running:
try:
user, info = _tg_queue.get_nowait()
tg_send_and_mark(cfg, user, info)
info = _tg_queue.get_nowait()
tg_send_and_mark(cfg, 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:
@@ -142,13 +139,10 @@ 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.json"
cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
cfg = load_config(cfg_path)
init_db()
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")
logger.info(f"AI邮件摘要机器人已启动，轮询间隔: {cfg.polling.interval_seconds}s")
threads = [
threading.Thread(target=_email_poller, args=(cfg,), daemon=True),
src/ai_client.py
+19 -27
View File
@@ -8,7 +8,7 @@ from src.retry import retry
logger = logging.getLogger(__name__)
SYSTEM_PROMPT = """你是一个邮件摘要助手。请分析邮件内容并以 JSON 格式返回结构化摘要。所有输出（subject、summary 等）必须使用简体中文。
SYSTEM_PROMPT = """你是一个邮件摘要助手。请分析邮件内容并以 JSON 格式返回结构化摘要。所有输出（subject、summary 等）必须使用简体中文，但 reply_suggestions（智能回复）请使用与邮件相同的语言，以便用户直接回复。
返回格式必须严格遵循以下 JSON schema：
{
@@ -21,7 +21,12 @@ SYSTEM_PROMPT = """你是一个邮件摘要助手。请分析邮件内容并以
"verification_code": "",
"app_name": "",
"links": [{"text": "按钮文字", "url": "https://..."}],
"can_reply": true/false
"can_reply": true/false,
"reply_suggestions": [
{"text": "回复内容", "is_simple": true/false},
{"text": "回复内容", "is_simple": true/false},
{"text": "回复内容", "is_simple": true/false}
]
}
其中：
@@ -29,33 +34,26 @@ SYSTEM_PROMPT = """你是一个邮件摘要助手。请分析邮件内容并以
- category: 邮件类型。推广邮件（营销、促销、广告、 newsletter 推广等）标记为 "promotion"；验证码邮件标记为 "notification"；普通邮件标记为 "normal"
- verification_code: 仅当 category 为 "notification" 且邮件包含验证码时填写，提取完整的验证码字符串（纯数字、字母数字混合等均可，如「123456」「abc-789」「701383」）。非验证码邮件留空字符串
- app_name: 仅当 category 为 "notification" 且邮件包含验证码时填写，填写发送验证码的应用/平台名称（如「GitHub」「微信」「Kraken」）。非验证码邮件留空字符串
- links: 提取用户需要点击的链接。其中退订/取消订阅链接按下方「退订链接提取规则」强制执行，其余仅提取邮件明确要求用户点击才能完成某操作的链接（如确认邮箱、重置密码、查看账单、审批请求、查看通知详情等）。不要提取推广性质的链接（如下载App、注册领优惠、查看更多商品等营销链接）、追踪像素、邮件签名中的社交链接。每条包含简短描述文字和完整 URL。最多 3 条。没有符合条件的链接时为空数组
- 退订链接提取规则（最高优先级，仅适用于 category 为 "promotion" 的邮件）：
* 若邮件头提供了 List-Unsubscribe 头（用户消息中「List-Unsubscribe:」行），该链接是权威退订入口，优先使用它作为退订链接（除非它明显是 mailto: 形式而正文里能找到更好的 http 链接，此时两者取其一即可）
* 推广邮件必须在 links 中输出退订链接，且放在 links 数组的第一位；只有邮件正文确实不存在任何退订入口时，links 才可以为空
* 退订链接通常位于邮件最底部/页脚区域，表现形式多样：文字为 unsubscribe、Unsubscribe、退订、取消订阅、点击这里退订、manage preferences、email preferences、update preferences、subscription settings、opt out 等；也可能是图片链接，或包裹在 HTML 标签中的 <a href="...">，必须提取 href 属性里的完整 URL
* 同一封邮件出现多个退订相关链接时，只提取最主要的那个（优先取正文最后出现的完整 unsubscribe 链接）
* 退订链接固定占用 1 个位置，其余链接最多 2 条，不得挤掉退订链接
* 严禁编造退订 URL，确实找不到时 links 输出空数组
- links: 仅提取以下两类链接：1）邮件明确要求用户点击才能完成某操作的链接（如确认邮箱、重置密码、查看账单、审批请求、查看通知详情等）；2）退订/取消订阅链接。不要提取推广性质的链接（如下载App、注册领优惠、查看更多商品等营销链接）、追踪像素、邮件签名中的社交链接。每条包含简短描述文字和完整 URL。最多 3 条。没有符合条件的链接时为空数组
- 退订链接提取规则：推广邮件必须提取退订链接。退订链接通常在邮件最底部，包含以下关键词之一：unsubscribe、退订、取消订阅、opt out、manage preferences、email preferences。即使链接被包裹在 HTML 标签中，也要提取 href 属性里的完整 URL。如果邮件底部同时有多个退订相关链接，只提取最主要的那个
- plain_text: 将邮件全文转换为可读的纯文本版本。去除所有HTML标签、CSS样式、脚本代码，保留完整的文字内容和原始段落结构。不删减、不修改、不总结任何内容，完整保留原文所有信息
- summary 要简洁，把需要知道的信息浓缩成一句话，不需要分条列出
- 当 category 为 "promotion" 时，summary 直接写成一句话：「[发送方名称] 在推广 [推广的产品/服务/活动]」，例如「腾讯云在推广双十一云服务器折扣」。此时 can_reply 为 false
- can_reply 为 false 表示无需回复或不适合建议回复
- 当 category 为 "promotion" 时，summary 直接写成一句话：「[发送方名称] 在推广 [推广的产品/服务/活动]」，例如「腾讯云在推广双十一云服务器折扣」。此时 can_reply 为 false，reply_suggestions 为空数组
- can_reply 为 false 表示无需回复或不适合建议回复（此时忽略 reply_suggestions）
- is_simple 为 true 表示纯确认性回复（如"好的""收到"），用户选择后可直发
- is_simple 为 false 表示涉及实质内容，需要用户确认再发送
- 每条回复控制在 50 字以内
只返回 JSON，不要包含任何其他文字。"""
@retry()
def summarize_email(ai_cfg: AIConfig, recipient: str, subject: str, sender: str, body: str,
account: str = "", list_unsubscribe: str = "") -> dict[str, Any]:
def summarize_email(ai_cfg: AIConfig, recipient: str, subject: str, sender: str, body: str, account: str = "") -> dict[str, Any]:
body_preview = body[:80].replace("\n", " ")
logger.info("AI 摘要请求: sender=%s subject=%s body_preview=%s", sender, subject, body_preview)
content = f"收件人: {recipient}\n发件人: {sender}\n主题: {subject}\n账号: {account}\n正文:\n{body}"
if list_unsubscribe:
content += f"\n\nList-Unsubscribe: {list_unsubscribe}"
payload = {
"model": ai_cfg.model,
"thinking": {"type": "disabled"},
"messages": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": content},
@@ -78,7 +76,6 @@ def expand_promo_detail(ai_cfg: AIConfig, sender: str, subject: str, body: str)
content = f"发件人: {sender}\n主题: {subject}\n正文:\n{body}"
payload = {
"model": ai_cfg.model,
"thinking": {"type": "disabled"},
"messages": [
{"role": "system", "content": PROMO_DETAIL_PROMPT},
{"role": "user", "content": content},
@@ -104,8 +101,7 @@ MORE_REPLIES_PROMPT_SHORT = """你是一个邮件回复助手。根据以下邮
{"text": "回复内容"},
{"text": "回复内容"},
{"text": "回复内容"}
],
"comparison": "用中文简要说明这三个回复的区别和适用场景，帮助用户选择。例如：回复1偏确认型，回复2带感谢，回复3更详细。"
]
}
每条回复控制在 50 字以内。
只返回 JSON，不要包含任何其他文字。"""
@@ -113,7 +109,7 @@ MORE_REPLIES_PROMPT_SHORT = """你是一个邮件回复助手。根据以下邮
MORE_REPLIES_PROMPT_FULL = """你是一个邮件回复助手。根据以下邮件内容生成3个不同的正式邮件回复。
规则：
- 完整模式：生成完整的邮件格式，包括称呼（如 Dear xxx,）和正文，结尾只写问候语逗号，不写落款，用户会自己编辑添加。
- 完整模式：生成完整的邮件格式，包括称呼（如 Dear xxx,）和正文，结尾只写问候语逗号（如 Best regards,）不写落款，用户会自己编辑添加。
- 必须使用与邮件相同的语言回复（英文邮件用英文，中文邮件用中文）。
- 语气正式专业，适合商务邮件。
@@ -125,8 +121,7 @@ MORE_REPLIES_PROMPT_FULL = """你是一个邮件回复助手。根据以下邮
{"text": "回复内容"},
{"text": "回复内容"},
{"text": "回复内容"}
],
"comparison": "用中文简要说明这三个回复的区别和适用场景，帮助用户选择。例如：回复1偏正式确认，回复2表达感谢并跟进，回复3提供详细方案。"
]
}
每条回复控制在 300 字以内。
只返回 JSON，不要包含任何其他文字。"""
@@ -154,10 +149,7 @@ def generate_more_replies(ai_cfg: AIConfig, sender: str, subject: str, body: str
"response_format": {"type": "json_object"},
}
result = _call_deepseek(ai_cfg, payload)
return {
"suggestions": result.get("suggestions", []),
"comparison": result.get("comparison", ""),
}
return result.get("suggestions", [])
def _call_deepseek(ai_cfg: AIConfig, payload: dict) -> dict:
src/config.py
+34 -49
View File
@@ -1,30 +1,25 @@
from dataclasses import dataclass, field
from dataclasses import dataclass
from typing import List, Optional
import json
@dataclass
class ImapConfig:
server: str
port: int = 993
ssl: bool = True
starttls: bool = False
import yaml
@dataclass
class SmtpConfig:
server: str
port: int = 465
ssl: bool = True
starttls: bool = False
port: int
use_ssl: bool = True
use_starttls: bool = False
@dataclass
class EmailAccount:
imap_server: str
imap_port: int
username: str
password: str
imap: ImapConfig
smtp: SmtpConfig
use_ssl: bool = True
use_starttls: bool = False
smtp: Optional[SmtpConfig] = None
@dataclass
@@ -35,50 +30,40 @@ class AIConfig:
@dataclass
class UserConfig:
class TelegramConfig:
bot_token: str
chat_id: str
ai: AIConfig
email_accounts: List[EmailAccount] = field(default_factory=list)
@dataclass
class PollingConfig:
interval_seconds: int
@dataclass
class Config:
bot_token: str
polling_interval: int
users: List[UserConfig] = field(default_factory=list)
email_accounts: List[EmailAccount]
ai: AIConfig
telegram: TelegramConfig
polling: PollingConfig
def load_config(path: str) -> Config:
with open(path, "r", encoding="utf-8") as f:
raw = json.load(f)
raw = yaml.safe_load(f)
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,
))
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"])
return Config(
bot_token=raw["bot_token"],
polling_interval=raw.get("polling_interval", 30),
users=users,
email_accounts=accounts,
ai=ai,
telegram=tg,
polling=polling,
)
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
src/email_client.py
+7 -37
View File
@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
class Email:
def __init__(self, uid: bytes, subject: str, sender: str, recipient: str,
body: str, date: str, reply_to: str = "", account_email: str = "",
folder: str = "INBOX", cc: str = "", list_unsubscribe: str = ""):
folder: str = "INBOX"):
self.uid = uid
self.subject = subject
self.sender = sender
@@ -24,8 +24,6 @@ class Email:
self.reply_to = reply_to
self.account_email = account_email
self.folder = folder
self.cc = cc
self.list_unsubscribe = list_unsubscribe
def _decode_str(s: str) -> str:
@@ -73,31 +71,6 @@ def _get_text_from_msg(msg) -> str:
return ""
def _parse_list_unsubscribe(header_value: str) -> str:
"""解析 List-Unsubscribe 头（RFC 2369），返回退订链接。
优先返回 http(s) 链接（可作按钮直接打开），没有则退回 mailto: 链接。
示例头: <https://x.com/unsub>, <mailto:unsub@x.com>
"""
if not header_value:
return ""
# 标准格式是尖括号包裹的 URL 列表，逐个提取
urls = re.findall(r"<([^>]+)>", header_value)
if not urls:
# 兼容无尖括号的裸格式，按逗号切分后取 http/mailto
urls = [u.strip() for u in header_value.split(",")]
http_url = ""
mailto_url = ""
for u in urls:
u = u.strip()
if u.startswith(("http://", "https://")):
http_url = u
break
if u.startswith("mailto:"):
mailto_url = u
return http_url or mailto_url
_PROVIDERS_NEED_ID = ("163.com", "126.com")
@@ -118,13 +91,12 @@ def _select_mailbox(conn, mailbox: str = "INBOX"):
def _login_and_prepare(account: EmailAccount, timeout: int = 30, select_inbox: bool = True):
imap = account.imap
logger.info("连接 %s:%d (超时%ds)", imap.server, imap.port, timeout)
if imap.ssl:
conn = imaplib.IMAP4_SSL(imap.server, imap.port, timeout=timeout)
logger.info("连接 %s:%d (超时%ds)", account.imap_server, account.imap_port, timeout)
if account.use_ssl:
conn = imaplib.IMAP4_SSL(account.imap_server, account.imap_port, timeout=timeout)
else:
conn = imaplib.IMAP4(imap.server, imap.port, timeout=timeout)
if imap.starttls:
conn = imaplib.IMAP4(account.imap_server, account.imap_port, timeout=timeout)
if account.use_starttls:
conn.starttls()
conn.login(account.username, account.password)
logger.info("登录成功: %s", account.username)
@@ -192,17 +164,15 @@ def fetch_unseen_emails(account: EmailAccount, skip_uids: set[bytes] | None = No
subject = _decode_str(msg.get("Subject", ""))
sender = _decode_address_header(msg.get("From", ""))
recipient = _decode_address_header(msg.get("To", ""))
cc = _decode_address_header(msg.get("Cc", ""))
reply_to = _decode_address_header(msg.get("Reply-To", ""))
date_str = msg.get("Date", "")
body = _get_text_from_msg(msg)
body_len = len(body)
list_unsubscribe = _parse_list_unsubscribe(msg.get("List-Unsubscribe", ""))
logger.info(" 邮件: [%s] from=%s to=%s reply-to=%s (%d 字符)", subject, sender, recipient, reply_to, body_len)
emails.append(Email(uid=uid, subject=subject, sender=sender, recipient=recipient,
body=body, date=date_str, reply_to=reply_to, account_email=account.username,
folder=folder, cc=cc, list_unsubscribe=list_unsubscribe))
folder=folder))
# FETCH 可能被服务器自动标已读，立即标回未读
conn.uid("STORE", uid, "-FLAGS", "\\Seen")
src/smtp_client.py
+3 -3
View File
@@ -23,12 +23,12 @@ def send_reply(account: EmailAccount, to_addr: str, subject: str, body: str,
raise RuntimeError(f"账号 {account.username} 未配置 SMTP")
s = account.smtp
logger.info("SMTP 连接: %s:%d (ssl=%s)", s.server, s.port, s.ssl)
if s.ssl:
logger.info("SMTP 连接: %s:%d (ssl=%s)", s.server, s.port, s.use_ssl)
if s.use_ssl:
conn = smtplib.SMTP_SSL(s.server, s.port, timeout=15)
else:
conn = smtplib.SMTP(s.server, s.port, timeout=15)
if s.starttls:
if s.use_starttls:
conn.starttls()
conn.login(account.username, account.password)
src/summarizer.py
+24 -33
View File
@@ -2,7 +2,7 @@ import logging
import re
import threading
from typing import Generator, Optional
from src.config import Config, UserConfig
from src.config import Config
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,33 +21,31 @@ def _mail_label(mail: Email) -> str:
return f"{_escape_md(sender_short)} · {_escape_md(subject_short)}"
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):
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):
try:
logger.info("检查邮箱: %s (user=%s)", acct.username, user.chat_id)
logger.info(f"检查邮箱: {acct.username}")
emails = fetch_unseen_emails(acct, skip_uids=skip_uids)
if emails:
logger.info(" 发现 %d 封新邮件", len(emails))
logger.info(f" 发现 {len(emails)} 封新邮件")
for mail in emails:
yield idx, mail
except Exception as e:
logger.error("轮询 %s 失败: %s", acct.username, e, exc_info=True)
logger.error(f"轮询 {acct.username} 失败: {e}", exc_info=True)
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)
def ai_process(cfg: Config, acct_idx: int, mail: Email, tg_msg_id: int = 0) -> Optional[dict]:
logger.info(f" 正在摘要: {mail.subject}")
label = _mail_label(mail)
# 阶段1：显示等待处理
if tg_msg_id:
edit_message(cfg.bot_token, user.chat_id, tg_msg_id, f"⏳ *等待处理…*\n{label}")
edit_message(cfg.telegram, cfg.telegram.chat_id, tg_msg_id, f"⏳ *等待处理…*\n{label}")
else:
tg_msg_id = send_thinking(cfg.bot_token, user.chat_id, f"⏳ *等待处理…*\n{label}")
tg_msg_id = send_thinking(cfg.telegram, cfg.telegram.chat_id, f"⏳ *等待处理…*\n{label}")
# 阶段2：AI 处理，更新为思考中
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, mail.list_unsubscribe)
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)
summary["recipient"] = mail.recipient
if "@" not in summary.get("sender", ""):
summary["sender"] = mail.sender
@@ -60,31 +58,27 @@ def ai_process(cfg: Config, user: UserConfig, acct_idx: int, mail: Email, tg_msg
"original_sender": mail.sender,
"original_recipient": mail.recipient,
"original_reply_to": mail.reply_to,
"original_cc": mail.cc,
"original_subject": mail.subject,
"acct_idx": acct_idx,
"account_email": mail.account_email,
"uid": mail.uid,
"folder": mail.folder,
"list_unsubscribe": mail.list_unsubscribe,
"thinking_msg_id": tg_msg_id if tg_msg_id else 0,
}
def tg_send_and_mark(cfg: Config, user: UserConfig, info: dict):
# 找到用户对应的邮箱账户
acct = user.email_accounts[info["acct_idx"]]
def tg_send_and_mark(cfg: Config, info: dict):
acct = cfg.email_accounts[info["acct_idx"]]
thinking_msg_id = info.get("thinking_msg_id")
if thinking_msg_id:
from src.tg_bot import _summary_keyboard, _ensure_unsub_link
# 编辑思考中消息为实际摘要
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"
# AI 漏提取退订链接时，用 List-Unsubscribe 头兜底补上
links = _ensure_unsub_link(info["data"].get("links", []), info.get("list_unsubscribe", ""))
info["data"]["links"] = links
edit_message(cfg.bot_token, user.chat_id, thinking_msg_id,
links = info["data"].get("links", [])
edit_message(cfg.telegram, cfg.telegram.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"],
@@ -92,30 +86,27 @@ def tg_send_and_mark(cfg: Config, user: UserConfig, info: dict):
"plain_text": info["data"].get("plain_text", ""),
"original_sender": info.get("original_sender", ""),
"original_recipient": info.get("original_recipient", ""),
"original_cc": info.get("original_cc", ""),
"original_reply_to": info.get("original_reply_to", ""),
"original_subject": info.get("original_subject", ""),
"account_email": info.get("account_email", ""),
"account_idx": info["acct_idx"],
"sender": info["data"].get("sender", ""),
"subject": info["data"].get("subject", ""),
"cc": info.get("original_cc", ""),
"can_reply": can_reply,
})
else:
send_summary(cfg.bot_token, user.chat_id,
# 兜底：没有思考中消息就正常发送
send_summary(cfg.telegram, cfg.telegram.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", ""),
list_unsubscribe=info.get("list_unsubscribe", ""))
# 标记已读放到后台线程
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(" [%s] TG 发送成功，标记已读已提交后台", user.chat_id)
logger.info(f" TG 发送成功，标记已读已提交后台")
src/tg_bot.py
+127 -249
View File
@@ -4,7 +4,7 @@ import threading
from concurrent.futures import ThreadPoolExecutor
from email.utils import parseaddr
import requests
from src.config import Config, UserConfig, find_user_by_chat_id
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
@@ -18,20 +18,14 @@ _ai_pool = ThreadPoolExecutor(max_workers=3, thread_name_prefix="ai_regen")
# ── Public API ──────────────────────────────────────────
@retry()
def send_summary(bot_token: str, chat_id: str, summary_text: str,
def send_summary(tg_cfg: TelegramConfig, 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 = "",
original_cc: str = "", list_unsubscribe: str = "") -> int:
original_reply_to: str = "", account_email: str = "") -> int:
can_reply = summary_data.get("can_reply", True)
is_promotion = summary_data.get("category") == "promotion"
# AI 漏提取退订链接时，用 List-Unsubscribe 头兜底补上
links = _ensure_unsub_link(summary_data.get("links", []), list_unsubscribe)
summary_data["links"] = links
# summary_data 里可能有 cc
if original_cc:
summary_data["cc"] = original_cc
result = _tg_req(bot_token, "sendMessage", {
links = summary_data.get("links", [])
result = _tg_req(tg_cfg, "sendMessage", {
"chat_id": chat_id,
"text": summary_text,
"parse_mode": "MarkdownV2",
@@ -44,37 +38,35 @@ def send_summary(bot_token: str, chat_id: str, summary_text: str,
"original_body": original_body,
"original_sender": original_sender or summary_data.get("sender", ""),
"original_recipient": original_recipient or summary_data.get("recipient", ""),
"original_cc": original_cc,
"original_reply_to": original_reply_to,
"account_idx": account_idx,
"account_email": account_email,
"sender": summary_data.get("sender", ""),
"subject": summary_data.get("subject", ""),
"cc": original_cc,
"can_reply": can_reply,
})
return msg_id
@retry()
def send_text(bot_token: str, chat_id: str, text: str,
def send_text(tg_cfg: TelegramConfig, 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(bot_token, "sendMessage", payload)
result = _tg_req(tg_cfg, "sendMessage", payload)
return result["result"]["message_id"]
def send_thinking(bot_token: str, chat_id: str, text: str = "💭 思考中…") -> int:
result = _tg_req(bot_token, "sendMessage", {
def send_thinking(tg_cfg: TelegramConfig, chat_id: str, text: str = "💭 思考中…") -> int:
result = _tg_req(tg_cfg, "sendMessage", {
"chat_id": chat_id,
"text": text,
})
return result["result"]["message_id"]
def edit_message(bot_token: str, chat_id: str, msg_id: int,
def edit_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int,
text: str, reply_markup: dict = None):
payload = {
"chat_id": chat_id, "message_id": msg_id,
@@ -82,22 +74,22 @@ def edit_message(bot_token: str, chat_id: str, msg_id: int,
}
if reply_markup:
payload["reply_markup"] = reply_markup
_tg_req(bot_token, "editMessageText", payload)
_tg_req(tg_cfg, "editMessageText", payload)
def delete_message(bot_token: str, chat_id: str, msg_id: int):
def delete_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int):
try:
_tg_req(bot_token, "deleteMessage", {
_tg_req(tg_cfg, "deleteMessage", {
"chat_id": chat_id, "message_id": msg_id,
})
except Exception:
pass
def poll_telegram(bot_token: str, cfg: Config, last_update_id: int) -> int:
def poll_telegram(tg_cfg: TelegramConfig, cfg: Config, last_update_id: int) -> int:
try:
resp = requests.get(
f"https://api.telegram.org/bot{bot_token}/getUpdates",
f"https://api.telegram.org/bot{tg_cfg.bot_token}/getUpdates",
params={"offset": last_update_id, "timeout": 1},
timeout=10,
)
@@ -108,7 +100,7 @@ def poll_telegram(bot_token: str, cfg: Config, last_update_id: int) -> int:
for update in data.get("result", []):
uid = update["update_id"]
try:
_handle_update(bot_token, cfg, update)
_handle_update(tg_cfg, cfg, update)
except Exception as e:
logger.warning("处理 update %d 失败: %s", uid, e)
last_update_id = uid + 1
@@ -154,21 +146,15 @@ def _format_verification(data: dict, account_email: str = "") -> str:
def _format_normal(data: dict, account_email: str = "") -> str:
recipient = data.get("recipient", "未知")
cc = data.get("cc", "")
lines = [
f"📧 *{_escape(data.get('subject', '邮件摘要'))}*",
f"━━━━━━━━━━━━━━━━━━",
f"*收件账户:* {_escape(account_email)}",
f"*收件人:* {_escape(recipient)}",
]
if cc:
lines.append(f"*抄送:* {_escape(cc)}")
lines.extend([
f"*收件人:* {_escape(data.get('recipient', '未知'))}",
f"*发件人:* {_escape(data.get('sender', '未知'))}",
"",
_escape(data.get("summary", "")),
])
]
return "\n".join(lines)
@@ -193,75 +179,12 @@ CALLBACK_CANCEL = "c"
CALLBACK_PROMO_DETAIL = "pd"
# ── Link buttons ────────────────────────────────────────
_UNSUB_KEYWORDS = (
"unsubscribe", "opt out", "unsub", "退订", "取消订阅",
"manage preferences", "email preferences", "subscription settings",
"subscription preferences", "preferences", "email-pref", "manage-pref",
"偏好设置", "管理偏好", "订阅设置",
)
def _is_unsub_link(link: dict) -> bool:
"""判断链接是否为退订/取消订阅类链接（按文字或 URL 关键词）"""
text = (link.get("text") or "").lower()
url = (link.get("url") or "").lower()
return any(k in text or k in url for k in _UNSUB_KEYWORDS)
def _link_buttons(links: list | None) -> list:
"""把 AI 提取的链接转成键盘按钮行；退订链接排最前并用 🔕 标识，其余用 🔗"""
if not links:
return []
unsub_links = [l for l in links if _is_unsub_link(l)]
other_links = [l for l in links if not _is_unsub_link(l)]
rows = []
for link in unsub_links + other_links:
url = link.get("url", "")
# 校验 URL：http/https 需无多重协议头；mailto: 仅用于退订场景（如 mailto:x@y.com?subject=Unsubscribe）
if url.startswith("http"):
valid = url.count("://") == 1 and len(url) < 2048
elif url.startswith("mailto:"):
valid = len(url) < 2048
else:
valid = False
if not valid:
continue
text = (link.get("text") or "").strip()[:30]
if not text:
text = "退订" if _is_unsub_link(link) else "打开链接"
icon = "🔕" if _is_unsub_link(link) else "🔗"
rows.append([{"text": f"{icon} {text}", "url": url}])
return rows
def _ensure_unsub_link(links: list, list_unsubscribe: str = "") -> list:
"""AI 未提取到退订链接时，用邮件头 List-Unsubscribe 的链接兜底补一个退订按钮。
仅当 links 中没有任何退订类链接且头部确实有可用的 http(s)/mailto 链接时生效。
"""
if any(_is_unsub_link(l) for l in links):
return links
url = (list_unsubscribe or "").strip()
if url.startswith("http"):
valid = url.count("://") == 1 and len(url) < 2048
elif url.startswith("mailto:"):
valid = len(url) < 2048
else:
valid = False
if valid:
return [{"text": "退订", "url": url}] + list(links)
return links
def _summary_keyboard(can_reply: bool = True, is_promotion: bool = False, links: list | None = None) -> dict:
def _summary_keyboard(can_reply: bool = True, is_promotion: bool = False, links: list = None) -> dict:
if is_promotion:
kb = []
kb.extend(_link_buttons(links))
kb.append([{"text": "告诉我更多", "callback_data": CALLBACK_PROMO_DETAIL}])
kb.append([{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG}])
return {"inline_keyboard": kb}
return {"inline_keyboard": [
[{"text": "告诉我更多", "callback_data": CALLBACK_PROMO_DETAIL}],
[{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG}],
]}
kb = [
[
{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG},
@@ -273,7 +196,12 @@ def _summary_keyboard(can_reply: bool = True, is_promotion: bool = False, links:
{"text": "AI完整回复", "callback_data": CALLBACK_AI_REPLY_FULL},
{"text": "AI简洁回复", "callback_data": CALLBACK_AI_REPLY_SHORT},
])
kb.extend(_link_buttons(links))
for link in (links or []):
url = link.get("url", "")
text = link.get("text", "打开链接")[:30]
# 校验 URL：合法 http/https，无多重协议头，长度合理
if url.startswith("http") and url.count("://") == 1 and len(url) < 200:
kb.append([{"text": f"🔗 {text}", "url": url}])
return {"inline_keyboard": kb}
@@ -319,8 +247,8 @@ def _ai_reply_keyboard(msg_id: int, suggestions: list) -> dict:
# ── Telegram API ────────────────────────────────────────
def _tg_req(bot_token: str, method: str, payload: dict) -> dict:
url = f"https://api.telegram.org/bot{bot_token}/{method}"
def _tg_req(tg_cfg: TelegramConfig, method: str, payload: dict) -> dict:
url = f"https://api.telegram.org/bot{tg_cfg.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:
@@ -332,90 +260,64 @@ def _tg_req(bot_token: str, method: str, payload: dict) -> dict:
return result
def _cleanup_conv_messages(bot_token: str, chat_id: int, conv: dict):
"""清理回复流程中产生的所有中间消息"""
chat_id_str = str(chat_id)
for key in ("prompt_msg_id",):
mid = conv.get(key)
if mid:
delete_message(bot_token, chat_id_str, mid)
for mid in conv.get("extra_msg_ids", []):
delete_message(bot_token, chat_id_str, mid)
# ── Update handling ─────────────────────────────────────
def _do_regen(bot_token: str, cfg: Config, chat_id: int, msg_id: int, ctx: dict, conv: dict):
def _do_regen(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:
result = generate_more_replies(
user.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
new_suggestions = generate_more_replies(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []), style=style,
)
new_suggestions = result["suggestions"]
comparison = result.get("comparison", "")
except Exception as e:
send_text(bot_token, str(chat_id), f"生成失败: {e}")
send_text(cfg.telegram, str(chat_id), f"生成失败: {e}")
return
if new_suggestions:
conv["ai_suggestions"] = new_suggestions
conv["all_suggestions"].extend(new_suggestions)
conv["comparison"] = comparison
save_conversation(chat_id, conv)
logger.info("[ai_regen] 新生成 %d 条回复 (style=%s)", len(new_suggestions), style)
_show_ai_suggestions(bot_token, chat_id, msg_id, new_suggestions, comparison)
_show_ai_suggestions(cfg.telegram, chat_id, msg_id, new_suggestions)
def _do_ai_hint(bot_token: str, cfg: Config, chat_id: int, conv: dict, ctx: dict, hint_history: list[str]):
"""在 AI 线程池中根据用户提示生成回复（支持多轮追问）"""
user = find_user_by_chat_id(cfg, chat_id)
if not user:
return
def _do_ai_hint(cfg: Config, chat_id: int, conv: dict, ctx: dict, hint: str):
"""在 AI 线程池中根据用户提示生成回复"""
style = conv.get("reply_style", "short")
# 拼接历史提示作为上下文
combined_hint = " → ".join(hint_history)
try:
result = generate_more_replies(
user.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []), user_hint=combined_hint, style=style,
new_suggestions = generate_more_replies(
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
conv.get("all_suggestions", []), user_hint=hint, style=style,
)
new_suggestions = result["suggestions"]
comparison = result.get("comparison", "")
except Exception as e:
send_text(bot_token, str(chat_id), f"生成失败: {e}")
send_text(cfg.telegram, str(chat_id), f"生成失败: {e}")
delete_conversation(chat_id)
return
if new_suggestions:
conv["ai_suggestions"] = new_suggestions
conv["all_suggestions"].extend(new_suggestions)
conv["comparison"] = comparison
conv["state"] = "ai_reply_selection"
conv["hint_history"] = hint_history
conv.pop("prompt_msg_id", None)
save_conversation(chat_id, conv)
logger.info("[ai_hint] 根据提示生成 %d 条新回复 (历史 %d 条)", len(new_suggestions), len(hint_history))
_show_ai_suggestions(bot_token, chat_id, conv["summary_msg_id"], new_suggestions, comparison)
logger.info("[ai_hint] 根据提示生成 %d 条新回复", len(new_suggestions))
_show_ai_suggestions(cfg.telegram, chat_id, conv["summary_msg_id"], new_suggestions)
def _handle_update(bot_token: str, cfg: Config, update: dict):
def _handle_update(tg_cfg: TelegramConfig, cfg: Config, update: dict):
if "callback_query" in update:
_handle_callback(bot_token, cfg, update["callback_query"])
_handle_callback(tg_cfg, cfg, update["callback_query"])
elif "message" in update:
_handle_message(bot_token, cfg, update["message"])
_handle_message(tg_cfg, cfg, update["message"])
def _handle_callback(bot_token: str, cfg: Config, cb: dict):
def _handle_callback(tg_cfg: TelegramConfig, 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(bot_token, "answerCallbackQuery", {"callback_query_id": cb_id})
_tg_req(tg_cfg, "answerCallbackQuery", {"callback_query_id": cb_id})
except Exception:
pass
@@ -423,19 +325,13 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
action = parts[0]
ctx = load_email_context(msg_id)
subject = ctx.get("original_subject") or ctx["subject"] if ctx else "?"
subject = ctx["subject"] if ctx else "?"
logger.info("TG 回调: action=%s msg_id=%d subject=%s", action, msg_id, subject)
if action == CALLBACK_VIEW_ORIG and ctx:
can_reply = ctx.get("can_reply", True)
subject = ctx.get("subject", "邮件")
body = ctx.get("plain_text", ctx["original_body"])
if len(body) > 3500:
body = body[:3500] + "\n\n... (内容过长已截断)"
# 用代码块包裹原文，避免 MarkdownV2 解析特殊字符
escaped_body = body.replace("\\", "\\\\").replace("`", "\\`")
text = f"*📄 原文 \\- {_escape(subject)}*\n\n`{escaped_body}`"
_tg_req(bot_token, "editMessageText", {
text = f"*📄 原文 \\- {_escape(ctx['subject'])}*\n\n{_escape(ctx.get('plain_text', ctx['original_body']))}"
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": text, "parse_mode": "MarkdownV2",
"reply_markup": _orig_keyboard(can_reply),
@@ -446,7 +342,7 @@ def _handle_callback(bot_token: str, 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(bot_token, "editMessageText", {
_tg_req(tg_cfg, "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),
@@ -455,7 +351,7 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
elif action == CALLBACK_REPLY and ctx:
prompt_msg_id = send_text(
bot_token, str(chat_id), "请输入你的回复内容：",
tg_cfg, str(chat_id), "请输入你的回复内容：",
reply_markup=_cancel_keyboard(),
)
save_conversation(chat_id, {
@@ -466,19 +362,33 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
elif action in (CALLBACK_AI_REPLY_FULL, CALLBACK_AI_REPLY_SHORT) and ctx:
style = "full" if action == CALLBACK_AI_REPLY_FULL else "short"
# 总是用对应风格重新生成，不使用摘要时的通用缓存
logger.info(" 重新生成 AI 回复 (style=%s)", style)
_tg_req(bot_token, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
suggestions = ctx["summary_data"].get("reply_suggestions", [])
# 缓存为空时重新生成
if not suggestions:
logger.info(" 缓存为空，重新生成 AI 回复 (style=%s)", style)
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
conv = load_conversation(chat_id) or {}
conv["state"] = "ai_reply_selection"
conv["summary_msg_id"] = msg_id
conv["reply_style"] = style
conv["all_suggestions"] = []
save_conversation(chat_id, conv)
_ai_pool.submit(_do_regen, cfg, chat_id, msg_id, ctx, conv)
return
save_conversation(chat_id, {
"state": "ai_reply_selection",
"summary_msg_id": msg_id,
"ai_suggestions": suggestions,
"all_suggestions": list(suggestions),
"reply_style": style,
})
conv = load_conversation(chat_id) or {}
conv["state"] = "ai_reply_selection"
conv["summary_msg_id"] = msg_id
conv["reply_style"] = style
conv["all_suggestions"] = []
save_conversation(chat_id, conv)
_ai_pool.submit(_do_regen, bot_token, cfg, chat_id, msg_id, ctx, conv)
logger.info(" 显示 AI 回复建议 (%d 条, style=%s)", len(suggestions), style)
_show_ai_suggestions(tg_cfg, chat_id, msg_id, suggestions)
elif action == CALLBACK_SEND_REPLY and ctx:
idx = int(parts[2])
@@ -490,19 +400,14 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
return
sel_text = suggestions[idx]["text"]
logger.info(" 直接发送回复 #%d: %s", idx, sel_text[:60])
# 清理中间消息
_cleanup_conv_messages(bot_token, chat_id, conv)
_do_send_reply(bot_token, cfg, chat_id, msg_id, ctx, sel_text)
send_text(bot_token, str(chat_id), f"✅ 回复已发送")
_do_send_reply(tg_cfg, cfg, chat_id, msg_id, ctx, sel_text)
send_text(tg_cfg, 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(bot_token, "editMessageText", {
_tg_req(tg_cfg, "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),
"reply_markup": _summary_keyboard(can_reply),
})
elif action == CALLBACK_SELECT_REPLY and ctx:
@@ -519,16 +424,15 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
# 用代码块显示，点击可复制（MarkdownV2）
escaped = sel_text.replace("\\", "\\\\").replace("`", "\\`")
copy_result = _tg_req(bot_token, "sendMessage", {
_tg_req(tg_cfg, "sendMessage", {
"chat_id": str(chat_id),
"text": f"*建议回复（点击代码块复制）:*\n`{escaped}`",
"parse_mode": "MarkdownV2",
})
copy_msg_id = copy_result["result"]["message_id"]
# 进入 awaiting_reply 状态，用户可编辑后发送
prompt_msg_id = send_text(
bot_token, str(chat_id),
tg_cfg, str(chat_id),
"请编辑后发送（或直接回复附件）：",
reply_markup=_cancel_keyboard(),
)
@@ -536,7 +440,6 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
conv["summary_msg_id"] = msg_id
conv["prompt_msg_id"] = prompt_msg_id
conv["draft_text"] = sel_text
conv["extra_msg_ids"] = [copy_msg_id]
save_conversation(chat_id, conv)
logger.info(" 建议已展示，等待用户编辑发送")
@@ -546,61 +449,53 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
return
logger.info(" 换一批 (已有 %d 条历史)", len(conv.get("all_suggestions", [])))
# AI 处理中
_tg_req(bot_token, "editMessageText", {
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
# 丢到 AI 线程池处理，不阻塞按钮轮询
_ai_pool.submit(_do_regen, bot_token, cfg, chat_id, msg_id, ctx, conv)
_ai_pool.submit(_do_regen, cfg, chat_id, msg_id, ctx, conv)
elif action == CALLBACK_HINT and ctx:
prompt_msg_id = send_text(
bot_token, str(chat_id), "请输入你对回复的提示/要求：",
tg_cfg, str(chat_id), "请输入你对回复的提示/要求：",
reply_markup=_cancel_keyboard(),
)
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
if "hint_history" not in conv:
conv["hint_history"] = []
save_conversation(chat_id, conv)
logger.info(" 等待用户输入 AI 提示 (历史 %d 条)", len(conv["hint_history"]))
logger.info(" 等待用户输入 AI 提示")
elif action == CALLBACK_CANCEL:
conv = load_conversation(chat_id)
if conv:
_cleanup_conv_messages(bot_token, chat_id, conv)
delete_conversation(chat_id)
# 清除回复建议缓存，重新选择时会重新生成
if ctx:
ctx["summary_data"]["reply_suggestions"] = []
save_email_context(msg_id, ctx)
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(bot_token, "editMessageText", {
_tg_req(tg_cfg, "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),
"reply_markup": _summary_keyboard(can_reply),
})
logger.info(" 取消, 恢复摘要视图")
else:
delete_message(bot_token, str(chat_id), msg_id)
delete_message(tg_cfg, str(chat_id), msg_id)
logger.info(" 取消, 删除提示消息")
elif action == CALLBACK_PROMO_DETAIL and ctx:
_tg_req(bot_token, "editMessageText", {
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
try:
user = find_user_by_chat_id(cfg, chat_id)
details = expand_promo_detail(
user.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
)
except Exception as e:
send_text(bot_token, str(chat_id), f"获取详情失败: {e}")
send_text(tg_cfg, 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 "暂无更多详情"
@@ -611,19 +506,18 @@ def _handle_callback(bot_token: str, cfg: Config, cb: dict):
f"\n{_escape(summary)}\n"
f"\n*📋 详情:*\n{detail_lines}"
)
promo_kb = []
promo_kb.extend(_link_buttons(ctx.get("summary_data", {}).get("links", [])))
promo_kb.append([{"text": "收起", "callback_data": CALLBACK_VIEW_SUMM}])
promo_kb.append([{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG}])
_tg_req(bot_token, "editMessageText", {
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": msg_id,
"text": expanded, "parse_mode": "MarkdownV2",
"reply_markup": {"inline_keyboard": promo_kb},
"reply_markup": {"inline_keyboard": [
[{"text": "收起", "callback_data": CALLBACK_VIEW_SUMM}],
[{"text": "查看原文", "callback_data": CALLBACK_VIEW_ORIG}],
]},
})
logger.info(" 显示推广详情")
def _handle_message(bot_token: str, cfg: Config, msg: dict):
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
chat_id = msg["chat"]["id"]
text = msg.get("text") or msg.get("caption") or ""
text = text.strip()
@@ -644,29 +538,24 @@ def _handle_message(bot_token: str, cfg: Config, msg: dict):
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
delete_conversation(chat_id)
delete_message(bot_token, chat_id_str, user_msg_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(bot_token, chat_id_str, prompt_msg_id)
for mid in conv.get("extra_msg_ids", []):
delete_message(bot_token, chat_id_str, mid)
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
logger.info(" 上下文已丢失，清理")
return
delete_message(bot_token, chat_id_str, user_msg_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(bot_token, chat_id_str, prompt_msg_id)
# 清理 extra 消息（代码块等）
for mid in conv.get("extra_msg_ids", []):
delete_message(bot_token, chat_id_str, mid)
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
# 合并累积的附件
all_attachments = conv.get("pending_attachments", [])
all_attachments.extend(media)
_do_send_reply(bot_token, cfg, chat_id, conv["summary_msg_id"], ctx, text,
_do_send_reply(tg_cfg, 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(bot_token, chat_id_str, f"✅ 回复已发送{attach_info}")
send_text(tg_cfg, chat_id_str, f"✅ 回复已发送{attach_info}")
delete_conversation(chat_id)
logger.info(" 回复发送完成")
@@ -674,57 +563,50 @@ def _handle_message(bot_token: str, cfg: Config, msg: dict):
ctx = load_email_context(conv["summary_msg_id"])
if not ctx:
delete_conversation(chat_id)
delete_message(bot_token, chat_id_str, user_msg_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(bot_token, chat_id_str, prompt_msg_id)
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
return
delete_message(bot_token, chat_id_str, user_msg_id)
delete_message(tg_cfg, chat_id_str, user_msg_id)
if prompt_msg_id:
delete_message(bot_token, chat_id_str, prompt_msg_id)
# 追加到 hint 历史
hint_history = conv.get("hint_history", [])
hint_history.append(text)
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
# 更新消息为 AI 思考中
_tg_req(bot_token, "editMessageText", {
_tg_req(tg_cfg, "editMessageText", {
"chat_id": chat_id, "message_id": conv["summary_msg_id"],
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
})
# 丢到 AI 线程池处理，传入完整 hint 历史
_ai_pool.submit(_do_ai_hint, bot_token, cfg, chat_id, conv, ctx, hint_history)
# 丢到 AI 线程池处理
_ai_pool.submit(_do_ai_hint, cfg, chat_id, conv, ctx, text)
def _show_ai_suggestions(bot_token: str, chat_id: int, msg_id: int,
suggestions: list, comparison: str = ""):
text = ""
if comparison:
text = f"*💡 三个回复的区别:*\n{_escape(comparison)}\n\n"
text += "*AI 建议回复，请选择：*"
def _show_ai_suggestions(tg_cfg: TelegramConfig, 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(bot_token, "editMessageText", {
_tg_req(tg_cfg, "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(bot_token: str, file_id: str, filename: str) -> str:
def _download_tg_file(tg_cfg: TelegramConfig, 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{bot_token}/getFile",
f"https://api.telegram.org/bot{tg_cfg.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{bot_token}/{file_path}",
f"https://api.telegram.org/file/bot{tg_cfg.bot_token}/{file_path}",
timeout=60,
)
dl.raise_for_status()
@@ -778,16 +660,12 @@ def _extract_media(msg: dict) -> list[dict]:
return attachments
def _do_send_reply(bot_token: str, cfg: Config,
def _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
chat_id: int, msg_id: int, ctx: dict, reply_text: str,
attachments: list[dict] | None = None):
user = find_user_by_chat_id(cfg, chat_id)
if not user:
send_text(bot_token, str(chat_id), "❌ 用户未找到。")
return
acct = user.email_accounts[ctx["account_idx"]]
acct = cfg.email_accounts[ctx["account_idx"]]
if not acct.smtp:
send_text(bot_token, str(chat_id), "❌ 该邮箱未配置 SMTP，无法发送回复。")
send_text(tg_cfg, str(chat_id), "❌ 该邮箱未配置 SMTP，无法发送回复。")
return
to_addr = parseaddr(ctx.get("original_reply_to", ""))[1]
if not to_addr:
@@ -796,14 +674,14 @@ def _do_send_reply(bot_token: str, cfg: Config,
to_addr = parseaddr(ctx["sender"])[1]
if not to_addr:
to_addr = ctx["sender"]
subject = ctx.get("original_subject") or ctx["subject"]
subject = ctx["subject"]
# 下载 Telegram 附件到本地
local_files = []
if attachments:
for att in attachments:
try:
path = _download_tg_file(bot_token, att["file_id"], att["filename"])
path = _download_tg_file(tg_cfg, att["file_id"], att["filename"])
local_files.append((path, att["mime_type"]))
except Exception as e:
logger.error(" 附件下载失败: %s - %s", att["filename"], e)
tests/test_unsub_keyboard.py
-149
View File
@@ -1,149 +0,0 @@
"""测试摘要键盘渲染：推广邮件的退订链接按钮必须显示且排最前。"""
from src.tg_bot import (
_ensure_unsub_link,
_is_unsub_link,
_link_buttons,
_summary_keyboard,
)
def _rows_of(kb: dict) -> list[list[dict]]:
return kb["inline_keyboard"]
def test_promotion_with_unsub_link_shows_button_first():
"""推广邮件：AI 提取到退订链接时必须显示，且排在第一行"""
links = [
{"text": "Manage Preferences", "url": "https://e.example.com/pref"},
{"text": "推广活动详情", "url": "https://e.example.com/deal"},
]
kb = _summary_keyboard(can_reply=False, is_promotion=True, links=links)
rows = _rows_of(kb)
first = rows[0][0]
assert first["url"] == "https://e.example.com/pref"
assert "🔕" in first["text"]
# 其余按钮仍然存在
texts = [b["text"] for r in rows for b in r]
assert "告诉我更多" in texts
assert "查看原文" in texts
def test_promotion_without_links_keeps_default_buttons():
"""推广邮件：没有链接时保持原有按钮"""
kb = _summary_keyboard(can_reply=False, is_promotion=True, links=None)
rows = _rows_of(kb)
assert len(rows) == 2
assert rows[0][0]["text"] == "告诉我更多"
assert rows[1][0]["text"] == "查看原文"
def test_promotion_with_only_non_unsub_link():
"""推广邮件：只有非退订链接时也正常渲染（🔗），不冒充退订按钮"""
links = [{"text": "查看活动", "url": "https://e.example.com/event"}]
kb = _summary_keyboard(can_reply=False, is_promotion=True, links=links)
rows = _rows_of(kb)
assert rows[0][0]["text"] == "🔗 查看活动"
assert rows[1][0]["text"] == "告诉我更多"
assert rows[2][0]["text"] == "查看原文"
def test_normal_email_links_still_rendered():
"""普通邮件：链接按钮行为不变"""
links = [{"text": "查看账单", "url": "https://e.example.com/bill"}]
kb = _summary_keyboard(can_reply=True, is_promotion=False, links=links)
rows = _rows_of(kb)
texts = [b["text"] for r in rows for b in r]
assert "🔗 查看账单" in texts
def test_unsub_detection_variants():
"""退订链接识别：文字与 URL 关键词变体"""
assert _is_unsub_link({"text": "Unsubscribe", "url": "https://x.com"})
assert _is_unsub_link({"text": "退订", "url": "https://x.com"})
assert _is_unsub_link({"text": "点击这里取消订阅", "url": "https://x.com"})
assert _is_unsub_link({"text": "管理偏好设置", "url": "https://x.com"})
assert _is_unsub_link({"text": "更多", "url": "https://x.com/unsubscribe"})
assert _is_unsub_link({"text": "Preferences", "url": "https://x.com/email-pref"})
assert not _is_unsub_link({"text": "查看详情", "url": "https://x.com/deal"})
def test_unsub_link_sorted_before_other_links():
"""退订链接即使不在 links 第一位也会被排到最前"""
links = [
{"text": "查看详情", "url": "https://x.com/a"},
{"text": "unsubscribe", "url": "https://x.com/b"},
]
rows = _link_buttons(links)
assert rows[0][0]["url"] == "https://x.com/b"
assert rows[1][0]["url"] == "https://x.com/a"
def test_invalid_url_filtered():
"""非法 URL（非 http、多重协议头）被过滤"""
links = [
{"text": "退订", "url": "javascript:alert(1)"},
{"text": "退订", "url": "https://x.com/https://evil.com"},
{"text": "退订", "url": "https://ok.example.com/unsub"},
]
rows = _link_buttons(links)
assert len(rows) == 1
assert rows[0][0]["url"] == "https://ok.example.com/unsub"
def test_long_unsub_url_rendered():
"""长 tracking 退订链接（如 OpenAI 的 >700 字符）必须渲染，不能被长度过滤误杀"""
url = "https://r.openai.com/asm/unsubscribe/?user_id=108370056&data=" + "x" * 700
assert len(url) > 200
links = [{"text": "取消订阅", "url": url}]
rows = _link_buttons(links)
assert len(rows) == 1
assert rows[0][0]["url"] == url
assert "🔕" in rows[0][0]["text"]
def test_oversized_url_filtered():
"""超过 2048 的 URL 仍被过滤"""
url = "https://example.com/" + "x" * 2100
links = [{"text": "退订", "url": url}]
rows = _link_buttons(links)
assert rows == []
def test_mailto_unsub_link_rendered():
"""mailto: 退订链接（如 InferX 的 mailto:team@inferx.net?subject=Unsubscribe）必须渲染"""
links = [{"text": "退订", "url": "mailto:team@inferx.net?subject=Unsubscribe"}]
kb = _summary_keyboard(can_reply=False, is_promotion=True, links=links)
rows = kb["inline_keyboard"]
assert rows[0][0]["url"] == "mailto:team@inferx.net?subject=Unsubscribe"
assert "🔕" in rows[0][0]["text"]
assert _is_unsub_link(links[0])
def test_ensure_unsub_link_http_header_fallback():
"""AI 没提取到退订链接时，用 List-Unsubscribe 头的 http 链接兜底"""
header = "https://links.iterable.com/s/uh/tKjpy4SHCKvbzE2votNyZEJDvlMdeodVyWFExNs310ugJ_RvZtoZq4gGLQ0BXMSL2lf3kEZsmr1XJmrODIFHpEwcypo3LR096pZ7jsfBynZBgGWz171Wc1E6ZIIC5Lu0e-Stq6p5-_Ad3CnNlUWrW1tnrLSeiGsOiQG_zMYKVbtfQZgvfm3QYfD9lQ/Cey29KU8yOzIUltaEJA8EMRFf-FKMZQy/24"
links = _ensure_unsub_link([], header)
assert links == [{"text": "退订", "url": header}]
kb = _summary_keyboard(can_reply=False, is_promotion=True, links=links)
assert kb["inline_keyboard"][0][0]["url"] == header
def test_ensure_unsub_link_mailto_header_fallback():
"""List-Unsubscribe 头只有 mailto 时也兜底"""
header = "mailto:unsubscribe+19693436+25822730@unsubscribe.iterable.com"
links = _ensure_unsub_link([], header)
assert links == [{"text": "退订", "url": header}]
def test_ensure_unsub_link_no_duplicate_when_ai_extracted():
"""AI 已提取退订链接时，不重复添加头部链接"""
ai_links = [{"text": "取消订阅", "url": "https://r.openai.com/asm/unsubscribe/?x=1"}]
links = _ensure_unsub_link(ai_links, "https://links.example.com/unsub")
assert links == ai_links
def test_ensure_unsub_link_ignores_invalid_header():
"""头部无有效链接时不兜底"""
assert _ensure_unsub_link([], "") == []
assert _ensure_unsub_link([], "javascript:alert(1)") == []
Reference in new issue
earthjasonlin/ai-mail-bot
Block a user
Blocking a user prevents them from interacting with repositories, such as opening or commenting on pull requests or issues. Learn more about blocking a user.

The note is not visible to the blocked user.