Compare commits
4 Commits
5ca739456d
...
d76a119190
| Author | SHA1 | Date | |
|---|---|---|---|
| d76a119190 | |||
| d05277581c | |||
| fe24a07d15 | |||
| 0abb7de5d3 |
6
MEMORY.md
Normal file
6
MEMORY.md
Normal file
@@ -0,0 +1,6 @@
|
||||
# AI Mail Bot - 记忆
|
||||
|
||||
## 固定规则
|
||||
|
||||
- **Telegram parse_mode 永远为 MarkdownV2**,不传 None。需要发纯文本时,用 Markdown 转义包装,而不是去掉 parse_mode。
|
||||
- **所有代码改动必须完整测试**,除非用户明确说不需要测试。没有测试的代码不算完成。
|
||||
@@ -2,6 +2,9 @@ import smtplib
|
||||
import logging
|
||||
from email.mime.text import MIMEText
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.base import MIMEBase
|
||||
from email import encoders
|
||||
from pathlib import Path
|
||||
from src.config import EmailAccount
|
||||
from src.retry import retry
|
||||
|
||||
@@ -9,7 +12,13 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@retry()
|
||||
def send_reply(account: EmailAccount, to_addr: str, subject: str, body: str):
|
||||
def send_reply(account: EmailAccount, to_addr: str, subject: str, body: str,
|
||||
attachments: list[tuple[str, str]] | None = None):
|
||||
"""发送回复邮件,支持附件。
|
||||
|
||||
Args:
|
||||
attachments: [(本地文件路径, mime_type), ...] 列表
|
||||
"""
|
||||
if not account.smtp:
|
||||
raise RuntimeError(f"账号 {account.username} 未配置 SMTP")
|
||||
|
||||
@@ -26,13 +35,32 @@ def send_reply(account: EmailAccount, to_addr: str, subject: str, body: str):
|
||||
logger.info("SMTP 登录成功")
|
||||
|
||||
reply_subject = subject if subject.lower().startswith("re:") else f"Re: {subject}"
|
||||
logger.info("SMTP 发送: to=%s subject=%s body=%d 字符", to_addr, reply_subject, len(body))
|
||||
logger.info("SMTP 发送: to=%s subject=%s body=%d 字符 attachments=%d",
|
||||
to_addr, reply_subject, len(body), len(attachments or []))
|
||||
|
||||
if attachments:
|
||||
msg = MIMEMultipart()
|
||||
msg.attach(MIMEText(body, "plain", "utf-8"))
|
||||
for filepath, mime_type in attachments:
|
||||
p = Path(filepath)
|
||||
if not p.exists():
|
||||
logger.warning(" 附件不存在: %s", filepath)
|
||||
continue
|
||||
maintype, subtype = mime_type.split("/", 1)
|
||||
with open(p, "rb") as f:
|
||||
part = MIMEBase(maintype, subtype)
|
||||
part.set_payload(f.read())
|
||||
encoders.encode_base64(part)
|
||||
part.add_header("Content-Disposition", "attachment", filename=p.name)
|
||||
msg.attach(part)
|
||||
logger.info(" 添加附件: %s (%s, %d bytes)", p.name, mime_type, p.stat().st_size)
|
||||
else:
|
||||
msg = MIMEMultipart("alternative")
|
||||
msg.attach(MIMEText(body, "plain", "utf-8"))
|
||||
|
||||
msg["From"] = account.username
|
||||
msg["To"] = to_addr
|
||||
msg["Subject"] = reply_subject
|
||||
msg.attach(MIMEText(body, "plain", "utf-8"))
|
||||
|
||||
conn.sendmail(account.username, [to_addr], msg.as_string())
|
||||
conn.quit()
|
||||
|
||||
201
src/tg_bot.py
201
src/tg_bot.py
@@ -1,5 +1,7 @@
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from email.utils import parseaddr
|
||||
import requests
|
||||
from src.config import TelegramConfig, Config
|
||||
@@ -10,6 +12,8 @@ from src.database import save_email_context, load_email_context, save_conversati
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ai_pool = ThreadPoolExecutor(max_workers=3, thread_name_prefix="ai_regen")
|
||||
|
||||
|
||||
# ── Public API ──────────────────────────────────────────
|
||||
|
||||
@@ -246,10 +250,49 @@ 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):
|
||||
"""在 AI 线程池中执行换一批,不阻塞按钮轮询"""
|
||||
try:
|
||||
new_suggestions = generate_more_replies(
|
||||
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
|
||||
conv.get("all_suggestions", []),
|
||||
)
|
||||
except Exception as 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)
|
||||
save_conversation(chat_id, conv)
|
||||
logger.info("[ai_regen] 新生成 %d 条回复", len(new_suggestions))
|
||||
_show_ai_suggestions(cfg.telegram, chat_id, msg_id, new_suggestions)
|
||||
|
||||
|
||||
def _do_ai_hint(cfg: Config, chat_id: int, conv: dict, ctx: dict, hint: str):
|
||||
"""在 AI 线程池中根据用户提示生成回复"""
|
||||
try:
|
||||
new_suggestions = generate_more_replies(
|
||||
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
|
||||
conv.get("all_suggestions", []), user_hint=hint,
|
||||
)
|
||||
except Exception as 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["state"] = "ai_reply_selection"
|
||||
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)
|
||||
|
||||
|
||||
def _handle_update(tg_cfg: TelegramConfig, cfg: Config, update: dict):
|
||||
if "callback_query" in update:
|
||||
_handle_callback(tg_cfg, cfg, update["callback_query"])
|
||||
elif "message" in update and "text" in update["message"]:
|
||||
elif "message" in update:
|
||||
_handle_message(tg_cfg, cfg, update["message"])
|
||||
|
||||
|
||||
@@ -345,21 +388,10 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
|
||||
# AI 处理中
|
||||
_tg_req(tg_cfg, "editMessageText", {
|
||||
"chat_id": chat_id, "message_id": msg_id,
|
||||
"text": "🤖 AI思考中...", "parse_mode": None,
|
||||
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
|
||||
})
|
||||
try:
|
||||
new_suggestions = generate_more_replies(
|
||||
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
|
||||
conv.get("all_suggestions", []),
|
||||
)
|
||||
except Exception as e:
|
||||
send_text(tg_cfg, str(chat_id), f"生成失败: {e}")
|
||||
return
|
||||
if new_suggestions:
|
||||
conv["ai_suggestions"] = new_suggestions
|
||||
conv["all_suggestions"].extend(new_suggestions)
|
||||
logger.info(" 新生成 %d 条回复", len(new_suggestions))
|
||||
_show_ai_suggestions(tg_cfg, chat_id, msg_id, new_suggestions)
|
||||
# 丢到 AI 线程池处理,不阻塞按钮轮询
|
||||
_ai_pool.submit(_do_regen, cfg, chat_id, msg_id, ctx, conv)
|
||||
|
||||
elif action == CALLBACK_HINT and ctx:
|
||||
prompt_msg_id = send_text(
|
||||
@@ -370,6 +402,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
|
||||
if conv and conv.get("summary_msg_id") == msg_id:
|
||||
conv["state"] = "awaiting_ai_hint"
|
||||
conv["prompt_msg_id"] = prompt_msg_id
|
||||
save_conversation(chat_id, conv)
|
||||
logger.info(" 等待用户输入 AI 提示")
|
||||
|
||||
elif action == CALLBACK_CANCEL:
|
||||
@@ -389,7 +422,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
|
||||
elif action == CALLBACK_PROMO_DETAIL and ctx:
|
||||
_tg_req(tg_cfg, "editMessageText", {
|
||||
"chat_id": chat_id, "message_id": msg_id,
|
||||
"text": "🤖 AI思考中...", "parse_mode": None,
|
||||
"text": "🤖 *AI思考中…*", "parse_mode": "MarkdownV2",
|
||||
})
|
||||
try:
|
||||
details = expand_promo_detail(
|
||||
@@ -420,7 +453,7 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
|
||||
|
||||
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
|
||||
chat_id = msg["chat"]["id"]
|
||||
text = msg["text"].strip()
|
||||
text = msg.get("text", "").strip()
|
||||
conv = load_conversation(chat_id)
|
||||
|
||||
if not conv:
|
||||
@@ -430,8 +463,9 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
|
||||
chat_id_str = str(chat_id)
|
||||
prompt_msg_id = conv.get("prompt_msg_id")
|
||||
state = conv.get("state")
|
||||
media = _extract_media(msg)
|
||||
|
||||
logger.info("TG 消息: state=%s text=%s", state, text[:60])
|
||||
logger.info("TG 消息: state=%s text=%s media=%d", state, text[:60] if text else "(无)", len(media))
|
||||
|
||||
if state == "awaiting_reply":
|
||||
ctx = load_email_context(conv["summary_msg_id"])
|
||||
@@ -443,11 +477,32 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
|
||||
logger.info(" 上下文已丢失,清理")
|
||||
return
|
||||
|
||||
# 有附件但没文字:累积附件,等用户输入文字
|
||||
if media and not text:
|
||||
pending = conv.get("pending_attachments", [])
|
||||
for att in media:
|
||||
pending.append(att)
|
||||
conv["pending_attachments"] = pending
|
||||
save_conversation(chat_id, conv)
|
||||
delete_message(tg_cfg, chat_id_str, user_msg_id)
|
||||
send_text(tg_cfg, chat_id_str,
|
||||
f"📎 已添加 {len(pending)} 个附件,请输入回复内容发送",
|
||||
reply_markup=_cancel_keyboard())
|
||||
return
|
||||
|
||||
# 有文字(可能也有附件):发送邮件
|
||||
delete_message(tg_cfg, chat_id_str, user_msg_id)
|
||||
if prompt_msg_id:
|
||||
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
|
||||
_do_send_reply(tg_cfg, cfg, chat_id, conv["summary_msg_id"], ctx, text)
|
||||
send_text(tg_cfg, chat_id_str, f"✅ 回复已发送: {text}")
|
||||
|
||||
# 合并累积的附件
|
||||
all_attachments = conv.get("pending_attachments", [])
|
||||
all_attachments.extend(media)
|
||||
|
||||
_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(tg_cfg, chat_id_str, f"✅ 回复已发送: {text}{attach_info}")
|
||||
delete_conversation(chat_id)
|
||||
logger.info(" 回复发送完成")
|
||||
|
||||
@@ -463,23 +518,14 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
|
||||
delete_message(tg_cfg, chat_id_str, user_msg_id)
|
||||
if prompt_msg_id:
|
||||
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
|
||||
try:
|
||||
new_suggestions = generate_more_replies(
|
||||
cfg.ai, ctx["sender"], ctx["subject"], ctx["original_body"],
|
||||
conv.get("all_suggestions", []), user_hint=text,
|
||||
)
|
||||
except Exception as e:
|
||||
send_text(tg_cfg, chat_id_str, f"生成失败: {e}")
|
||||
delete_conversation(chat_id)
|
||||
return
|
||||
|
||||
if new_suggestions:
|
||||
conv["ai_suggestions"] = new_suggestions
|
||||
conv["all_suggestions"].extend(new_suggestions)
|
||||
conv["state"] = "ai_reply_selection"
|
||||
conv.pop("prompt_msg_id", None)
|
||||
logger.info(" 根据提示生成 %d 条新回复", len(new_suggestions))
|
||||
_show_ai_suggestions(tg_cfg, chat_id, conv["summary_msg_id"], new_suggestions)
|
||||
# 更新消息为 AI 思考中
|
||||
_tg_req(tg_cfg, "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)
|
||||
|
||||
|
||||
def _show_ai_suggestions(tg_cfg: TelegramConfig, chat_id: int, msg_id: int,
|
||||
@@ -494,8 +540,76 @@ def _show_ai_suggestions(tg_cfg: TelegramConfig, chat_id: int, msg_id: int,
|
||||
})
|
||||
|
||||
|
||||
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{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{tg_cfg.bot_token}/{file_path}",
|
||||
timeout=60,
|
||||
)
|
||||
dl.raise_for_status()
|
||||
dest = media_dir / filename
|
||||
dest.write_bytes(dl.content)
|
||||
logger.info(" 下载附件: %s (%d bytes)", filename, len(dl.content))
|
||||
return str(dest)
|
||||
|
||||
|
||||
def _extract_media(msg: dict) -> list[dict]:
|
||||
"""从 Telegram 消息中提取附件信息列表 [{file_id, filename, mime_type}]"""
|
||||
attachments = []
|
||||
# 照片取最大尺寸
|
||||
if "photo" in msg:
|
||||
photo = max(msg["photo"], key=lambda p: p.get("file_size", 0))
|
||||
ext = ".jpg"
|
||||
attachments.append({
|
||||
"file_id": photo["file_id"],
|
||||
"filename": f"photo_{photo['file_id'][:8]}{ext}",
|
||||
"mime_type": "image/jpeg",
|
||||
})
|
||||
if "video" in msg:
|
||||
v = msg["video"]
|
||||
attachments.append({
|
||||
"file_id": v["file_id"],
|
||||
"filename": v.get("file_name", f"video_{v['file_id'][:8]}.mp4"),
|
||||
"mime_type": v.get("mime_type", "video/mp4"),
|
||||
})
|
||||
if "document" in msg:
|
||||
d = msg["document"]
|
||||
attachments.append({
|
||||
"file_id": d["file_id"],
|
||||
"filename": d.get("file_name", f"doc_{d['file_id'][:8]}"),
|
||||
"mime_type": d.get("mime_type", "application/octet-stream"),
|
||||
})
|
||||
if "voice" in msg:
|
||||
vc = msg["voice"]
|
||||
attachments.append({
|
||||
"file_id": vc["file_id"],
|
||||
"filename": f"voice_{vc['file_id'][:8]}.ogg",
|
||||
"mime_type": vc.get("mime_type", "audio/ogg"),
|
||||
})
|
||||
if "sticker" in msg:
|
||||
s = msg["sticker"]
|
||||
if not s.get("is_animated"):
|
||||
attachments.append({
|
||||
"file_id": s["file_id"],
|
||||
"filename": f"sticker_{s['file_id'][:8]}.webp",
|
||||
"mime_type": "image/webp",
|
||||
})
|
||||
return attachments
|
||||
|
||||
|
||||
def _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
|
||||
chat_id: int, msg_id: int, ctx: dict, reply_text: str):
|
||||
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,无法发送回复。")
|
||||
@@ -508,4 +622,15 @@ def _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,
|
||||
if not to_addr:
|
||||
to_addr = ctx["sender"]
|
||||
subject = ctx["subject"]
|
||||
send_reply(acct, to_addr, subject, reply_text)
|
||||
|
||||
# 下载 Telegram 附件到本地
|
||||
local_files = []
|
||||
if attachments:
|
||||
for att in attachments:
|
||||
try:
|
||||
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)
|
||||
|
||||
send_reply(acct, to_addr, subject, reply_text, attachments=local_files or None)
|
||||
|
||||
Reference in New Issue
Block a user