TelegramBotAdmin/wbb/modules/business_assistant.py

from __future__ import annotations

import asyncio
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from typing import Any

from pyrogram import filters
from pyrogram.enums import ChatMemberStatus, ChatType

import wbb
from wbb import SUDOERS, app
from wbb.services.business_assistant import BusinessAssistantRuntime, classify_handoff
from wbb.utils.dbassistant import (
    AssistantDataError,
    get_account_settings,
    get_business_connection,
    get_conversation,
    get_knowledge_source,
    list_enabled_sources_for_chat,
    list_sources_for_chat,
    mark_source_event_deleted,
    resume_conversation,
    touch_knowledge_source,
    update_knowledge_source_sync,
)

__MODULE__ = "智能接待"
__HELP__ = "Telegram Business 智能接待由 Web 管理后台配置。"

_runtime: BusinessAssistantRuntime | None = None
_source_sync_tasks: dict[str, asyncio.Task[None]] = {}


async def start_business_assistant_runtime() -> None:
    global _runtime
    if _runtime is not None:
        return
    session = getattr(wbb, "aiohttpsession", None)
    token = str(getattr(wbb, "BOT_TOKEN", ""))
    if not session or not token:
        wbb.log.error("智能接待未启动：Telegram Bot 会话或令牌不可用。")
        return
    _runtime = BusinessAssistantRuntime(token=token, session=session)
    status = await _runtime.start()
    if status.get("polling_state") == "running":
        wbb.log.info("Telegram Business 智能接待已启动。")
    else:
        wbb.log.error(f"Telegram Business 智能接待未运行：{status.get('last_error') or '未知原因'}")


async def stop_business_assistant_runtime() -> None:
    global _runtime
    tasks = [task for task in _source_sync_tasks.values() if not task.done()]
    for task in tasks:
        task.cancel()
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)
    _source_sync_tasks.clear()
    if _runtime is None:
        return
    await _runtime.stop()
    _runtime = None


def get_business_assistant_runtime() -> BusinessAssistantRuntime | None:
    return _runtime


async def validate_knowledge_source(values: dict[str, Any]) -> dict[str, Any]:
    source_type = str(values.get("source_type") or "")
    if source_type == "business":
        return {
            **values,
            "chat_id": 0,
            "title": str(values.get("title") or "Business 日常对话"),
            "linked_chat_id": 0,
            "linked_chat_title": "",
            "access_status": "business_connected",
        }
    try:
        chat_id = int(values.get("chat_id") or 0)
    except (TypeError, ValueError) as exc:
        raise AssistantDataError("invalid_source", "来源聊天 ID 必须是整数。") from exc
    if not chat_id:
        raise AssistantDataError("invalid_source", "请填写频道或群组聊天 ID。")
    try:
        chat = await app.get_chat(chat_id)
        member = await app.get_chat_member(chat_id, int(wbb.BOT_ID))
    except Exception as exc:
        raise AssistantDataError(
            "source_access_denied", "机器人无法访问该频道或群组，请先将机器人加入。"
        ) from exc
    chat_type = getattr(chat, "type", None)
    if source_type == "channel" and chat_type != ChatType.CHANNEL:
        raise AssistantDataError("source_type_mismatch", "该聊天不是频道。")
    if source_type == "group" and chat_type not in {ChatType.GROUP, ChatType.SUPERGROUP}:
        raise AssistantDataError("source_type_mismatch", "该聊天不是群组。")
    if source_type not in {"channel", "group"}:
        raise AssistantDataError("invalid_source", "知识来源类型无效。")
    if source_type == "channel" and member.status not in {
        ChatMemberStatus.OWNER,
        ChatMemberStatus.ADMINISTRATOR,
    }:
        raise AssistantDataError(
            "source_access_denied", "机器人必须是频道管理员才能持续采集频道消息。"
        )
    linked_chat = getattr(chat, "linked_chat", None)
    return {
        **values,
        "chat_id": chat_id,
        "title": str(getattr(chat, "title", None) or values.get("title") or chat_id),
        "linked_chat_id": int(getattr(linked_chat, "id", 0) or 0),
        "linked_chat_title": str(getattr(linked_chat, "title", "") or ""),
        "access_status": str(getattr(member.status, "value", member.status)),
    }


