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

"""添加订阅工具"""

from typing import Optional, Type, List

from pydantic import BaseModel, Field

from app.agent.tools.base import MoviePilotTool
from app.chain.subscribe import SubscribeChain
from app.log import logger
from app.schemas.types import MediaType


class AddSubscribeInput(BaseModel):
    """添加订阅工具的输入参数模型"""

    explanation: str = Field(
        ...,
        description="Clear explanation of why this tool is being used in the current context",
    )
    title: str = Field(
        ...,
        description="The title of the media to subscribe to (e.g., 'The Matrix', 'Breaking Bad')",
    )
    year: str = Field(
        ...,
        description="Release year of the media (required for accurate identification)",
    )
    media_type: str = Field(..., description="Allowed values: movie, tv")
    season: Optional[int] = Field(
        None,
        description=(
            "Season number for TV shows (optional). If omitted, the subscription defaults to season 1 only. "
            "To subscribe multiple seasons or the full series, call this tool separately for each season."
        ),
    )
    tmdb_id: Optional[int] = Field(
        None,
        description="TMDB database ID for precise media identification (optional, can be obtained from search_media tool)",
    )
    douban_id: Optional[str] = Field(
        None,
        description="Douban ID for precise media identification (optional, alternative to tmdb_id)",
    )
    start_episode: Optional[int] = Field(
        None,
        description="Starting episode number for TV shows (optional, defaults to 1 if not specified)",
    )
    total_episode: Optional[int] = Field(
        None,
        description="Total number of episodes for TV shows (optional, will be auto-detected from TMDB if not specified)",
    )
    quality: Optional[str] = Field(
        None,
        description="Quality filter as regular expression (optional, e.g., 'BluRay|WEB-DL|HDTV')",
    )
    resolution: Optional[str] = Field(
        None,
        description="Resolution filter as regular expression (optional, e.g., '1080p|720p|2160p')",
    )
    effect: Optional[str] = Field(
        None,
        description="Effect filter as regular expression (optional, e.g., 'HDR|DV|SDR')",
    )
    filter_groups: Optional[List[str]] = Field(
        None,
        description="List of filter rule group names to apply (optional, can be obtained from query_rule_groups tool)",
    )
    sites: Optional[List[int]] = Field(
        None,
        description="List of site IDs to search from (optional, can be obtained from query_sites tool)",
    )


class AddSubscribeTool(MoviePilotTool):
    name: str = "add_subscribe"
    description: str = (
        "Add media subscription to create automated download rules for movies and TV shows. "
        "The system will automatically search and download new episodes or releases based on the subscription criteria. "
        "For TV shows, omitting `season` subscribes season 1 only by default; to subscribe multiple seasons or "
        "the full series, call this tool once per season. Supports advanced filtering options like quality, "
        "resolution, and effect filters using regular expressions."
    )
    args_schema: Type[BaseModel] = AddSubscribeInput

    def get_tool_message(self, **kwargs) -> Optional[str]:
        """根据订阅参数生成友好的提示消息"""
        title = kwargs.get("title", "")
        year = kwargs.get("year", "")
        media_type = kwargs.get("media_type", "")
        season = kwargs.get("season")

        message = f"添加订阅: {title}"
        if year:
            message += f" ({year})"
        if media_type:
            message += f" [{media_type}]"
        if season:
            message += f" 第{season}季"
        elif media_type == "tv":
            message += " 第1季(默认)"

        return message

    async def run(
        self,
        title: str,
        year: str,
        media_type: str,
        season: Optional[int] = None,
        tmdb_id: Optional[int] = None,
        douban_id: Optional[str] = None,
        start_episode: Optional[int] = None,
        total_episode: Optional[int] = None,
        quality: Optional[str] = None,
        resolution: Optional[str] = None,
        effect: Optional[str] = None,
        filter_groups: Optional[List[str]] = None,
        sites: Optional[List[int]] = None,
        **kwargs,
    ) -> str:
        logger.info(
            f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, "
            f"season={season}, tmdb_id={tmdb_id}, douban_id={douban_id}, start_episode={start_episode}, "
            f"total_episode={total_episode}, quality={quality}, resolution={resolution}, "
            f"effect={effect}, filter_groups={filter_groups}, sites={sites}"
        )

        try:
            subscribe_chain = SubscribeChain()
            media_type_enum = MediaType.from_agent(media_type)
            if not media_type_enum:
                return f"错误：无效的媒体类型 '{media_type}'，支持的类型：'movie', 'tv'"
            effective_season = (
                season
                if season is not None
                else 1
                if media_type_enum == MediaType.TV
                else None
            )

            # 构建额外的订阅参数
            subscribe_kwargs = {}
            if start_episode is not None:
                subscribe_kwargs["start_episode"] = start_episode
            if total_episode is not None:
                subscribe_kwargs["total_episode"] = total_episode
            if quality:
                subscribe_kwargs["quality"] = quality
            if resolution:
                subscribe_kwargs["resolution"] = resolution
            if effect:
                subscribe_kwargs["effect"] = effect
            if filter_groups:
                subscribe_kwargs["filter_groups"] = filter_groups
            if sites:
                subscribe_kwargs["sites"] = sites

            sid, message = await subscribe_chain.async_add(
                mtype=media_type_enum,
                title=title,
                year=year,
                tmdbid=tmdb_id,
                doubanid=douban_id,
                season=season,
                username=self._user_id,
                **subscribe_kwargs,
            )
            if sid:
                if message and "已存在" in message:
                    result_msg = f"订阅已存在：{title} ({year})"
                    if effective_season is not None:
                        result_msg += f" 第{effective_season}季"
                    result_msg += "。如需修改参数请先删除旧订阅。"
                    return result_msg

                result_msg = f"成功添加订阅：{title} ({year})"
                if effective_season is not None:
                    result_msg += f" 第{effective_season}季"
                    if season is None:
                        result_msg += "（未指定季号，默认按第一季订阅）"
                if subscribe_kwargs:
                    params = []
                    if start_episode is not None:
                        params.append(f"开始集数: {start_episode}")
                    if total_episode is not None:
                        params.append(f"总集数: {total_episode}")
                    if quality:
                        params.append(f"质量过滤: {quality}")
                    if resolution:
                        params.append(f"分辨率过滤: {resolution}")
                    if effect:
                        params.append(f"特效过滤: {effect}")
                    if filter_groups:
                        params.append(f"规则组: {', '.join(filter_groups)}")
                    if sites:
                        params.append(f"站点: {', '.join(map(str, sites))}")
                    if params:
                        result_msg += f"\n配置参数: {', '.join(params)}"
                return result_msg
            else:
                return f"添加订阅失败：{message}"
        except Exception as e:
            logger.error(f"添加订阅失败: {e}", exc_info=True)
            return f"添加订阅时发生错误: {str(e)}"