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

"""修改下载任务工具"""

from typing import Optional, Type, List

from pydantic import BaseModel, Field

from app.agent.tools.base import MoviePilotTool
from app.chain.download import DownloadChain
from app.log import logger


class ModifyDownloadInput(BaseModel):
    """修改下载任务工具的输入参数模型"""

    explanation: str = Field(
        ...,
        description="Clear explanation of why this tool is being used in the current context",
    )
    hash: str = Field(
        ..., description="Task hash (can be obtained from query_download_tasks tool)"
    )
    action: Optional[str] = Field(
        None,
        description="Action to perform on the task: 'start' to resume downloading, 'stop' to pause downloading. "
        "If not provided, no start/stop action will be performed.",
    )
    tags: Optional[List[str]] = Field(
        None,
        description="List of tags to set on the download task. If provided, these tags will be added to the task. "
        "Example: ['movie', 'hd']",
    )
    downloader: Optional[str] = Field(
        None,
        description="Name of specific downloader (optional, if not provided will search all downloaders)",
    )


class ModifyDownloadTool(MoviePilotTool):
    """修改下载任务工具"""

    name: str = "modify_download"
    description: str = (
        "Modify a download task in the downloader by task hash. "
        "Supports: 1) Setting tags on a download task, "
        "2) Starting (resuming) a paused download task, "
        "3) Stopping (pausing) a downloading task. "
        "Multiple operations can be performed in a single call."
    )
    args_schema: Type[BaseModel] = ModifyDownloadInput

    def get_tool_message(self, **kwargs) -> Optional[str]:
        hash_value = kwargs.get("hash", "")
        action = kwargs.get("action")
        tags = kwargs.get("tags")
        downloader = kwargs.get("downloader")

        parts = [f"正在修改下载任务: {hash_value}"]
        if action == "start":
            parts.append("操作: 开始下载")
        elif action == "stop":
            parts.append("操作: 暂停下载")
        if tags:
            parts.append(f"标签: {', '.join(tags)}")
        if downloader:
            parts.append(f"下载器: {downloader}")
        return " | ".join(parts)

    async def run(
        self,
        hash: str,
        action: Optional[str] = None,
        tags: Optional[List[str]] = None,
        downloader: Optional[str] = None,
        **kwargs,
    ) -> str:
        logger.info(
            f"执行工具: {self.name}, 参数: hash={hash}, action={action}, tags={tags}, downloader={downloader}"
        )

        try:
            # 校验 hash 格式
            if len(hash) != 40 or not all(c in "0123456789abcdefABCDEF" for c in hash):
                return "参数错误：hash 格式无效，请先使用 query_download_tasks 工具获取正确的 hash。"

            # 校验参数：至少需要一个操作
            if not action and not tags:
                return "参数错误：至少需要指定 action（start/stop）或 tags 中的一个。"

            # 校验 action 参数
            if action and action not in ("start", "stop"):
                return f"参数错误：action 只支持 'start'（开始下载）或 'stop'（暂停下载），收到: '{action}'。"

            download_chain = DownloadChain()
            results = []

            # 设置标签
            if tags:
                tag_result = download_chain.set_torrents_tag(
                    hashs=[hash], tags=tags, downloader=downloader
                )
                if tag_result:
                    results.append(f"成功设置标签：{', '.join(tags)}")
                else:
                    results.append(f"设置标签失败，请检查任务是否存在或下载器是否可用")

            # 执行开始/暂停操作
            if action:
                action_result = download_chain.set_downloading(
                    hash_str=hash, oper=action, name=downloader
                )
                action_desc = "开始" if action == "start" else "暂停"
                if action_result:
                    results.append(f"成功{action_desc}下载任务")
                else:
                    results.append(
                        f"{action_desc}下载任务失败，请检查任务是否存在或下载器是否可用"
                    )

            return f"下载任务 {hash}：" + "；".join(results)

        except Exception as e:
            logger.error(f"修改下载任务失败: {e}", exc_info=True)
            return f"修改下载任务时发生错误: {str(e)}"