def _message_text(message: Any) -> str:
    return str(getattr(message, "text", None) or getattr(message, "caption", None) or "").strip()


async def _trusted_source_author(
    source: dict[str, Any], message: Any, *, effective_source_type: str
) -> bool:
    if effective_source_type == "channel":
        return True
    author = getattr(message, "from_user", None)
    author_id = int(getattr(author, "id", 0) or 0)
    policy = str(source.get("author_policy") or "admins_or_allowlist")
    allowlisted = author_id in {int(item) for item in source.get("allowed_user_ids") or []}
    if policy == "all":
        return True
    if policy == "allowlist":
        return allowlisted
    sender_chat_id = int(getattr(getattr(message, "sender_chat", None), "id", 0) or 0)
    is_admin = sender_chat_id == int(message.chat.id)
    if author_id:
        with suppress(Exception):
            member = await app.get_chat_member(int(message.chat.id), author_id)
            is_admin = member.status in {
                ChatMemberStatus.OWNER,
                ChatMemberStatus.ADMINISTRATOR,
            }
    if policy == "admins":
        return is_admin
    return allowlisted or is_admin


async def _ingest_source_message(source: dict[str, Any], message: Any) -> bool:
    runtime = get_business_assistant_runtime()
    if runtime is None:
        return False
    text = _message_text(message)
    author = getattr(message, "from_user", None)
    if (
        not text
        or bool(getattr(message, "outgoing", False))
        or bool(getattr(author, "is_bot", False))
        or classify_handoff(text) == "sensitive_request"
    ):
        return False
    chat_id = int(message.chat.id)
    linked_discussion = bool(
        source.get("include_linked_chat")
        and int(source.get("linked_chat_id") or 0) == chat_id
    )
    if linked_discussion and bool(getattr(message, "is_automatic_forward", False)):
        return False
    effective_source_type = "discussion" if linked_discussion else str(source["source_type"])
    trusted = await _trusted_source_author(
        source, message, effective_source_type=effective_source_type
    )
    author_id = int(getattr(author, "id", 0) or 0)
    message_id = int(getattr(message, "id", 0) or 0)
    if not message_id:
        return False
    try:
        await runtime.ingestion.ingest(
            source,
            event_key=f"telegram:{chat_id}:{message_id}",
            content=text,
            metadata={
                "effective_source_type": effective_source_type,
                "chat_id": chat_id,
                "message_id": message_id,
                "author_id": author_id,
                "trusted_author": trusted,
                "message_date": getattr(message, "date", None),
            },
            auto_publish=bool(
                source.get("publication_mode") == "auto" and trusted
            ),
        )
        await touch_knowledge_source(str(source["source_id"]))
        return True
    except Exception as exc:
        await touch_knowledge_source(str(source["source_id"]), error=str(exc))
        wbb.log.error(f"知识来源采集失败：{source['source_id']} / {exc}")
        return False


async def process_knowledge_source_message(message: Any) -> None:
    chat_id = int(getattr(getattr(message, "chat", None), "id", 0) or 0)
    if not chat_id:
        return
    for source in await list_enabled_sources_for_chat(chat_id):
        await _ingest_source_message(source, message)


async def _run_source_sync(source_id: str) -> None:
    processed = 0
    try:
        source = await get_knowledge_source(source_id)
        if not source:
            return
        await update_knowledge_source_sync(source_id, status="running", processed=0)
        cutoff = datetime.now(UTC) - timedelta(days=int(source.get("backfill_days") or 0))
        chat_ids = [int(source.get("chat_id") or 0)]
        if source.get("include_linked_chat") and source.get("linked_chat_id"):
            chat_ids.append(int(source["linked_chat_id"]))
        limit = int(source.get("backfill_limit") or 100)
        for chat_id in dict.fromkeys(item for item in chat_ids if item):
            remaining = limit - processed
            if remaining <= 0:
                break
            async for message in app.get_chat_history(chat_id, limit=remaining):
                message_date = getattr(message, "date", None)
                if isinstance(message_date, datetime):
                    aware_date = (
                        message_date.replace(tzinfo=UTC)
                        if message_date.tzinfo is None
                        else message_date.astimezone(UTC)
                    )
                    if int(source.get("backfill_days") or 0) and aware_date < cutoff:
                        break
                if await _ingest_source_message(source, message):
                    processed += 1
        await update_knowledge_source_sync(
            source_id, status="completed", processed=processed
        )
    except asyncio.CancelledError:
        raise
    except Exception as exc:
        with suppress(Exception):
            await update_knowledge_source_sync(
                source_id, status="failed", error=str(exc), processed=processed
            )
        wbb.log.error(f"知识来源历史同步失败：{source_id} / {exc}")


