From a8c6516b311d7b95852e1cb03122802bebf83ef0 Mon Sep 17 00:00:00 2001 From: jxxghp Date: Tue, 18 Nov 2025 16:07:11 +0800 Subject: [PATCH] refactor: remove deprecated tools and add RecognizeMediaTool - Removed the old tool definitions from __init__.py to streamline the module. - Added RecognizeMediaTool to factory.py for enhanced media recognition capabilities. - Updated tool exports to reflect the changes in available tools. --- app/agent/tools/__init__.py | 61 --------- app/agent/tools/factory.py | 2 + app/agent/tools/impl/recognize_media.py | 162 ++++++++++++++++++++++++ 3 files changed, 164 insertions(+), 61 deletions(-) create mode 100644 app/agent/tools/impl/recognize_media.py diff --git a/app/agent/tools/__init__.py b/app/agent/tools/__init__.py index 6ade0c60..e69de29b 100644 --- a/app/agent/tools/__init__.py +++ b/app/agent/tools/__init__.py @@ -1,61 +0,0 @@ -"""MoviePilot工具模块""" - -from .base import MoviePilotTool -from app.agent.tools.impl.search_media import SearchMediaTool -from app.agent.tools.impl.add_subscribe import AddSubscribeTool -from app.agent.tools.impl.search_torrents import SearchTorrentsTool -from app.agent.tools.impl.add_download import AddDownloadTool -from app.agent.tools.impl.query_subscribes import QuerySubscribesTool -from app.agent.tools.impl.query_subscribe_shares import QuerySubscribeSharesTool -from app.agent.tools.impl.query_popular_subscribes import QueryPopularSubscribesTool -from app.agent.tools.impl.query_subscribe_history import QuerySubscribeHistoryTool -from app.agent.tools.impl.delete_subscribe import DeleteSubscribeTool -from app.agent.tools.impl.query_downloads import QueryDownloadsTool -from app.agent.tools.impl.delete_download import DeleteDownloadTool -from app.agent.tools.impl.query_downloaders import QueryDownloadersTool -from app.agent.tools.impl.query_sites import QuerySitesTool -from app.agent.tools.impl.test_site import TestSiteTool -from app.agent.tools.impl.get_recommendations import GetRecommendationsTool -from app.agent.tools.impl.query_media_library import QueryMediaLibraryTool -from app.agent.tools.impl.query_directories import QueryDirectoriesTool -from app.agent.tools.impl.list_directory import ListDirectoryTool -from app.agent.tools.impl.query_transfer_history import QueryTransferHistoryTool -from app.agent.tools.impl.transfer_file import TransferFileTool -from app.agent.tools.impl.send_message import SendMessageTool -from app.agent.tools.impl.query_schedulers import QuerySchedulersTool -from app.agent.tools.impl.run_scheduler import RunSchedulerTool -from app.agent.tools.impl.query_workflows import QueryWorkflowsTool -from app.agent.tools.impl.run_workflow import RunWorkflowTool -from app.agent.tools.impl.update_site_cookie import UpdateSiteCookieTool -from .factory import MoviePilotToolFactory - -__all__ = [ - "MoviePilotTool", - "SearchMediaTool", - "AddSubscribeTool", - "SearchTorrentsTool", - "AddDownloadTool", - "QuerySubscribesTool", - "QuerySubscribeSharesTool", - "QueryPopularSubscribesTool", - "QuerySubscribeHistoryTool", - "DeleteSubscribeTool", - "QueryDownloadsTool", - "DeleteDownloadTool", - "QueryDownloadersTool", - "QuerySitesTool", - "TestSiteTool", - "UpdateSiteCookieTool", - "GetRecommendationsTool", - "QueryMediaLibraryTool", - "QueryDirectoriesTool", - "ListDirectoryTool", - "QueryTransferHistoryTool", - "TransferFileTool", - "SendMessageTool", - "QuerySchedulersTool", - "RunSchedulerTool", - "QueryWorkflowsTool", - "RunWorkflowTool", - "MoviePilotToolFactory" -] diff --git a/app/agent/tools/factory.py b/app/agent/tools/factory.py index e3eed516..449c95c8 100644 --- a/app/agent/tools/factory.py +++ b/app/agent/tools/factory.py @@ -16,6 +16,7 @@ from app.agent.tools.impl.query_popular_subscribes import QueryPopularSubscribes from app.agent.tools.impl.query_subscribe_history import QuerySubscribeHistoryTool from app.agent.tools.impl.delete_subscribe import DeleteSubscribeTool from app.agent.tools.impl.search_media import SearchMediaTool +from app.agent.tools.impl.recognize_media import RecognizeMediaTool from app.agent.tools.impl.search_torrents import SearchTorrentsTool from app.agent.tools.impl.send_message import SendMessageTool from app.agent.tools.impl.query_schedulers import QuerySchedulersTool @@ -44,6 +45,7 @@ class MoviePilotToolFactory: tools = [] tool_definitions = [ SearchMediaTool, + RecognizeMediaTool, AddSubscribeTool, SearchTorrentsTool, AddDownloadTool, diff --git a/app/agent/tools/impl/recognize_media.py b/app/agent/tools/impl/recognize_media.py new file mode 100644 index 00000000..bb352d7a --- /dev/null +++ b/app/agent/tools/impl/recognize_media.py @@ -0,0 +1,162 @@ +"""识别媒体信息工具""" + +import json +from typing import Optional, Type + +from pydantic import BaseModel, Field + +from app.agent.tools.base import MoviePilotTool +from app.chain.media import MediaChain +from app.core.context import Context +from app.core.metainfo import MetaInfo +from app.log import logger + + +class RecognizeMediaInput(BaseModel): + """识别媒体信息工具的输入参数模型""" + explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context") + title: Optional[str] = Field(None, description="The title of the torrent/media to recognize (required for torrent recognition)") + subtitle: Optional[str] = Field(None, description="The subtitle or description of the torrent (optional, helps improve recognition accuracy)") + path: Optional[str] = Field(None, description="The file path to recognize (required for file recognition, mutually exclusive with title)") + + +class RecognizeMediaTool(MoviePilotTool): + name: str = "recognize_media" + description: str = "Recognize media information from torrent titles or file paths. Supports two modes: 1) Recognize from torrent title and optional subtitle, 2) Recognize from file path. Returns detailed media information including title, year, type, TMDB ID, overview, and other metadata." + args_schema: Type[BaseModel] = RecognizeMediaInput + + def get_tool_message(self, **kwargs) -> Optional[str]: + """根据识别参数生成友好的提示消息""" + title = kwargs.get("title") + subtitle = kwargs.get("subtitle") + path = kwargs.get("path") + + if path: + message = f"正在识别文件媒体信息: {path}" + elif title: + message = f"正在识别种子媒体信息: {title}" + if subtitle: + message += f" ({subtitle})" + else: + message = "正在识别媒体信息" + + return message + + async def run(self, title: Optional[str] = None, subtitle: Optional[str] = None, + path: Optional[str] = None, **kwargs) -> str: + logger.info(f"执行工具: {self.name}, 参数: title={title}, subtitle={subtitle}, path={path}") + + try: + media_chain = MediaChain() + context = None + + # 根据提供的参数选择识别方式 + if path: + # 文件路径识别 + if not path: + return json.dumps({ + "success": False, + "message": "文件路径不能为空" + }, ensure_ascii=False) + + context = await media_chain.async_recognize_by_path(path) + if context: + return self._format_context_result(context, "文件") + else: + return json.dumps({ + "success": False, + "message": f"无法识别文件媒体信息: {path}", + "path": path + }, ensure_ascii=False) + + elif title: + # 种子标题识别 + metainfo = MetaInfo(title, subtitle) + mediainfo = await media_chain.async_recognize_by_meta(metainfo) + if mediainfo: + context = Context(meta_info=metainfo, media_info=mediainfo) + return self._format_context_result(context, "种子") + else: + return json.dumps({ + "success": False, + "message": f"无法识别种子媒体信息: {title}", + "title": title, + "subtitle": subtitle + }, ensure_ascii=False) + + else: + return json.dumps({ + "success": False, + "message": "必须提供 title（标题）或 path（文件路径）参数之一" + }, ensure_ascii=False) + + except Exception as e: + error_message = f"识别媒体信息失败: {str(e)}" + logger.error(f"识别媒体信息失败: {e}", exc_info=True) + return json.dumps({ + "success": False, + "message": error_message + }, ensure_ascii=False) + + def _format_context_result(self, context: Context, source_type: str) -> str: + """格式化识别结果为JSON字符串""" + if not context: + return json.dumps({ + "success": False, + "message": "识别结果为空" + }, ensure_ascii=False) + + context_dict = context.to_dict() + media_info = context_dict.get("media_info") + meta_info = context_dict.get("meta_info") + + # 构建简化的结果 + result = { + "success": True, + "source_type": source_type, + "media_info": None, + "meta_info": None + } + + # 处理媒体信息 + if media_info: + result["media_info"] = { + "title": media_info.get("title"), + "en_title": media_info.get("en_title"), + "year": media_info.get("year"), + "type": media_info.get("type"), + "season": media_info.get("season"), + "tmdb_id": media_info.get("tmdb_id"), + "imdb_id": media_info.get("imdb_id"), + "douban_id": media_info.get("douban_id"), + "bangumi_id": media_info.get("bangumi_id"), + "overview": media_info.get("overview"), + "vote_average": media_info.get("vote_average"), + "poster_path": media_info.get("poster_path"), + "backdrop_path": media_info.get("backdrop_path"), + "detail_link": media_info.get("detail_link"), + "title_year": media_info.get("title_year"), + "source": media_info.get("source") + } + + # 处理元数据信息 + if meta_info: + result["meta_info"] = { + "name": meta_info.get("name"), + "title": meta_info.get("title"), + "year": meta_info.get("year"), + "type": meta_info.get("type"), + "begin_season": meta_info.get("begin_season"), + "end_season": meta_info.get("end_season"), + "begin_episode": meta_info.get("begin_episode"), + "end_episode": meta_info.get("end_episode"), + "total_episode": meta_info.get("total_episode"), + "part": meta_info.get("part"), + "season_episode": meta_info.get("season_episode"), + "episode_list": meta_info.get("episode_list"), + "tmdbid": meta_info.get("tmdbid"), + "doubanid": meta_info.get("doubanid") + } + + return json.dumps(result, ensure_ascii=False, indent=2) +