"""添加订阅工具""" from typing import List, Optional, Type from pydantic import BaseModel, Field from app.agent.tools.base import MoviePilotTool from app.agent.tools.tags import ToolTag from app.chain.subscribe import SubscribeChain from app.db.user_oper import UserOper from app.log import logger from app.schemas.types import MUSIC_ENTITY_ALBUM, MediaSource, MediaType, MessageChannel from ._music_utils import normalize_music_type class AddSubscribeInput(BaseModel): """添加订阅工具的输入参数模型""" title: str = Field( ..., description="The title of the media to subscribe to (e.g., 'The Matrix', 'Breaking Bad')", ) year: Optional[str] = Field( None, description="Release year. Recommended for video; optional when a stable music ID is provided", ) media_type: str = Field(..., description="Allowed values: movie, tv, music") music_type: Optional[str] = Field( None, description="Required for music subscriptions: recording for one track or album for the complete album. Artists cannot be subscribed", ) 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." ), ) media_source: Optional[MediaSource] = Field( None, description="Media metadata source returned by search_media", ) media_id: Optional[str] = Field(None, description="Native ID for media_source") 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')", ) audio_quality: Optional[str] = Field( None, description="Music quality tier filter: hires, lossless, lossy, or a regular-expression combination", ) audio_format: Optional[str] = Field( None, description="Music audio-format filter as a regular expression, e.g. FLAC|ALAC|DSD", ) min_bitrate: Optional[int] = Field(None, description="Minimum music bitrate in bits per second") min_bit_depth: Optional[int] = Field(None, description="Minimum music bit depth") min_sample_rate: Optional[int] = Field(None, description="Minimum music sample rate in Hz") best_version: Optional[int] = Field( None, description="Enable quality upgrades: 0 for no, 1 for yes. Music upgrades use normalized audio quality", ) 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" tags: list[str] = [ ToolTag.Write, ToolTag.Subscription, ToolTag.Media, ] description: str = ( "Add automated subscriptions for movies, TV shows, one music recording, or one complete album. " "Music requires the source-native ID and music_type='recording' or 'album'; artists are browse-only. " "Album subscriptions only complete after a resource is confirmed to cover the album's expected track count. " "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", "") music_type = kwargs.get("music_type") season = kwargs.get("season") message = f"添加订阅: {title}" if year: message += f" ({year})" if media_type: message += f" [{media_type}]" if music_type: message += f" [{music_type}]" if season is not None: message += f" 第{season}季" elif media_type == "tv": message += " 第1季(默认)" return message async def _resolve_subscribe_username(self) -> Optional[str]: """优先映射为系统用户名，未绑定时回退当前渠道用户名。""" resolved_username = self._username if not self._channel or not self._user_id: return resolved_username try: channel = MessageChannel(self._channel) except ValueError: return resolved_username binding_keys = { MessageChannel.Telegram: ("telegram_userid",), MessageChannel.Discord: ("discord_userid",), MessageChannel.Wechat: ("wechat_userid",), MessageChannel.Feishu: ("feishu_userid", "feishu_openid"), MessageChannel.WechatClawBot: ("wechatclawbot_userid",), MessageChannel.Slack: ("slack_userid",), MessageChannel.VoceChat: ("vocechat_userid",), MessageChannel.SynologyChat: ("synologychat_userid",), MessageChannel.QQ: ("qq_userid", "qq_openid"), }.get(channel) if not binding_keys: return resolved_username mapped_username = await self.run_blocking( "db", UserOper().get_name, **{key: self._user_id for key in binding_keys}, ) return mapped_username or resolved_username async def run( self, title: str, media_type: str, year: Optional[str] = None, music_type: Optional[str] = None, season: Optional[int] = None, media_source: Optional[MediaSource] = None, media_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, audio_quality: Optional[str] = None, audio_format: Optional[str] = None, min_bitrate: Optional[int] = None, min_bit_depth: Optional[int] = None, min_sample_rate: Optional[int] = None, best_version: Optional[int] = 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"music_type={music_type}, season={season}, media_source={media_source}, " f"media_id={media_id}, " f"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', 'music'" normalized_music_type = None if media_type_enum == MediaType.MUSIC: normalized_music_type = normalize_music_type( music_type, allow_artist=False, ) if not normalized_music_type: return ( "错误：音乐订阅必须提供 music_type='recording'（单曲）" "或 music_type='album'（整张专辑），艺术家不能订阅" ) if not media_source or not media_id: return "错误：音乐订阅必须同时提供 search_media 返回的 media_source 和 media_id" if any(value is not None for value in (season, start_episode, total_episode)): return "错误：音乐订阅没有季集参数，不能传入 season、start_episode 或 total_episode" elif music_type: return "错误：music_type 仅能与 media_type='music' 一起使用" if (media_source is None) != (media_id is None): return "错误：media_source 和 media_id 必须同时提供" audio_filter_values = ( audio_quality, audio_format, min_bitrate, min_bit_depth, min_sample_rate ) if media_type_enum != MediaType.MUSIC and any( value is not None for value in audio_filter_values ): return "错误：audio_quality、audio_format 和音频技术参数仅用于音乐订阅" effective_season = ( season if season is not None else 1 if media_type_enum == MediaType.TV else None ) subscribe_username = await self._resolve_subscribe_username() # 构建额外的订阅参数 subscribe_kwargs = {} if normalized_music_type: subscribe_kwargs["music_type"] = normalized_music_type 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 audio_quality: subscribe_kwargs["audio_quality"] = audio_quality if audio_format: subscribe_kwargs["audio_format"] = audio_format if min_bitrate is not None: subscribe_kwargs["min_bitrate"] = min_bitrate if min_bit_depth is not None: subscribe_kwargs["min_bit_depth"] = min_bit_depth if min_sample_rate is not None: subscribe_kwargs["min_sample_rate"] = min_sample_rate if best_version is not None: subscribe_kwargs["best_version"] = best_version 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 or "", media_source=media_source, media_id=media_id, season=season, username=subscribe_username, **subscribe_kwargs, ) if sid: display_year = f" ({year})" if year else "" music_label = ( "专辑" if normalized_music_type == MUSIC_ENTITY_ALBUM else "单曲" ) if normalized_music_type else "" if message and "已存在" in message: result_msg = f"{music_label}订阅已存在：{title}{display_year}" if effective_season is not None: result_msg += f" 第{effective_season}季" result_msg += "。如需修改参数请先删除旧订阅。" return result_msg result_msg = f"成功添加{music_label}订阅：{title}{display_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 audio_quality: params.append(f"音质等级: {audio_quality}") if audio_format: params.append(f"音频格式: {audio_format}") if min_bitrate is not None: params.append(f"最低码率: {round(min_bitrate / 1000)}kbps") if min_bit_depth is not None: params.append(f"最低位深: {min_bit_depth}bit") if min_sample_rate is not None: params.append(f"最低采样率: {min_sample_rate / 1000:g}kHz") if best_version is not None: params.append(f"音质洗版: {'开启' if best_version else '关闭'}") 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)}"