async def schedule_knowledge_source_sync(source_id: str) -> dict[str, Any]:
    if get_business_assistant_runtime() is None:
        raise AssistantDataError(
            "assistant_runtime_unavailable", "智能接待运行时尚未启动。"
        )
    source = await get_knowledge_source(source_id)
    if not source:
        raise AssistantDataError("source_not_found", "未找到知识来源。")
    current = _source_sync_tasks.get(source_id)
    if current and not current.done():
        return source
    queued = await update_knowledge_source_sync(source_id, status="queued", processed=0)
    task = asyncio.create_task(
        _run_source_sync(source_id), name=f"business-knowledge-sync-{source_id}"
    )
    _source_sync_tasks[source_id] = task
    task.add_done_callback(lambda _task: _source_sync_tasks.pop(source_id, None))
    return queued


async def cancel_knowledge_source_sync(source_id: str) -> None:
    task = _source_sync_tasks.pop(source_id, None)
    if task and not task.done():
        task.cancel()
        await asyncio.gather(task, return_exceptions=True)


async def regenerate_knowledge_candidate(candidate_id: str) -> list[dict[str, Any]]:
    runtime = get_business_assistant_runtime()
    if runtime is None:
        raise AssistantDataError(
            "assistant_runtime_unavailable", "智能接待运行时尚未启动。"
        )
    return await runtime.ingestion.regenerate_candidate(candidate_id)


@app.on_message(filters.channel | filters.group, group=-30)
async def collect_knowledge_source_message(_, message):
    await process_knowledge_source_message(message)


@app.on_edited_message(filters.channel | filters.group, group=-30)
async def collect_edited_knowledge_source_message(_, message):
    await process_knowledge_source_message(message)


@app.on_deleted_messages(filters.channel | filters.group, group=-30)
async def collect_deleted_knowledge_source_messages(_, messages):
    for message in messages:
        chat_id = int(getattr(getattr(message, "chat", None), "id", 0) or 0)
        message_id = int(getattr(message, "id", 0) or 0)
        if not chat_id or not message_id:
            continue
        for source in await list_sources_for_chat(chat_id, enabled_only=False):
            await mark_source_event_deleted(
                str(source["source_id"]), f"telegram:{chat_id}:{message_id}"
            )


async def _ops_group_admin_allowed(query, settings: dict) -> bool:
    ops_group_id = int(settings.get("ops_group_id") or 0)
    if not ops_group_id or not query.message or int(query.message.chat.id) != ops_group_id:
        return False
    with suppress(Exception):
        member = await app.get_chat_member(ops_group_id, query.from_user.id)
        return member.status in {
            ChatMemberStatus.OWNER,
            ChatMemberStatus.ADMINISTRATOR,
        }
    return False


@app.on_callback_query(filters.regex(r"^ba:resume:[a-f0-9]{32}$"), group=-20)
async def resume_business_assistant_callback(_, query):
    conversation_id = str(query.data).rsplit(":", 1)[-1]
    conversation = await get_conversation(conversation_id)
    if not conversation:
        return await query.answer("会话不存在或已清理。", show_alert=True)
    connection = await get_business_connection(conversation["connection_id"])
    if not connection:
        return await query.answer("Business 连接不存在。", show_alert=True)
    settings = await get_account_settings(connection["connection_id"])
    owner_id = int((connection.get("user") or {}).get("id") or 0)
    allowed = (
        query.from_user.id == owner_id
        or query.from_user.id in SUDOERS
        or await _ops_group_admin_allowed(query, settings)
    )
    if not allowed:
        return await query.answer("只有账号本人或运营群管理员可以恢复。", show_alert=True)
    await resume_conversation(conversation_id)
    with suppress(Exception):
        await query.message.edit_reply_markup(None)
    await query.answer("已恢复自动回复。", show_alert=True)