tg-spam-guard/handlers/knowledge.py

"""知识库问答 handler"""
import asyncio
import logging
import re

from telegram import Update
from telegram.ext import CommandHandler, ContextTypes, filters

import config

logger = logging.getLogger("spam_guard")


def looks_like_markdown(text: str) -> bool:
    return bool(re.search(r"(^|\n)(#{1,3}\s|[-*]\s|\d+\.\s|```)|\*\*[^*]+\*\*|\*[^*]+\*|`[^`]+`", text))


def optimize_plain_answer(text: str) -> str:
    """非 Markdown 回答也做一点可读性整理。"""
    text = text.strip()
    if looks_like_markdown(text):
        return text
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if len(lines) <= 1:
        return f"*回答*\n{text}"
    return "*回答*\n" + "\n".join(f"- {line}" if not line.startswith(("-", "•", "✅", "⚠️", "📚")) else line for line in lines)


def telegram_markdown(text: str) -> str:
    """把常见 Markdown 调整为 Telegram legacy Markdown 更容易接受的形式。"""
    text = re.sub(r"\*\*([^*]+)\*\*", r"*\1*", text)
    text = re.sub(r"^#{1,3}\s+(.+)$", r"*\1*", text, flags=re.MULTILINE)
    return text


def compact_answer(text: str, max_len: int = 3800) -> str:
    """尽量压到一条 Telegram 消息；超长时保留开头和关键尾部。"""
    text = optimize_plain_answer(text)
    if len(text) <= max_len:
        return text
    tail_markers = ["**文档**", "📚", "相关文档", "文档"]
    tail = ""
    for marker in tail_markers:
        idx = text.rfind(marker)
        if idx > max_len // 2:
            tail = text[idx:].strip()
            break
    head_len = max_len - len(tail) - 80 if tail else max_len - 40
    head = text[:max(1000, head_len)].rstrip()
    return head + "\n\n...（内容较长已压缩，可追问具体步骤）" + (("\n\n" + tail) if tail else "")


def split_answer(text: str, max_len: int = 3900, max_parts: int = 2) -> list[str]:
    """优先单条；实在超长最多拆 2 条。"""
    if len(text) <= max_len:
        return [text]
    parts = []
    rest = text
    while rest and len(parts) < max_parts:
        cut = min(len(rest), max_len)
        if cut < len(rest):
            better = rest.rfind("\n\n", 0, cut)
            if better > max_len // 2:
                cut = better
        parts.append(rest[:cut].strip())
        rest = rest[cut:].strip()
    if rest and parts:
        parts[-1] += "\n\n...（后续内容已省略，可继续追问）"
    return parts


async def safe_send(update: Update, context: ContextTypes.DEFAULT_TYPE, text: str, markdown: bool | None = None):
    """优先回复原消息；Markdown 失败或原消息被删时自动降级。"""
    chat_id = update.effective_chat.id
    use_markdown = looks_like_markdown(text) if markdown is None else markdown
    send_text = telegram_markdown(text) if use_markdown else text

    async def _send_plain():
        try:
            return await update.message.reply_text(text)
        except Exception:
            return await context.bot.send_message(chat_id=chat_id, text=text)

    if use_markdown:
        try:
            return await update.message.reply_text(send_text, parse_mode="Markdown", disable_web_page_preview=True)
        except Exception as e:
            if "Message to be replied not found" in str(e):
                try:
                    return await context.bot.send_message(chat_id=chat_id, text=send_text, parse_mode="Markdown", disable_web_page_preview=True)
                except Exception:
                    return await context.bot.send_message(chat_id=chat_id, text=text)
            logger.warning(f"Markdown 发送失败，降级普通文本: {e}")
            return await _send_plain()

    try:
        return await update.message.reply_text(text, disable_web_page_preview=True)
    except Exception as e:
        if "Message to be replied not found" not in str(e):
            logger.warning(f"reply_text 失败，改用 send_message: {e}")
        return await context.bot.send_message(chat_id=chat_id, text=text, disable_web_page_preview=True)


async def cmd_ask(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """向知识库提问"""
    query = " ".join(context.args) if context.args else ""
    if not query:
        reply = await safe_send(update, context,
            "❓ 用法: /ask <你的问题>\n"
            "例如: /ask 怎么安装收割机\n"
            f"文档地址: {config.SITE_URL}"
        )
        asyncio.create_task(auto_delete(update.message, seconds=180))
        asyncio.create_task(auto_delete(reply, seconds=180))
        return

    logger.info(f"📩 用户提问: [{update.message.from_user.username}] {query}")
    searching_msg = await safe_send(update, context, "🔍 正在搜索知识库并整理详细答案...")
    asyncio.create_task(auto_delete(update.message, seconds=180))

    try:
        from services.knowledge_base import ask_knowledge_base

        answer = compact_answer(await ask_knowledge_base(query))
        replies = []
        parts = split_answer(answer)
        for i, part in enumerate(parts, 1):
            prefix = f"*回答 {i}/{len(parts)}*\n\n" if len(parts) > 1 else ""
            reply = await safe_send(update, context, prefix + part, markdown=looks_like_markdown(prefix + part))
            replies.append(reply)
        for reply in replies:
            asyncio.create_task(auto_delete(reply, seconds=180))
        asyncio.create_task(auto_delete(searching_msg, seconds=30))
    except Exception as e:
        logger.error(f"❌ 回答失败: {e}")
        reply = await safe_send(update, context, f"❌ 回答失败: {e}")
        asyncio.create_task(auto_delete(reply, seconds=180))
        asyncio.create_task(auto_delete(searching_msg, seconds=30))


async def auto_delete(message, seconds=180):
    try:
        await asyncio.sleep(seconds)
        await message.delete()
    except Exception:
        pass


def register_knowledge(app):
    app.add_handler(CommandHandler("ask", cmd_ask, filters.ChatType.GROUPS))