AstrBot/astrbot/dashboard/services/conversation_service.py

from __future__ import annotations

import json
import traceback
from dataclasses import asdict, dataclass
from datetime import datetime
from io import BytesIO

from astrbot.core import logger
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
from astrbot.core.db import BaseDatabase
from astrbot.core.umo_alias import build_umo_alias_map, parse_umo, serialize_umo_alias


class ConversationServiceError(Exception):
    pass


@dataclass
class ConversationExport:
    file_obj: BytesIO
    filename: str
    mimetype: str = "application/jsonl"


class ConversationService:
    def __init__(
        self,
        db_helper: BaseDatabase,
        core_lifecycle: AstrBotCoreLifecycle,
    ) -> None:
        self.db_helper = db_helper
        self.conv_mgr = core_lifecycle.conversation_manager

    async def list_conversations(
        self,
        *,
        page: int,
        page_size: int,
        platforms: str,
        message_types: str,
        search_query: str,
        exclude_ids: str,
        exclude_platforms: str,
    ) -> dict:
        platform_list = platforms.split(",") if platforms else []
        message_type_list = message_types.split(",") if message_types else []
        exclude_id_list = exclude_ids.split(",") if exclude_ids else []
        exclude_platform_list = (
            exclude_platforms.split(",") if exclude_platforms else []
        )

        page = max(page, 1)
        if page_size < 1:
            page_size = 20
        page_size = min(page_size, 100)

        try:
            conversations, total_count = await self.conv_mgr.get_filtered_conversations(
                page=page,
                page_size=page_size,
                platforms=platform_list,
                message_types=message_type_list,
                search_query=search_query,
                exclude_ids=exclude_id_list,
                exclude_platforms=exclude_platform_list,
            )
        except Exception as exc:
            logger.error(f"数据库查询出错: {exc!s}\n{traceback.format_exc()}")
            raise ConversationServiceError(f"数据库查询出错: {exc!s}") from exc

        total_pages = (
            (total_count + page_size - 1) // page_size if total_count > 0 else 1
        )
        umos = sorted({conv.user_id for conv in conversations if conv.user_id})
        alias_map = build_umo_alias_map(await self.db_helper.get_umo_aliases(umos))

        return {
            "conversations": [
                self._serialize_conversation(conversation, alias_map)
                for conversation in conversations
            ],
            "pagination": {
                "page": page,
                "page_size": page_size,
                "total": total_count,
                "total_pages": total_pages,
            },
        }

    async def get_conversation_detail(self, data: object) -> dict:
        payload = self._payload(data)
        user_id, cid = self._require_user_and_cid(payload)

        conversation = await self.conv_mgr.get_conversation(
            unified_msg_origin=user_id,
            conversation_id=cid,
        )
        if not conversation:
            raise ConversationServiceError("对话不存在")

        alias_map = build_umo_alias_map(await self.db_helper.get_umo_aliases([user_id]))
        return {
            "user_id": user_id,
            "cid": cid,
            "title": conversation.title,
            "persona_id": conversation.persona_id,
            "history": conversation.history,
            "created_at": conversation.created_at,
            "updated_at": conversation.updated_at,
            "umo_info": self._build_umo_info(user_id, alias_map),
        }

    async def update_conversation(self, data: object) -> dict:
        payload = self._payload(data)
        user_id, cid = self._require_user_and_cid(payload)
        title = payload.get("title")

        conversation = await self.conv_mgr.get_conversation(
            unified_msg_origin=user_id,
            conversation_id=cid,
        )
        if not conversation:
            raise ConversationServiceError("对话不存在")

        persona_id = payload.get("persona_id", conversation.persona_id)

        if title is not None or persona_id is not None:
            await self.conv_mgr.update_conversation(
                unified_msg_origin=user_id,
                conversation_id=cid,
                title=title,
                persona_id=persona_id,
            )
        return {"message": "对话信息更新成功"}

    async def delete_conversation(self, data: object) -> dict:
        payload = self._payload(data)
        if "conversations" in payload:
            return await self._delete_conversations(payload.get("conversations", []))

        user_id, cid = self._require_user_and_cid(payload)
        await self.conv_mgr.delete_conversation(
            unified_msg_origin=user_id,
            conversation_id=cid,
        )
        return {"message": "对话删除成功"}

    async def update_history(self, data: object) -> dict:
        payload = self._payload(data)
        user_id, cid = self._require_user_and_cid(payload)
        history = payload.get("history")

        if history is None:
            raise ConversationServiceError("缺少必要参数: history")

        history = self._normalize_history(history)

        conversation = await self.conv_mgr.get_conversation(
            unified_msg_origin=user_id,
            conversation_id=cid,
        )
        if not conversation:
            raise ConversationServiceError("对话不存在")

        await self.conv_mgr.update_conversation(
            unified_msg_origin=user_id,
            conversation_id=cid,
            history=history,
        )

        return {"message": "对话历史更新成功"}

    async def export_conversations(self, data: object) -> ConversationExport:
        payload = self._payload(data)
        conversations_to_export = payload.get("conversations", [])

        if not conversations_to_export:
            raise ConversationServiceError("导出列表不能为空")

        jsonl_lines = []
        exported_count = 0
        failed_items = []

        for conv_info in conversations_to_export:
            user_id = conv_info.get("user_id")
            cid = conv_info.get("cid")

            if not user_id or not cid:
                failed_items.append(f"user_id:{user_id}, cid:{cid} - 缺少必要参数")
                continue

            try:
                conversation = await self.conv_mgr.get_conversation(
                    unified_msg_origin=user_id,
                    conversation_id=cid,
                )

                if not conversation:
                    failed_items.append(f"user_id:{user_id}, cid:{cid} - 对话不存在")
                    continue

                content = json.loads(conversation.history)
                export_record = {
                    "cid": cid,
                    "user_id": user_id,
                    "platform_id": conversation.platform_id,
                    "title": conversation.title,
                    "persona_id": conversation.persona_id,
                    "created_at": conversation.created_at,
                    "updated_at": conversation.updated_at,
                    "content": content,
                }
                jsonl_lines.append(json.dumps(export_record, ensure_ascii=False))
                exported_count += 1
            except Exception as exc:
                failed_items.append(f"user_id:{user_id}, cid:{cid} - {exc!s}")
                logger.error(
                    f"导出对话失败: user_id={user_id}, cid={cid}, error={exc!s}"
                )

        if exported_count == 0:
            raise ConversationServiceError("没有成功导出任何对话")

        jsonl_content = "\n".join(jsonl_lines)
        file_obj = BytesIO(jsonl_content.encode("utf-8"))
        file_obj.seek(0)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return ConversationExport(
            file_obj=file_obj,
            filename=f"astrbot_conversations_export_{timestamp}.jsonl",
        )

    async def _delete_conversations(self, conversations: object) -> dict:
        if not isinstance(conversations, list) or not conversations:
            raise ConversationServiceError("批量删除时conversations参数不能为空")

        deleted_count = 0
        failed_items = []

        for conv in conversations:
            if not isinstance(conv, dict):
                failed_items.append(f"{conv!r} - 格式错误")
                continue
            user_id = conv.get("user_id")
            cid = conv.get("cid")

            if not user_id or not cid:
                failed_items.append(f"user_id:{user_id}, cid:{cid} - 缺少必要参数")
                continue

            try:
                await self.conv_mgr.delete_conversation(
                    unified_msg_origin=user_id,
                    conversation_id=cid,
                )
                deleted_count += 1
            except Exception as exc:
                failed_items.append(f"user_id:{user_id}, cid:{cid} - {exc!s}")

        message = f"成功删除 {deleted_count} 个对话"
        if failed_items:
            message += f"，失败 {len(failed_items)} 个"

        return {
            "message": message,
            "deleted_count": deleted_count,
            "failed_count": len(failed_items),
            "failed_items": failed_items,
        }

    def _serialize_conversation(self, conversation, alias_map: dict) -> dict:
        return {
            **asdict(conversation),
            "umo_info": self._build_umo_info(conversation.user_id, alias_map),
        }

    @staticmethod
    def _build_umo_info(umo: str | None, alias_map: dict) -> dict:
        umo_str = umo or ""
        return {
            "umo": umo_str,
            **parse_umo(umo_str),
            **serialize_umo_alias(alias_map.get(umo_str), umo_str),
        }

    @staticmethod
    def _require_user_and_cid(payload: dict) -> tuple[str, str]:
        user_id = payload.get("user_id")
        cid = payload.get("cid")
        if not user_id or not cid:
            raise ConversationServiceError("缺少必要参数: user_id 和 cid")
        return user_id, cid

    @staticmethod
    def _normalize_history(history):
        try:
            if isinstance(history, list):
                history = json.dumps(history)
            else:
                json.loads(history)
        except json.JSONDecodeError as exc:
            raise ConversationServiceError(
                "history 必须是有效的 JSON 字符串或数组"
            ) from exc

        return json.loads(history) if isinstance(history, str) else history

    @staticmethod
    def _payload(data: object) -> dict:
        return data if isinstance(data, dict) else {}