From cbff2fed17ae6c1d29f1d0d38d850644154fe621 Mon Sep 17 00:00:00 2001 From: jxxghp Date: Mon, 30 Mar 2026 11:54:48 +0800 Subject: [PATCH] =?UTF-8?q?agent=E5=B7=A5=E5=85=B7=E5=A2=9E=E5=8A=A0?= =?UTF-8?q?=E7=AE=A1=E7=90=86=E5=91=98=E6=9D=83=E9=99=90=E6=A0=A1=E9=AA=8C?= =?UTF-8?q?=EF=BC=9A=E6=9F=A5=E8=AF=A2=E7=AB=99=E7=82=B9=E3=80=81=E6=9F=A5?= =?UTF-8?q?=E8=AF=A2=E5=B7=B2=E5=AE=89=E8=A3=85=E6=8F=92=E4=BB=B6=E3=80=81?= =?UTF-8?q?=E6=9F=A5=E8=AF=A2=E6=8F=92=E4=BB=B6=E8=83=BD=E5=8A=9B=E3=80=81?= =?UTF-8?q?=E6=9F=A5=E8=AF=A2=E7=AB=99=E7=82=B9=E7=94=A8=E6=88=B7=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E3=80=81=E5=88=AE=E5=89=8A=E5=85=83=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../tools/impl/query_installed_plugins.py | 1 + .../tools/impl/query_plugin_capabilities.py | 1 + app/agent/tools/impl/query_site_userdata.py | 129 +++++++++++------- app/agent/tools/impl/query_sites.py | 1 + app/agent/tools/impl/scrape_metadata.py | 112 +++++++++------ 5 files changed, 151 insertions(+), 93 deletions(-) diff --git a/app/agent/tools/impl/query_installed_plugins.py b/app/agent/tools/impl/query_installed_plugins.py index f79a0438..6363d64e 100644 --- a/app/agent/tools/impl/query_installed_plugins.py +++ b/app/agent/tools/impl/query_installed_plugins.py @@ -26,6 +26,7 @@ class QueryInstalledPluginsTool(MoviePilotTool): "description, version, author, running state, and other information. " "Use this tool to discover what plugins are available before querying plugin capabilities or running plugin commands." ) + require_admin: bool = True args_schema: Type[BaseModel] = QueryInstalledPluginsInput def get_tool_message(self, **kwargs) -> Optional[str]: diff --git a/app/agent/tools/impl/query_plugin_capabilities.py b/app/agent/tools/impl/query_plugin_capabilities.py index b83f0c82..264c285d 100644 --- a/app/agent/tools/impl/query_plugin_capabilities.py +++ b/app/agent/tools/impl/query_plugin_capabilities.py @@ -33,6 +33,7 @@ class QueryPluginCapabilitiesTool(MoviePilotTool): "Scheduled services are periodic tasks that can be triggered via the run_scheduler tool. " "Optionally specify a plugin_id to query a specific plugin, or omit to query all running plugins." ) + require_admin: bool = True args_schema: Type[BaseModel] = QueryPluginCapabilitiesInput def get_tool_message(self, **kwargs) -> Optional[str]: diff --git a/app/agent/tools/impl/query_site_userdata.py b/app/agent/tools/impl/query_site_userdata.py index e4ea4aa5..4a7315e0 100644 --- a/app/agent/tools/impl/query_site_userdata.py +++ b/app/agent/tools/impl/query_site_userdata.py @@ -14,60 +14,74 @@ from app.log import logger class QuerySiteUserdataInput(BaseModel): """查询站点用户数据工具的输入参数模型""" - explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context") - site_id: int = Field(..., description="The ID of the site to query user data for (can be obtained from query_sites tool)") - workdate: Optional[str] = Field(None, description="Work date to query (optional, format: 'YYYY-MM-DD', if not specified returns latest data)") + + explanation: str = Field( + ..., + description="Clear explanation of why this tool is being used in the current context", + ) + site_id: int = Field( + ..., + description="The ID of the site to query user data for (can be obtained from query_sites tool)", + ) + workdate: Optional[str] = Field( + None, + description="Work date to query (optional, format: 'YYYY-MM-DD', if not specified returns latest data)", + ) class QuerySiteUserdataTool(MoviePilotTool): name: str = "query_site_userdata" description: str = "Query user data for a specific site including username, user level, upload/download statistics, seeding information, bonus points, and other account details. Supports querying data for a specific date or latest data." + require_admin: bool = True args_schema: Type[BaseModel] = QuerySiteUserdataInput def get_tool_message(self, **kwargs) -> Optional[str]: """根据查询参数生成友好的提示消息""" site_id = kwargs.get("site_id") workdate = kwargs.get("workdate") - + message = f"正在查询站点 #{site_id} 的用户数据" if workdate: message += f" (日期: {workdate})" else: message += " (最新数据)" - + return message async def run(self, site_id: int, workdate: Optional[str] = None, **kwargs) -> str: - logger.info(f"执行工具: {self.name}, 参数: site_id={site_id}, workdate={workdate}") - + logger.info( + f"执行工具: {self.name}, 参数: site_id={site_id}, workdate={workdate}" + ) + try: # 获取数据库会话 async with AsyncSessionFactory() as db: # 获取站点 site = await Site.async_get(db, site_id) if not site: - return json.dumps({ - "success": False, - "message": f"站点不存在: {site_id}" - }, ensure_ascii=False) - + return json.dumps( + {"success": False, "message": f"站点不存在: {site_id}"}, + ensure_ascii=False, + ) + # 获取站点用户数据 user_data_list = await SiteUserData.async_get_by_domain( - db, - domain=site.domain, - workdate=workdate + db, domain=site.domain, workdate=workdate ) - + if not user_data_list: - return json.dumps({ - "success": False, - "message": f"站点 {site.name} ({site.domain}) 暂无用户数据", - "site_id": site_id, - "site_name": site.name, - "site_domain": site.domain, - "workdate": workdate - }, ensure_ascii=False) - + return json.dumps( + { + "success": False, + "message": f"站点 {site.name} ({site.domain}) 暂无用户数据", + "site_id": site_id, + "site_name": site.name, + "site_domain": site.domain, + "workdate": workdate, + }, + ensure_ascii=False, + ) + # 格式化用户数据 result = { "success": True, @@ -76,16 +90,26 @@ class QuerySiteUserdataTool(MoviePilotTool): "site_domain": site.domain, "workdate": workdate, "data_count": len(user_data_list), - "user_data": [] + "user_data": [], } - + for user_data in user_data_list: # 格式化上传/下载量（转换为可读格式） - upload_gb = user_data.upload / (1024 ** 3) if user_data.upload else 0 - download_gb = user_data.download / (1024 ** 3) if user_data.download else 0 - seeding_size_gb = user_data.seeding_size / (1024 ** 3) if user_data.seeding_size else 0 - leeching_size_gb = user_data.leeching_size / (1024 ** 3) if user_data.leeching_size else 0 - + upload_gb = user_data.upload / (1024**3) if user_data.upload else 0 + download_gb = ( + user_data.download / (1024**3) if user_data.download else 0 + ) + seeding_size_gb = ( + user_data.seeding_size / (1024**3) + if user_data.seeding_size + else 0 + ) + leeching_size_gb = ( + user_data.leeching_size / (1024**3) + if user_data.leeching_size + else 0 + ) + user_data_dict = { "domain": user_data.domain, "name": user_data.name, @@ -100,37 +124,46 @@ class QuerySiteUserdataTool(MoviePilotTool): "download_gb": round(download_gb, 2), "ratio": round(user_data.ratio, 2) if user_data.ratio else 0, "seeding": int(user_data.seeding) if user_data.seeding else 0, - "leeching": int(user_data.leeching) if user_data.leeching else 0, + "leeching": int(user_data.leeching) + if user_data.leeching + else 0, "seeding_size": user_data.seeding_size, "seeding_size_gb": round(seeding_size_gb, 2), "leeching_size": user_data.leeching_size, "leeching_size_gb": round(leeching_size_gb, 2), - "seeding_info": user_data.seeding_info if user_data.seeding_info else [], + "seeding_info": user_data.seeding_info + if user_data.seeding_info + else [], "message_unread": user_data.message_unread, - "message_unread_contents": user_data.message_unread_contents if user_data.message_unread_contents else [], + "message_unread_contents": user_data.message_unread_contents + if user_data.message_unread_contents + else [], "err_msg": user_data.err_msg, "updated_day": user_data.updated_day, - "updated_time": user_data.updated_time + "updated_time": user_data.updated_time, } result["user_data"].append(user_data_dict) - + # 如果有多条数据，只返回最新的（按更新时间排序） if len(result["user_data"]) > 1: result["user_data"].sort( - key=lambda x: (x.get("updated_day", ""), x.get("updated_time", "")), - reverse=True + key=lambda x: ( + x.get("updated_day", ""), + x.get("updated_time", ""), + ), + reverse=True, + ) + result["message"] = ( + f"找到 {len(result['user_data'])} 条数据，显示最新的一条" ) - result["message"] = f"找到 {len(result['user_data'])} 条数据，显示最新的一条" result["user_data"] = [result["user_data"][0]] - + return json.dumps(result, ensure_ascii=False, indent=2) - + except Exception as e: error_message = f"查询站点用户数据失败: {str(e)}" logger.error(f"查询站点用户数据失败: {e}", exc_info=True) - return json.dumps({ - "success": False, - "message": error_message, - "site_id": site_id - }, ensure_ascii=False) - + return json.dumps( + {"success": False, "message": error_message, "site_id": site_id}, + ensure_ascii=False, + ) diff --git a/app/agent/tools/impl/query_sites.py b/app/agent/tools/impl/query_sites.py index c8b7c8c2..d858baf6 100644 --- a/app/agent/tools/impl/query_sites.py +++ b/app/agent/tools/impl/query_sites.py @@ -29,6 +29,7 @@ class QuerySitesInput(BaseModel): class QuerySitesTool(MoviePilotTool): name: str = "query_sites" description: str = "Query site status and list all configured sites. Shows site name, domain, status, priority, and basic configuration. Site priority (pri): smaller values have higher priority (e.g., pri=1 has higher priority than pri=10)." + require_admin: bool = True args_schema: Type[BaseModel] = QuerySitesInput def get_tool_message(self, **kwargs) -> Optional[str]: diff --git a/app/agent/tools/impl/scrape_metadata.py b/app/agent/tools/impl/scrape_metadata.py index 6101e2d1..e6b6526b 100644 --- a/app/agent/tools/impl/scrape_metadata.py +++ b/app/agent/tools/impl/scrape_metadata.py @@ -16,18 +16,29 @@ from app.schemas import FileItem class ScrapeMetadataInput(BaseModel): """刮削媒体元数据工具的输入参数模型""" - explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context") - path: str = Field(..., - description="Path to the file or directory to scrape metadata for (e.g., '/path/to/file.mkv' or '/path/to/directory')") - storage: Optional[str] = Field("local", - description="Storage type: 'local' for local storage, 'smb', 'alist', etc. for remote storage (default: 'local')") - overwrite: Optional[bool] = Field(False, - description="Whether to overwrite existing metadata files (default: False)") + + explanation: str = Field( + ..., + description="Clear explanation of why this tool is being used in the current context", + ) + path: str = Field( + ..., + description="Path to the file or directory to scrape metadata for (e.g., '/path/to/file.mkv' or '/path/to/directory')", + ) + storage: Optional[str] = Field( + "local", + description="Storage type: 'local' for local storage, 'smb', 'alist', etc. for remote storage (default: 'local')", + ) + overwrite: Optional[bool] = Field( + False, + description="Whether to overwrite existing metadata files (default: False)", + ) class ScrapeMetadataTool(MoviePilotTool): name: str = "scrape_metadata" description: str = "Generate metadata files (NFO files, posters, backgrounds, etc.) for existing media files or directories. Automatically recognizes media information from the file path and creates metadata files. Supports both local and remote storage. Use 'search_media' to search TMDB database, or 'recognize_media' to extract info from torrent titles/file paths without generating files." + require_admin: bool = True args_schema: Type[BaseModel] = ScrapeMetadataInput def get_tool_message(self, **kwargs) -> Optional[str]: @@ -44,33 +55,38 @@ class ScrapeMetadataTool(MoviePilotTool): return message - async def run(self, path: str, storage: Optional[str] = "local", - overwrite: Optional[bool] = False, **kwargs) -> str: - logger.info(f"执行工具: {self.name}, 参数: path={path}, storage={storage}, overwrite={overwrite}") + async def run( + self, + path: str, + storage: Optional[str] = "local", + overwrite: Optional[bool] = False, + **kwargs, + ) -> str: + logger.info( + f"执行工具: {self.name}, 参数: path={path}, storage={storage}, overwrite={overwrite}" + ) try: # 验证路径 if not path: - return json.dumps({ - "success": False, - "message": "刮削路径不能为空" - }, ensure_ascii=False) + return json.dumps( + {"success": False, "message": "刮削路径不能为空"}, + ensure_ascii=False, + ) # 创建 FileItem fileitem = FileItem( - storage=storage, - path=path, - type="file" if Path(path).suffix else "dir" + storage=storage, path=path, type="file" if Path(path).suffix else "dir" ) # 检查本地存储路径是否存在 if storage == "local": scrape_path = Path(path) if not scrape_path.exists(): - return json.dumps({ - "success": False, - "message": f"刮削路径不存在: {path}" - }, ensure_ascii=False) + return json.dumps( + {"success": False, "message": f"刮削路径不存在: {path}"}, + ensure_ascii=False, + ) # 识别媒体信息 media_chain = MediaChain() @@ -79,11 +95,14 @@ class ScrapeMetadataTool(MoviePilotTool): mediainfo = await media_chain.async_recognize_by_meta(meta) if not mediainfo: - return json.dumps({ - "success": False, - "message": f"刮削失败，无法识别媒体信息: {path}", - "path": path - }, ensure_ascii=False) + return json.dumps( + { + "success": False, + "message": f"刮削失败，无法识别媒体信息: {path}", + "path": path, + }, + ensure_ascii=False, + ) # 在线程池中执行同步的刮削操作 await global_vars.loop.run_in_executor( @@ -92,28 +111,31 @@ class ScrapeMetadataTool(MoviePilotTool): fileitem=fileitem, meta=meta, mediainfo=mediainfo, - overwrite=overwrite - ) + overwrite=overwrite, + ), ) - return json.dumps({ - "success": True, - "message": f"{path} 刮削完成", - "path": path, - "media_info": { - "title": mediainfo.title, - "year": mediainfo.year, - "type": mediainfo.type.value if mediainfo.type else None, - "tmdb_id": mediainfo.tmdb_id, - "season": mediainfo.season - } - }, ensure_ascii=False, indent=2) + return json.dumps( + { + "success": True, + "message": f"{path} 刮削完成", + "path": path, + "media_info": { + "title": mediainfo.title, + "year": mediainfo.year, + "type": mediainfo.type.value if mediainfo.type else None, + "tmdb_id": mediainfo.tmdb_id, + "season": mediainfo.season, + }, + }, + ensure_ascii=False, + indent=2, + ) except Exception as e: error_message = f"刮削媒体元数据失败: {str(e)}" logger.error(f"刮削媒体元数据失败: {e}", exc_info=True) - return json.dumps({ - "success": False, - "message": error_message, - "path": path - }, ensure_ascii=False) + return json.dumps( + {"success": False, "message": error_message, "path": path}, + ensure_ascii=False, + )