"""添加下载工具""" import re from typing import Optional, Type from pydantic import BaseModel, Field from app.agent.tools.base import MoviePilotTool, ToolChain from app.chain.search import SearchChain 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 from app.utils.crypto import HashUtils class AddDownloadInput(BaseModel): """添加下载工具的输入参数模型""" explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context") site_name: Optional[str] = Field(None, description="Name of the torrent site/source (e.g., 'The Pirate Bay')") torrent_title: Optional[str] = Field( None, description="The display name/title of the torrent (e.g., 'The.Matrix.1999.1080p.BluRay.x264')" ) torrent_url: str = Field( ..., description="Torrent download link, magnet URI, or search result reference returned by search_torrents" ) torrent_description: Optional[str] = Field(None, description="Brief description of the torrent content (optional)") downloader: Optional[str] = Field(None, description="Name of the downloader to use (optional, uses default if not specified)") save_path: Optional[str] = Field(None, description="Directory path where the downloaded files should be saved. Using `:` for remote storage. e.g. rclone:/MP, smb:/server/share/Movies. (optional, uses default path if not specified)") labels: Optional[str] = Field(None, description="Comma-separated list of labels/tags to assign to the download (optional, e.g., 'movie,hd,bluray')") class AddDownloadTool(MoviePilotTool): name: str = "add_download" 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 def get_tool_message(self, **kwargs) -> Optional[str]: """根据下载参数生成友好的提示消息""" torrent_title = kwargs.get("torrent_title", "") torrent_url = kwargs.get("torrent_url") site_name = kwargs.get("site_name", "") downloader = kwargs.get("downloader") message = f"正在添加下载任务: {torrent_title or f'资源 {torrent_url}'}" if site_name: message += f" (来源: {site_name})" if downloader: message += f" [下载器: {downloader}]" return message @staticmethod def _build_torrent_ref(context: Context) -> str: """生成用于校验缓存项的短引用""" if not context or not context.torrent_info: return "" return HashUtils.sha1(context.torrent_info.enclosure or "")[:7] @staticmethod def _is_torrent_ref(torrent_ref: Optional[str]) -> bool: """判断是否为内部搜索结果引用""" if not torrent_ref: return False return bool(re.fullmatch(r"[0-9a-f]{7}:\d+", str(torrent_ref).strip())) @classmethod def _resolve_cached_context(cls, torrent_ref: str) -> Optional[Context]: """从最近一次搜索缓存中解析种子上下文,仅支持 hash:id 格式""" ref = str(torrent_ref).strip() if ":" not in ref: return None try: ref_hash, ref_index = ref.split(":", 1) index = int(ref_index) except (TypeError, ValueError): return None if index < 1: return None results = SearchChain().last_search_results() or [] if index > len(results): return None context = results[index - 1] if not ref_hash or cls._build_torrent_ref(context) != ref_hash: return None return context async def run(self, site_name: Optional[str] = None, torrent_title: Optional[str] = None, torrent_url: Optional[str] = None, torrent_description: Optional[str] = None, downloader: Optional[str] = None, save_path: Optional[str] = None, labels: Optional[str] = None, **kwargs) -> str: logger.info( f"执行工具: {self.name}, 参数: site_name={site_name}, torrent_title={torrent_title}, torrent_url={torrent_url}, downloader={downloader}, save_path={save_path}, labels={labels}") try: cached_context = None if torrent_url: is_torrent_ref = self._is_torrent_ref(torrent_url) cached_context = self._resolve_cached_context(torrent_url) if is_torrent_ref and (not cached_context or not cached_context.torrent_info): return "错误:torrent_url 无效,请重新使用 search_torrents 搜索" if not cached_context or not cached_context.torrent_info: cached_context = None if cached_context and cached_context.torrent_info: cached_torrent = cached_context.torrent_info site_name = site_name or cached_torrent.site_name torrent_title = torrent_title or cached_torrent.title torrent_url = cached_torrent.enclosure torrent_description = torrent_description or cached_torrent.description if not torrent_title: return "错误:必须提供种子标题和下载链接" if not torrent_url: return "错误:必须提供种子标题和下载链接" # 使用DownloadChain添加下载 download_chain = DownloadChain() # 根据站点名称查询站点cookie if not site_name: return "错误:必须提供站点名称,请从搜索资源结果信息中获取" siteinfo = await SiteOper().async_get_by_name(site_name) if not siteinfo: return f"错误:未找到站点信息:{site_name}" # 创建下载上下文 torrent_info = TorrentInfo( title=torrent_title, description=torrent_description, enclosure=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 = cached_context.media_info if cached_context and cached_context.media_info else None if not media_info: media_info = await ToolChain().async_recognize_media(meta=meta_info) if not media_info: return "错误:无法识别媒体信息,无法添加下载任务" context = Context( torrent_info=torrent_info, meta_info=meta_info, media_info=media_info ) did = download_chain.download_single( context=context, downloader=downloader, save_path=save_path, label=labels ) if did: return f"成功添加下载任务:{torrent_title}" else: return "添加下载任务失败" except Exception as e: logger.error(f"添加下载任务失败: {e}", exc_info=True) return f"添加下载任务时发生错误: {str(e)}"