feat: 优化工具和SKILL.md

app/agent/tools/impl/search_torrents.py

@@ -1,15 +1,16 @@
 """搜索种子工具"""
 
 import json
-import re
 from typing import List, Optional, Type
 
 from pydantic import BaseModel, Field
 
 from app.agent.tools.base import MoviePilotTool
 from app.chain.search import SearchChain
+from app.db.systemconfig_oper import SystemConfigOper
+from app.helper.sites import SitesHelper
 from app.log import logger
-from app.schemas.types import MediaType
+from app.schemas.types import MediaType, SystemConfigKey
 from ._torrent_search_utils import (
     SEARCH_RESULT_CACHE_FILE,
     build_filter_options,
@@ -19,87 +20,94 @@ from ._torrent_search_utils import (
 class SearchTorrentsInput(BaseModel):
     """搜索种子工具的输入参数模型"""
     explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
-    title: str = Field(...,
-                       description="The title of the media resource to search for (e.g., 'The Matrix 1999', 'Breaking Bad S01E01')")
-    year: Optional[str] = Field(None,
-                                description="Release year of the media (optional, helps narrow down search results)")
-    media_type: Optional[str] = Field(None,
-                                      description="Type of media content: '电影' for films, '电视剧' for television series or anime series")
-    season: Optional[int] = Field(None, description="Season number for TV shows (optional, only applicable for series)")
+    tmdb_id: Optional[int] = Field(None, description="TMDB ID (can be obtained from search_media tool). Either tmdb_id or douban_id must be provided.")
+    douban_id: Optional[str] = Field(None, description="Douban ID (can be obtained from search_media tool). Either tmdb_id or douban_id must be provided.")
+    media_type: Optional[str] = Field(None, description="Allowed values: movie, tv")
+    area: Optional[str] = Field(None, description="Search scope: 'title' (default) or 'imdbid'")
     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):
     name: str = "search_torrents"
-    description: str = "Search for torrent files across configured indexer sites based on media information. Returns available frontend-style filter options for the most recent search and caches the underlying results for get_search_results."
+    description: str = ("Search for torrent files by media ID across configured indexer sites, cache the matched results, "
+                        "and return available filter options for follow-up selection. "
+                        "Requires tmdb_id or douban_id (can be obtained from search_media tool) for accurate matching.")
     args_schema: Type[BaseModel] = SearchTorrentsInput
 
     def get_tool_message(self, **kwargs) -> Optional[str]:
         """根据搜索参数生成友好的提示消息"""
-        title = kwargs.get("title", "")
-        year = kwargs.get("year")
+        tmdb_id = kwargs.get("tmdb_id")
+        douban_id = kwargs.get("douban_id")
         media_type = kwargs.get("media_type")
-        season = kwargs.get("season")
-        filter_pattern = kwargs.get("filter_pattern")
-        
-        message = f"正在搜索种子: {title}"
-        if year:
-            message += f" ({year})"
+
+        if tmdb_id:
+            message = f"正在搜索种子: TMDB={tmdb_id}"
+        elif douban_id:
+            message = f"正在搜索种子: 豆瓣={douban_id}"
+        else:
+            message = "正在搜索种子"
         if media_type:
             message += f" [{media_type}]"
-        if season:
-            message += f" 第{season}季"
-        if filter_pattern:
-            message += f" 过滤: {filter_pattern}"
-        
         return message
 
-    async def run(self, title: str, year: Optional[str] = None,
-                  media_type: Optional[str] = None, season: Optional[int] = None,
-                  sites: Optional[List[int]] = None, filter_pattern: Optional[str] = None, **kwargs) -> str:
+    async def run(self, tmdb_id: Optional[int] = None, douban_id: Optional[str] = None,
+                  media_type: Optional[str] = None, area: Optional[str] = None,
+                  sites: Optional[List[int]] = None, **kwargs) -> str:
         logger.info(
-            f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}, filter_pattern={filter_pattern}")
+            f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}, douban_id={douban_id}, media_type={media_type}, area={area}, sites={sites}")
+
+        if not tmdb_id and not douban_id:
+            return "参数错误：tmdb_id 和 douban_id 至少需要提供一个，请先使用 search_media 工具获取媒体 ID。"
 
         try:
             search_chain = SearchChain()
-            torrents = await search_chain.async_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:
-                    continue
-                if media_type and torrent.meta_info and torrent.meta_info.type:
-                    if torrent.meta_info.type != MediaType(media_type):
-                        continue
-                if season is not None 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)
+            media_type_enum = None
+            if media_type:
+                media_type_key = media_type.strip().lower()
+                if media_type_key == "movie":
+                    media_type_enum = MediaType.MOVIE
+                elif media_type_key == "tv":
+                    media_type_enum = MediaType.TV
+                else:
+                    return f"错误：无效的媒体类型 '{media_type}'，支持的类型：'movie', 'tv'"
+
+            filtered_torrents = await search_chain.async_search_by_id(
+                tmdbid=tmdb_id,
+                doubanid=douban_id,
+                mtype=media_type_enum,
+                area=area or "title",
+                sites=sites,
+                cache_local=False,
+            )
+
+            # 获取站点信息
+            all_indexers = await SitesHelper().async_get_indexers()
+            all_sites = [{"id": indexer.get("id"), "name": indexer.get("name")} for indexer in (all_indexers or [])]
+
+            if sites:
+                search_site_ids = sites
+            else:
+                configured_sites = SystemConfigOper().get(SystemConfigKey.IndexerSites)
+                search_site_ids = configured_sites if configured_sites else []
 
             if filtered_torrents:
                 await search_chain.async_save_cache(filtered_torrents, SEARCH_RESULT_CACHE_FILE)
                 result_json = json.dumps({
                     "total_count": len(filtered_torrents),
                     "message": "搜索完成。请使用 get_search_results 工具获取搜索结果。",
+                    "all_sites": all_sites,
+                    "search_site_ids": search_site_ids,
                     "filter_options": build_filter_options(filtered_torrents),
                 }, ensure_ascii=False, indent=2)
                 return result_json
             else:
-                return f"未找到相关种子资源: {title}"
+                media_id = f"TMDB={tmdb_id}" if tmdb_id else f"豆瓣={douban_id}"
+                result_json = json.dumps({
+                    "message": f"未找到相关种子资源: {media_id}",
+                    "all_sites": all_sites,
+                    "search_site_ids": search_site_ids,
+                }, ensure_ascii=False, indent=2)
+                return result_json
         except Exception as e:
             error_message = f"搜索种子时发生错误: {str(e)}"
             logger.error(f"搜索种子失败: {e}", exc_info=True)