MoviePilot/app/agent/tools/impl/search_torrents.py

"""搜索种子工具"""

import json
import re
from typing import List, Optional, Type

from pydantic import BaseModel, Field, field_validator

from app.agent.tools.base import MoviePilotTool
from app.chain.search import SearchChain
from app.log import logger
from app.schemas.types import MediaType
from app.utils.string import StringUtils


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)")
    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)")

    @field_validator("sites", mode="before")
    @classmethod
    def normalize_sites(cls, value):
        """兼容字符串格式的站点列表（如 "[28]"、"28,30"）"""
        if value is None:
            return value
        if isinstance(value, str):
            value = value.strip()
            if not value:
                return None
            try:
                parsed = json.loads(value)
                if isinstance(parsed, list):
                    return parsed
            except Exception:
                pass
            if "," in value:
                return [v.strip() for v in value.split(",") if v.strip()]
            if value.isdigit():
                return [value]
        return value


class SearchTorrentsTool(MoviePilotTool):
    name: str = "search_torrents"
    description: str = "Search for torrent files across configured indexer sites based on media information. Returns available torrent downloads with details like file size, quality, and download links."
    args_schema: Type[BaseModel] = SearchTorrentsInput

    def get_tool_message(self, **kwargs) -> Optional[str]:
        """根据搜索参数生成友好的提示消息"""
        title = kwargs.get("title", "")
        year = kwargs.get("year")
        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 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:
        logger.info(
            f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}, filter_pattern={filter_pattern}")

        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.media_info:
                    if torrent.media_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)

            if filtered_torrents:
                # 限制最多50条结果
                total_count = len(filtered_torrents)
                limited_torrents = filtered_torrents[:50]
                # 精简字段，只保留关键信息
                simplified_torrents = []
                for t in limited_torrents:
                    simplified = {}
                    # 精简 torrent_info
                    if t.torrent_info:
                        simplified["torrent_info"] = {
                            "title": t.torrent_info.title,
                            "size": StringUtils.format_size(t.torrent_info.size),
                            "seeders": t.torrent_info.seeders,
                            "peers": t.torrent_info.peers,
                            "site_name": t.torrent_info.site_name,
                            "enclosure": t.torrent_info.enclosure,
                            "page_url": t.torrent_info.page_url,
                            "volume_factor": t.torrent_info.volume_factor,
                            "pubdate": t.torrent_info.pubdate
                        }
                    # 精简 media_info
                    if t.media_info:
                        simplified["media_info"] = {
                            "title": t.media_info.title,
                            "en_title": t.media_info.en_title,
                            "year": t.media_info.year,
                            "type": t.media_info.type.value if t.media_info.type else None,
                            "season": t.media_info.season,
                            "tmdb_id": t.media_info.tmdb_id
                        }
                    # 精简 meta_info
                    if t.meta_info:
                        simplified["meta_info"] = {
                            "name": t.meta_info.name,
                            "cn_name": t.meta_info.cn_name,
                            "en_name": t.meta_info.en_name,
                            "year": t.meta_info.year,
                            "type": t.meta_info.type.value if t.meta_info.type else None,
                            "begin_season": t.meta_info.begin_season
                        }
                    simplified_torrents.append(simplified)
                result_json = json.dumps(simplified_torrents, ensure_ascii=False, indent=2)
                # 如果结果被裁剪，添加提示信息
                if total_count > 50:
                    return f"注意：搜索结果共找到 {total_count} 条，为节省上下文空间，仅显示前 50 条结果。\n\n{result_json}"
                return result_json
            else:
                return f"未找到相关种子资源: {title}"
        except Exception as e:
            error_message = f"搜索种子时发生错误: {str(e)}"
            logger.error(f"搜索种子失败: {e}", exc_info=True)
            return error_message