Enhance agent workflows and tools: unify subscription and download processes, add site querying functionality, and improve error handling in download operations.

app/agent/tools/impl/add_download.py

@@ -8,6 +8,7 @@ from app.agent.tools.base import MoviePilotTool, ToolChain
 from app.chain.download import DownloadChain
 from app.core.context import Context
 from app.core.metainfo import MetaInfo
+from app.db.site_oper import SiteOper
 from app.log import logger
 from app.schemas import TorrentInfo
 
@@ -15,6 +16,7 @@ from app.schemas import TorrentInfo
 class AddDownloadInput(BaseModel):
     """添加下载工具的输入参数模型"""
     explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
+    site_name: str = Field(..., description="Name of the torrent site/source (e.g., 'The Pirate Bay')")
     torrent_title: str = Field(...,
                                description="The display name/title of the torrent (e.g., 'The.Matrix.1999.1080p.BluRay.x264')")
     torrent_url: str = Field(..., description="Direct URL to the torrent file (.torrent) or magnet link")
@@ -33,7 +35,7 @@ class AddDownloadTool(MoviePilotTool):
     description: str = "Add torrent download task to the configured downloader (qBittorrent, Transmission, etc.). Downloads the torrent file and starts the download process with specified settings."
     args_schema: Type[BaseModel] = AddDownloadInput
 
