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 not specified will subscribe to all seasons)")
    tmdb_id: Optional[int] = Field(None,
                                   description="TMDB database ID for precise media identification (optional, can be obtained from search_media tool)")
    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. 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}季"
        
        return message

    async def run(self, title: str, year: str, media_type: str,
                  season: Optional[int] = None, tmdb_id: Optional[int] = 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}, 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_key = media_type.strip().lower()
            if media_type_key == "movie":
                media_type_enum = MediaType.MOVIE
            elif media_type_key == "tv":
                media_type_enum = MediaType.TV
            else:
                return f"错误：无效的媒体类型 '{media_type}'，支持的类型：'movie', 'tv'"

            # 构建额外的订阅参数
            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,
                season=season,
                username=self._user_id,
                **subscribe_kwargs
            )
            if sid:
                if message and "已存在" in message:
                    return f"订阅已存在：{title} ({year})。如需修改参数请先删除旧订阅。"

                result_msg = f"成功添加订阅：{title} ({year})"
                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)}"