"""查询媒体库工具""" import asyncio import json from collections import OrderedDict from typing import Optional, Type, Any from pydantic import BaseModel, Field from app.agent.tools.base import MoviePilotTool from app.agent.tools.tags import ToolTag from app.chain.mediaserver import MediaServerChain from app.helper.mediaserver import MediaServerHelper from app.log import logger from app.schemas.types import MUSIC_ENTITY_ALBUM, MediaType, media_type_to_agent from ._music_utils import normalize_music_type def _sort_seasons(seasons: Optional[dict]) -> dict: """按季号、集号升序整理季集信息,保证输出稳定。""" if not seasons: return {} def _sort_key(value): try: return int(value) except (TypeError, ValueError): return str(value) return OrderedDict( (season, sorted(episodes, key=_sort_key)) for season, episodes in sorted(seasons.items(), key=lambda item: _sort_key(item[0])) ) def _filter_regular_seasons(seasons: Optional[dict]) -> OrderedDict: """仅保留正片季,忽略 season 0 等特殊季。""" sorted_seasons = _sort_seasons(seasons) regular_seasons = OrderedDict() for season, episodes in sorted_seasons.items(): try: season_number = int(season) except (TypeError, ValueError): continue if season_number > 0: regular_seasons[season_number] = episodes return regular_seasons def _build_tv_server_result(existing_seasons: OrderedDict, total_seasons: OrderedDict) -> dict[str, Any]: """构建单个服务器的电视剧存在性结果。""" seasons_result = OrderedDict() missing_seasons = [] all_seasons = sorted(set(total_seasons.keys()) | set(existing_seasons.keys())) for season in all_seasons: existing_episodes = existing_seasons.get(season, []) total_episodes = total_seasons.get(season) if total_episodes is not None: missing_episodes = [episode for episode in total_episodes if episode not in existing_episodes] total_episode_count = len(total_episodes) else: missing_episodes = None total_episode_count = None seasons_result[str(season)] = { "existing_episodes": existing_episodes, "total_episodes": total_episode_count, "missing_episodes": missing_episodes } if total_episodes is not None and not existing_episodes: missing_seasons.append(season) return { "seasons": seasons_result, "missing_seasons": missing_seasons } class QueryLibraryExistsInput(BaseModel): """查询媒体库工具的输入参数模型""" tmdb_id: Optional[int] = Field(None, description="TMDB media ID") douban_id: Optional[str] = Field(None, description="Douban media ID") bangumi_id: Optional[int] = Field(None, description="Bangumi media ID") anilist_id: Optional[int] = Field(None, description="AniList media ID") media_source: Optional[str] = Field(None, description="Media metadata source") media_id: Optional[str] = Field(None, description="Native ID for media_source") media_type: Optional[str] = Field(None, description="Allowed values: movie, tv, music") music_type: Optional[str] = Field( None, description="For music: recording or album. Artists are not library acquisition targets", ) class QueryLibraryExistsTool(MoviePilotTool): """查询影视、单曲或完整专辑是否已存在于媒体服务器。""" name: str = "query_library_exists" tags: list[str] = [ ToolTag.Read, ToolTag.Library, ToolTag.Media, ] description: str = ( "Check whether movies, TV, music recordings, or complete albums exist on configured media servers. " "TV results include episode coverage. Album existence is true only when the server confirms complete " "track coverage for the expected album size; artists cannot be queried as acquisition targets." ) args_schema: Type[BaseModel] = QueryLibraryExistsInput def get_tool_message(self, **kwargs) -> Optional[str]: """根据查询参数生成友好的提示消息""" media_type = kwargs.get("media_type") identities = ( ("TMDB", kwargs.get("tmdb_id")), ("豆瓣", kwargs.get("douban_id")), ("Bangumi", kwargs.get("bangumi_id")), ("AniList", kwargs.get("anilist_id")), (kwargs.get("media_source") or "媒体源", kwargs.get("media_id")), ) label, identity = next( ((label, identity) for label, identity in identities if identity is not None), (None, None), ) message = f"查询媒体库: {label}={identity}" if label else "查询媒体库" if media_type: message += f" [{media_type}]" return message @staticmethod def _get_media_server_names() -> list[str]: """同步读取已加载媒体服务器名称。""" return sorted(MediaServerHelper().get_services().keys()) @staticmethod def _query_media_exists(mediainfo, server: Optional[str] = None): """同步查询单个媒体服务器的存在性信息。""" return MediaServerChain().media_exists(mediainfo=mediainfo, server=server) async def run(self, tmdb_id: Optional[int] = None, douban_id: Optional[str] = None, bangumi_id: Optional[int] = None, anilist_id: Optional[int] = None, media_source: Optional[str] = None, media_id: Optional[str] = None, media_type: Optional[str] = None, music_type: Optional[str] = None, **kwargs) -> str: """识别精确媒体身份并按服务器汇总存在性或完整性。""" logger.info(f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}, douban_id={douban_id}, media_type={media_type}") try: if not any((tmdb_id, douban_id, bangumi_id, anilist_id, media_id)): return "参数错误:至少需要提供一个媒体 ID,请先使用 search_media 工具获取媒体信息。" media_type_enum = None if media_type: 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 music_type: normalized_music_type = normalize_music_type( music_type, allow_artist=False, ) if not normalized_music_type: return ( f"错误:无效的音乐实体类型 '{music_type}'," "支持的类型:'recording', 'album'" ) if media_type_enum != MediaType.MUSIC: return "错误:music_type 仅能与 media_type='music' 一起使用" if media_type_enum == MediaType.MUSIC and (not media_source or not media_id): return "错误:音乐媒体库查询必须同时提供 media_source 和 media_id" media_chain = MediaServerChain() mediainfo = await media_chain.async_recognize_media( tmdbid=tmdb_id, doubanid=douban_id, bangumiid=bangumi_id, anilistid=anilist_id, source=media_source, mediaid=media_id, mtype=media_type_enum, ) if not mediainfo: identity = media_id or tmdb_id or douban_id or bangumi_id or anilist_id return f"未识别到媒体信息: {identity}" if normalized_music_type and ( getattr(mediainfo, "music_type", None) != normalized_music_type ): return ( f"音乐实体类型不匹配:请求 {normalized_music_type}," f"实际 {getattr(mediainfo, 'music_type', None) or 'unknown'}" ) # 2. 遍历所有媒体服务器,分别查询存在性信息 server_results = OrderedDict() total_seasons = _filter_regular_seasons( mediainfo.seasons if mediainfo.type == MediaType.TV else None ) service_names = self._get_media_server_names() server_checks = await asyncio.gather( *[ self.run_blocking( "mediaserver", self._query_media_exists, mediainfo, service_name, ) for service_name in service_names ] ) for service_name, existsinfo in zip(service_names, server_checks): if not existsinfo: continue if existsinfo.type == MediaType.TV: existing_seasons = _filter_regular_seasons(existsinfo.seasons) server_results[service_name] = _build_tv_server_result( existing_seasons=existing_seasons, total_seasons=total_seasons ) else: server_results[service_name] = { "exists": True, **( { "complete": True, "expected_tracks": mediainfo.total_tracks, } if mediainfo.type == MediaType.MUSIC and getattr(mediainfo, "music_type", None) == MUSIC_ENTITY_ALBUM else {} ), } if not server_results: global_existsinfo = await self.run_blocking( "mediaserver", self._query_media_exists, mediainfo, None ) if not global_existsinfo: return "媒体库中未找到相关媒体" fallback_server_name = global_existsinfo.server or "local" if global_existsinfo.type == MediaType.TV: server_results[fallback_server_name] = _build_tv_server_result( existing_seasons=_filter_regular_seasons(global_existsinfo.seasons), total_seasons=total_seasons ) else: server_results[fallback_server_name] = { "exists": True, **( { "complete": True, "expected_tracks": mediainfo.total_tracks, } if mediainfo.type == MediaType.MUSIC and getattr(mediainfo, "music_type", None) == MUSIC_ENTITY_ALBUM else {} ), } # 3. 组装统一的存在性结果,不查询媒体服务器详情 result_dict = { "title": mediainfo.title, "year": mediainfo.year, "type": media_type_to_agent(mediainfo.type), "servers": server_results } if mediainfo.type == MediaType.MUSIC: result_dict.update({ "music_type": mediainfo.music_type, "artists": list(mediainfo.artists or []), "album": mediainfo.album, "total_tracks": mediainfo.total_tracks, "media_source": mediainfo.source, "media_id": mediainfo.media_id, }) return json.dumps([result_dict], ensure_ascii=False) except Exception as e: logger.error(f"查询媒体库失败: {e}", exc_info=True) return f"查询媒体库时发生错误: {str(e)}"