diff --git a/src/ai_client.py b/src/ai_client.py index afcc1f6..bcc2e4c 100644 --- a/src/ai_client.py +++ b/src/ai_client.py @@ -49,6 +49,43 @@ def summarize_email(ai_cfg: AIConfig, recipient: str, subject: str, sender: str, return _call_deepseek(ai_cfg, payload) +MORE_REPLIES_PROMPT = """你是一个邮件回复助手。根据以下邮件内容生成3个完全不同的建议回复。 + +以下是一些已经生成过的回复，请生成全新的回复，不要与已有回复重复。 + +返回 JSON: +{ + "suggestions": [ + {"text": "回复内容"}, + {"text": "回复内容"}, + {"text": "回复内容"} + ] +} +每条回复控制在 50 字以内。 +只返回 JSON，不要包含任何其他文字。""" + + +@retry() +def generate_more_replies(ai_cfg: AIConfig, sender: str, subject: str, body: str, + previous_suggestions: list[dict], user_hint: str = "") -> list[dict]: + content = f"发件人: {sender}\n主题: {subject}\n正文:\n{body[:4000]}\n\n已生成过的回复:\n" + for i, s in enumerate(previous_suggestions, 1): + content += f"{i}. {s['text']}\n" + if user_hint: + content += f"\n用户要求: {user_hint}" + + payload = { + "model": ai_cfg.model, + "messages": [ + {"role": "system", "content": MORE_REPLIES_PROMPT}, + {"role": "user", "content": content}, + ], + "response_format": {"type": "json_object"}, + } + result = _call_deepseek(ai_cfg, payload) + return result.get("suggestions", []) + + def _call_deepseek(ai_cfg: AIConfig, payload: dict) -> dict: headers = { "Authorization": f"Bearer {ai_cfg.api_key}", diff --git a/src/tg_bot.py b/src/tg_bot.py index be49fe0..116e11a 100644 --- a/src/tg_bot.py +++ b/src/tg_bot.py @@ -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,