MoviePilot/app/agent/tools/impl/query_custom_identifiers.py

"""查询自定义识别词工具"""

import json
from typing import Optional, Type

from pydantic import BaseModel, Field

from app.agent.tools.base import MoviePilotTool
from app.db.systemconfig_oper import SystemConfigOper
from app.log import logger
from app.schemas.types import SystemConfigKey


class QueryCustomIdentifiersInput(BaseModel):
    """查询自定义识别词工具的输入参数模型"""

    explanation: str = Field(
        ...,
        description="Clear explanation of why this tool is being used in the current context",
    )


class QueryCustomIdentifiersTool(MoviePilotTool):
    name: str = "query_custom_identifiers"
    description: str = (
        "Query all currently configured custom identifiers (自定义识别词). "
        "Returns the list of identifier rules used for preprocessing torrent/file names before media recognition. "
        "Use this tool to check existing rules before adding new ones to avoid duplicates."
    )
    args_schema: Type[BaseModel] = QueryCustomIdentifiersInput

    def get_tool_message(self, **kwargs) -> Optional[str]:
        """生成友好的提示消息"""
        return "查询自定义识别词"

    @staticmethod
    def _load_custom_identifiers():
        """从内存配置缓存中读取自定义识别词。"""
        return SystemConfigOper().get(SystemConfigKey.CustomIdentifiers)

    async def run(self, **kwargs) -> str:
        logger.info(f"执行工具: {self.name}")
        try:
            identifiers = self._load_custom_identifiers()
            if identifiers:
                return json.dumps(
                    {
                        "success": True,
                        "count": len(identifiers),
                        "identifiers": identifiers,
                    },
                    ensure_ascii=False,
                    indent=2,
                )
            return json.dumps(
                {
                    "success": True,
                    "count": 0,
                    "identifiers": [],
                    "message": "当前没有配置自定义识别词",
                },
                ensure_ascii=False,
                indent=2,
            )
        except Exception as e:
            logger.error(f"查询自定义识别词失败: {e}")
            return json.dumps(
                {"success": False, "message": f"查询自定义识别词时发生错误: {str(e)}"},
                ensure_ascii=False,
            )