-    async def run(self, torrent_title: str, torrent_url: str, torrent_description: Optional[str] = None,
+    async def run(self, site_name: str, torrent_title: str, torrent_url: str, torrent_description: Optional[str] = None,
                   downloader: Optional[str] = None, save_path: Optional[str] = None,
                   labels: Optional[str] = None, **kwargs) -> str:
         logger.info(
@@ -46,16 +48,30 @@ class AddDownloadTool(MoviePilotTool):
             # 使用DownloadChain添加下载
             download_chain = DownloadChain()
 
+            # 根据站点名称查询站点cookie
+            siteinfo = await SiteOper().async_get_by_name(site_name)
+            if not siteinfo:
+                return f"错误：未找到站点信息：{site_name}"
+
             # 创建下载上下文
             torrent_info = TorrentInfo(
                 title=torrent_title,
-                download_url=torrent_url
+                download_url=torrent_url,
+                site_name=site_name,
+                site_ua=siteinfo.ua,
+                site_cookie=siteinfo.cookie,
+                site_proxy=siteinfo.proxy,
+                site_order=siteinfo.pri,
+                site_downloader=siteinfo.downloader
             )
             meta_info = MetaInfo(title=torrent_title, subtitle=torrent_description)
+            media_info = ToolChain().recognize_media(meta=meta_info)
+            if not media_info:
+                return "错误：无法识别媒体信息，无法添加下载任务"
             context = Context(
                 torrent_info=torrent_info,
                 meta_info=meta_info,
-                media_info=ToolChain().recognize_media(meta=meta_info)
+                media_info=media_info
             )
 
             did = download_chain.download_single(

app/agent/tools/impl/query_media_library.py

@@ -17,14 +17,14 @@ class QueryMediaLibraryInput(BaseModel):
     media_type: Optional[str] = Field("all",
                                       description="Type of media content: '电影' for films, '电视剧' for television series or anime series, 'all' for all types")
     title: Optional[str] = Field(None,
-                                 description="Specific media title to search for (optional, if provided returns detailed info for that specific media)")
+                                 description="Specific media title to check if it exists in the media library (optional, if provided checks for that specific media)")
     year: Optional[str] = Field(None,
                                 description="Release year of the media (optional, helps narrow down search results)")
 
 
 class QueryMediaLibraryTool(MoviePilotTool):
     name: str = "query_media_library"
-    description: str = "Query media library status and list all media files that have been successfully processed and added to the media server (Plex, Emby, Jellyfin). Shows library statistics and file details."
+    description: str = "Check if a specific media resource already exists in the media library (Plex, Emby, Jellyfin). Use this tool to verify whether a movie or TV series has been successfully processed and added to the media server before performing operations like downloading or subscribing."
     args_schema: Type[BaseModel] = QueryMediaLibraryInput
 
     async def run(self, media_type: Optional[str] = "all",

app/agent/tools/impl/query_sites.py

@@ -0,0 +1,66 @@
+"""查询站点工具"""
+
+import json
+from typing import Optional, Type
+
+from pydantic import BaseModel, Field
+
+from app.agent.tools.base import MoviePilotTool
+from app.db.site_oper import SiteOper
+from app.log import logger
+
+
+class QuerySitesInput(BaseModel):
+    """查询站点工具的输入参数模型"""
+    explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
+    status: Optional[str] = Field("all",
+                                  description="Filter sites by status: 'active' for enabled sites, 'inactive' for disabled sites, 'all' for all sites")
+    name: Optional[str] = Field(None,
+                                description="Filter sites by name (partial match, optional)")
+
+
+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."
+    args_schema: Type[BaseModel] = QuerySitesInput
+
+    async def run(self, status: Optional[str] = "all", name: Optional[str] = None, **kwargs) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: status={status}, name={name}")
+        try:
+            site_oper = SiteOper()
+            # 获取所有站点（按优先级排序）
+            sites = site_oper.list_order_by_pri()
+            filtered_sites = []
+            for site in sites:
+                # 按状态过滤
+                if status == "active" and not site.is_active:
+                    continue
+                if status == "inactive" and site.is_active:
+                    continue
+                # 按名称过滤（部分匹配）
+                if name and name.lower() not in (site.name or "").lower():
+                    continue
+                filtered_sites.append(site)
+            if filtered_sites:
+                # 精简字段，只保留关键信息
+                simplified_sites = []
+                for s in filtered_sites:
+                    simplified = {
+                        "id": s.id,
+                        "name": s.name,
+                        "domain": s.domain,
+                        "url": s.url,
+                        "pri": s.pri,
+                        "is_active": s.is_active,
+                        "downloader": s.downloader,
+                        "proxy": s.proxy,
+                        "timeout": s.timeout
+                    }
+                    simplified_sites.append(simplified)
+                result_json = json.dumps(simplified_sites, ensure_ascii=False, indent=2)
+                return result_json
+            return "未找到相关站点。"
+        except Exception as e:
+            logger.error(f"查询站点失败: {e}", exc_info=True)
+            return f"查询站点时发生错误: {str(e)}"
+

app/agent/tools/impl/search_torrents.py

@@ -1,6 +1,7 @@
 """搜索种子工具"""
 
 import json
+import re
 from typing import List, Optional, Type
 
 from pydantic import BaseModel, Field
@@ -23,6 +24,8 @@ class SearchTorrentsInput(BaseModel):
     season: Optional[int] = Field(None, description="Season number for TV shows (optional, only applicable for series)")
     sites: Optional[List[int]] = Field(None,
                                        description="Array of specific site IDs to search on (optional, if not provided searches all configured sites)")
+    filter_pattern: Optional[str] = Field(None,
+                                          description="Regular expression pattern to filter torrent titles by resolution, quality, or other keywords (e.g., '4K|2160p|UHD' for 4K content, '1080p|BluRay' for 1080p BluRay)")
 
 
 class SearchTorrentsTool(MoviePilotTool):
@@ -32,14 +35,23 @@ class SearchTorrentsTool(MoviePilotTool):
 
     async def run(self, title: str, year: Optional[str] = None,
                   media_type: Optional[str] = None, season: Optional[int] = None,
-                  sites: Optional[List[int]] = None, **kwargs) -> str:
+                  sites: Optional[List[int]] = None, filter_pattern: Optional[str] = None, **kwargs) -> str:
         logger.info(
-            f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}")
+            f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}, filter_pattern={filter_pattern}")
 
         try:
             search_chain = SearchChain()
             torrents = search_chain.search_by_title(title=title, sites=sites)
             filtered_torrents = []
+            # 编译正则表达式（如果提供）
+            regex_pattern = None
+            if filter_pattern:
+                try:
+                    regex_pattern = re.compile(filter_pattern, re.IGNORECASE)
+                except re.error as e:
+                    logger.warning(f"正则表达式编译失败: {filter_pattern}, 错误: {e}")
+                    return f"正则表达式格式错误: {str(e)}"
+            
             for torrent in torrents:
                 # torrent 是 Context 对象，需要通过 meta_info 和 media_info 访问属性
                 if year and torrent.meta_info and torrent.meta_info.year != year:
@@ -49,6 +61,10 @@ class SearchTorrentsTool(MoviePilotTool):
                         continue
                 if season and torrent.meta_info and torrent.meta_info.begin_season != season:
                     continue
+                # 使用正则表达式过滤标题（分辨率、质量等关键字）
+                if regex_pattern and torrent.torrent_info and torrent.torrent_info.title:
+                    if not regex_pattern.search(torrent.torrent_info.title):
+                        continue
                 filtered_torrents.append(torrent)
 
             if filtered_torrents: