"""添加下载工具""" import re from pathlib import Path from typing import List, 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.config import settings from app.core.context import Context from app.core.metainfo import MetaInfo from app.db.site_oper import SiteOper from app.helper.directory import DirectoryHelper from app.log import logger from app.schemas import TorrentInfo, FileURI 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") torrent_url: List[str] = Field( ..., description="One or more torrent_url values. Supports refs from get_search_results (`hash:id`) and magnet links." ) 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 tasks using refs from get_search_results or magnet links." args_schema: Type[BaseModel] = AddDownloadInput def get_tool_message(self, **kwargs) -> Optional[str]: """根据下载参数生成友好的提示消息""" torrent_urls = self._normalize_torrent_urls(kwargs.get("torrent_url")) downloader = kwargs.get("downloader") if torrent_urls: if len(torrent_urls) == 1: if self._is_torrent_ref(torrent_urls[0]): message = f"正在添加下载任务: 资源 {torrent_urls[0]}" else: message = "正在添加下载任务: 磁力链接" else: message = f"正在批量添加下载任务: 共 {len(torrent_urls)} 个资源" else: message = "正在添加下载任务" 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())) @staticmethod def _is_magnet_link_input(torrent_url: Optional[str]) -> bool: """判断输入是否为允许直接添加的磁力链接""" if not torrent_url: return False value = str(torrent_url).strip() return value.startswith("magnet:") @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 @staticmethod def _merge_labels_with_system_tag(labels: Optional[str]) -> Optional[str]: """合并用户标签与系统默认标签，确保任务可被系统管理""" system_tag = (settings.TORRENT_TAG or "").strip() user_labels = [item.strip() for item in (labels or "").split(",") if item.strip()] if system_tag and system_tag not in user_labels: user_labels.append(system_tag) return ",".join(user_labels) if user_labels else None @staticmethod def _format_failed_result(failed_messages: List[str]) -> str: """统一格式化失败结果""" return ", ".join([message for message in failed_messages if message]) @staticmethod def _build_failure_message(torrent_ref: str, error_msg: Optional[str] = None) -> str: """构造失败提示""" normalized_error = (error_msg or "").strip() prefix = "添加种子任务失败：" if normalized_error.startswith(prefix): normalized_error = normalized_error[len(prefix):].lstrip() if AddDownloadTool._is_magnet_link_input(normalized_error): normalized_error = "" if normalized_error: return f"{torrent_ref} {normalized_error}" if AddDownloadTool._is_torrent_ref(torrent_ref): return torrent_ref return "" @classmethod def _normalize_torrent_urls(cls, torrent_url: Optional[List[str] | str]) -> List[str]: """统一规范 torrent_url 输入，保留所有非空值""" if torrent_url is None: return [] if isinstance(torrent_url, str): candidates = torrent_url.split(",") else: candidates = torrent_url return [str(item).strip() for item in candidates if item and str(item).strip()] @staticmethod def _resolve_direct_download_dir(save_path: Optional[str]) -> Optional[Path]: """解析直接下载使用的目录，优先使用 save_path，其次使用默认下载目录""" if save_path: return Path(save_path) download_dirs = DirectoryHelper().get_download_dirs() if not download_dirs: return None dir_conf = download_dirs[0] if not dir_conf.download_path: return None return Path(FileURI(storage=dir_conf.storage or "local", path=dir_conf.download_path).uri) async def run(self, torrent_url: Optional[List[str]] = None, downloader: Optional[str] = None, save_path: Optional[str] = None, labels: Optional[str] = None, **kwargs) -> str: logger.info( f"执行工具: {self.name}, 参数: torrent_url={torrent_url}, downloader={downloader}, save_path={save_path}, labels={labels}") try: torrent_inputs = self._normalize_torrent_urls(torrent_url) if not torrent_inputs: return "错误：torrent_url 不能为空。" download_chain = DownloadChain() merged_labels = self._merge_labels_with_system_tag(labels) success_count = 0 failed_messages = [] for torrent_input in torrent_inputs: if self._is_torrent_ref(torrent_input): cached_context = self._resolve_cached_context(torrent_input) if not cached_context or not cached_context.torrent_info: failed_messages.append(f"{torrent_input} 引用无效，请重新使用 get_search_results 查看搜索结果") continue cached_torrent = cached_context.torrent_info site_name = cached_torrent.site_name torrent_title = cached_torrent.title or torrent_input torrent_description = cached_torrent.description enclosure = cached_torrent.enclosure if not site_name: failed_messages.append(f"{torrent_input} 缺少站点名称") continue siteinfo = await SiteOper().async_get_by_name(site_name) if not siteinfo: failed_messages.append(f"{torrent_input} 未找到站点信息 {site_name}") continue torrent_info = TorrentInfo( title=torrent_title, description=torrent_description, enclosure=enclosure, 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.media_info else None if not media_info: media_info = await ToolChain().async_recognize_media(meta=meta_info) if not media_info: failed_messages.append(f"{torrent_input} 无法识别媒体信息") continue context = Context( torrent_info=torrent_info, meta_info=meta_info, media_info=media_info ) else: if not self._is_magnet_link_input(torrent_input): failed_messages.append( f"{torrent_input} 不是有效的下载内容，非 hash:id 时仅支持 magnet: 开头" ) continue download_dir = self._resolve_direct_download_dir(save_path) if not download_dir: failed_messages.append(f"{torrent_input} 缺少保存路径，且系统未配置可用下载目录") continue result = download_chain.download( content=torrent_input, download_dir=download_dir, cookie=None, label=merged_labels, downloader=downloader ) if result: _, did, _, error_msg = result else: did, error_msg = None, "未找到下载器" if did: success_count += 1 else: failed_messages.append(self._build_failure_message(torrent_input, error_msg)) continue did, error_msg = download_chain.download_single( context=context, downloader=downloader, save_path=save_path, label=merged_labels, return_detail=True ) if did: success_count += 1 else: failed_messages.append(self._build_failure_message(torrent_input, error_msg)) if success_count and not failed_messages: return "任务添加成功" if success_count: return f"部分任务添加失败：{self._format_failed_result(failed_messages)}" return f"任务添加失败：{self._format_failed_result(failed_messages)}" except Exception as e: logger.error(f"添加下载任务失败: {e}", exc_info=True) return f"添加下载任务时发生错误: {str(e)}"