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

feat: cancel button, regenerate AI replies, custom reply hints

Browse Source 
- Reply prompt now has an inline '取消回复' button instead of typing text
- AI reply selection adds '换一批' button: passes all previous suggestions
  to AI to generate different replies
- '我想说：' button: lets user type a hint/tone instruction, AI tailors
  suggestions accordingly
- Cancel button also works on prompt/hint messages (deletes them)
- Extract _show_ai_suggestions helper for reuse
This commit is contained in:
Zichao Lin
2026-07-02 20:45:40 +08:00
parent 18db9caa8b
commit 09d11a6c03
2 changed files with 145 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
src/tg_bot.py
View File

@@ -3,6 +3,7 @@ import logging
from email.utils import parseaddr
import requests
from src.config import TelegramConfig, Config
from src.ai_client import generate_more_replies
from src.smtp_client import send_reply
from src.retry import retry
@@ -42,19 +43,19 @@ def send_summary(tg_cfg: TelegramConfig, chat_id: str, summary_text: str,
@retry()
def send_text(tg_cfg: TelegramConfig, chat_id: str, text: str) -> int:
result = _tg_req(tg_cfg, "sendMessage", {
"chat_id": chat_id,
"text": text,
})
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(tg_cfg, "sendMessage", payload)
return result["result"]["message_id"]
def delete_message(tg_cfg: TelegramConfig, chat_id: str, msg_id: int):
try:
_tg_req(tg_cfg, "deleteMessage", {
"chat_id": chat_id,
"message_id": msg_id,
"chat_id": chat_id, "message_id": msg_id,
})
except Exception:
pass
@@ -121,6 +122,8 @@ CALLBACK_VIEW_SUMM = "s"
CALLBACK_REPLY = "r"
CALLBACK_AI_REPLY = "a"
CALLBACK_SELECT_REPLY = "sr"
CALLBACK_REGEN = "rg"
CALLBACK_HINT = "h"
CALLBACK_CANCEL = "c"
@@ -148,11 +151,23 @@ def _orig_keyboard(can_reply: bool = True) -> dict:
return {"inline_keyboard": kb}
def _cancel_keyboard() -> dict:
return {
"inline_keyboard": [
[{"text": "取消回复", "callback_data": CALLBACK_CANCEL}],
]
}
def _ai_reply_keyboard(msg_id: int, suggestions: list) -> dict:
kb = []
for i, s in enumerate(suggestions):
data = f"{CALLBACK_SELECT_REPLY}|{msg_id}|{i}"
kb.append([{"text": s["text"][:30], "callback_data": data}])
kb.append([
{"text": "换一批", "callback_data": f"{CALLBACK_REGEN}|{msg_id}"},
{"text": "我想说：", "callback_data": f"{CALLBACK_HINT}|{msg_id}"},
])
kb.append([{"text": "取消", "callback_data": CALLBACK_CANCEL}])
return {"inline_keyboard": kb}
@@ -212,37 +227,34 @@ 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), "请输入你的回复内容（或发送 \"取消\" 取消）：")
prompt_msg_id = send_text(
tg_cfg, str(chat_id), "请输入你的回复内容：",
reply_markup=_cancel_keyboard(),
)
_conversations[chat_id] = {
"state": "awaiting_reply", "msg_id": msg_id,
"state": "awaiting_reply", "summary_msg_id": msg_id,
"prompt_msg_id": prompt_msg_id,
}
elif action == CALLBACK_AI_REPLY and ctx:
suggestions = ctx["summary_data"].get("reply_suggestions", [])
if not suggestions:
send_text(tg_cfg, str(chat_id), "此邮件无需回复，已隐藏 AI 回复功能。")
send_text(tg_cfg, str(chat_id), "此邮件无需回复。")
return
_conversations[chat_id] = {
"state": "ai_reply_selection",
"msg_id": msg_id,
"summary_msg_id": msg_id,
"ai_suggestions": suggestions,
"all_suggestions": list(suggestions),
}
text = "*AI 建议回复，请选择：*\n\n"
for i, s in enumerate(suggestions, 1):
text += f"{i}\\. {_escape(s['text'])}\n"
_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),
})
_show_ai_suggestions(tg_cfg, chat_id, msg_id, suggestions)
elif action == CALLBACK_SELECT_REPLY and ctx:
idx = int(parts[2])
conv = _conversations.get(chat_id)
if not conv or conv.get("msg_id") != msg_id:
if not conv or conv.get("summary_msg_id") != msg_id:
return
suggestions = conv.get("ai_suggestions", [])
if idx >= len(suggestions):
@@ -258,6 +270,33 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
"reply_markup": _summary_keyboard(can_reply),
})
elif action == CALLBACK_REGEN and ctx:
conv = _conversations.get(chat_id)
if not conv or conv.get("summary_msg_id") != msg_id:
return
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)
_show_ai_suggestions(tg_cfg, chat_id, msg_id, new_suggestions)
elif action == CALLBACK_HINT and ctx:
prompt_msg_id = send_text(
tg_cfg, str(chat_id), "请输入你对回复的提示/要求：",
reply_markup=_cancel_keyboard(),
)
conv = _conversations.get(chat_id)
if conv and conv.get("summary_msg_id") == msg_id:
conv["state"] = "awaiting_ai_hint"
conv["prompt_msg_id"] = prompt_msg_id
elif action == CALLBACK_CANCEL:
_conversations.pop(chat_id, None)
if ctx:
@@ -267,6 +306,8 @@ def _handle_callback(tg_cfg: TelegramConfig, cfg: Config, cb: dict):
"text": ctx["summary_text"], "parse_mode": "MarkdownV2",
"reply_markup": _summary_keyboard(can_reply),
})
else:
delete_message(tg_cfg, str(chat_id), msg_id)
def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
@@ -274,30 +315,64 @@ def _handle_message(tg_cfg: TelegramConfig, cfg: Config, msg: dict):
text = msg["text"].strip()
conv = _conversations.get(chat_id)
if not conv or conv.get("state") != "awaiting_reply":
if not conv:
return
prompt_msg_id = conv.get("prompt_msg_id")
chat_id_str = str(chat_id)
prompt_msg_id = conv.get("prompt_msg_id")
if conv.get("state") == "awaiting_reply":
ctx = _email_contexts.get(conv["summary_msg_id"])
if not ctx:
_conversations.pop(chat_id, None)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
return
if text in ("取消", "/cancel"):
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, "✅ 回复已发送！")
_conversations.pop(chat_id, None)
return
ctx = _email_contexts.get(conv["msg_id"])
if not ctx:
_conversations.pop(chat_id, None)
elif conv.get("state") == "awaiting_ai_hint":
ctx = _email_contexts.get(conv["summary_msg_id"])
if not ctx:
_conversations.pop(chat_id, None)
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
return
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
return
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}")
_conversations.pop(chat_id, None)
return
if prompt_msg_id:
delete_message(tg_cfg, chat_id_str, prompt_msg_id)
_do_send_reply(tg_cfg, cfg, chat_id, conv["msg_id"], ctx, text)
send_text(tg_cfg, chat_id_str, "✅ 回复已发送！")
_conversations.pop(chat_id, None)
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)
_show_ai_suggestions(tg_cfg, chat_id, conv["summary_msg_id"], new_suggestions)
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(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 _do_send_reply(tg_cfg: TelegramConfig, cfg: Config,