mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-06-28 03:02:34 +08:00
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.
This commit is contained in:
jxxghp
@@ -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"
|
||||
]
|
||||
|
||||
@@ -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,
|
||||
|
||||
162
app/agent/tools/impl/recognize_media.py
Normal file
162
app/agent/tools/impl/recognize_media.py
Normal file
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user