feat(agent): 为需要管理员权限的工具添加 require_admin 字段

- ExecuteCommandTool: 执行命令行
- DeleteDownloadHistoryTool: 删除下载历史
- EditFileTool: 编辑文件
- WriteFileTool: 写入文件
- TransferFileTool: 传输文件
- UpdateSiteTool: 更新站点
- UpdateSiteCookieTool: 更新站点Cookie
- UpdateSubscribeTool: 更新订阅
- DeleteSubscribeTool: 删除订阅
- DeleteDownloadTool: 删除下载
- ModifyDownloadTool: 修改下载
- RunSchedulerTool: 运行定时任务
- RunWorkflowTool: 运行工作流
- RunPluginCommandTool: 运行插件命令
- SendMessageTool: 发送消息
This commit is contained in:
jxxghp
2026-03-29 10:46:35 +08:00
parent ca9cbc1160
commit 0cab21b83c
16 changed files with 541 additions and 291 deletions
+10 -4
View File
@@ -30,11 +30,13 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
_source: Optional[str] = PrivateAttr(default=None) _source: Optional[str] = PrivateAttr(default=None)
_username: Optional[str] = PrivateAttr(default=None) _username: Optional[str] = PrivateAttr(default=None)
_stream_handler: Optional[StreamingHandler] = PrivateAttr(default=None) _stream_handler: Optional[StreamingHandler] = PrivateAttr(default=None)
_require_admin: bool = PrivateAttr(default=False)
def __init__(self, session_id: str, user_id: str, **kwargs): def __init__(self, session_id: str, user_id: str, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
self._session_id = session_id self._session_id = session_id
self._user_id = user_id self._user_id = user_id
self._require_admin = getattr(self.__class__, "require_admin", False)
def _run(self, *args: Any, **kwargs: Any) -> Any: def _run(self, *args: Any, **kwargs: Any) -> Any:
raise NotImplementedError("MoviePilotTool 只支持异步调用,请使用 _arun") raise NotImplementedError("MoviePilotTool 只支持异步调用,请使用 _arun")
@@ -143,11 +145,15 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
async def _check_permission(self) -> Optional[str]: async def _check_permission(self) -> Optional[str]:
""" """
检查用户权限: 检查用户权限:
1. 首先检查用户是否是渠道管理员 1. 首先检查工具是否需要管理员权限
2. 如果渠道没有设置管理员名单,则检查用户是否是系统管理员 2. 如果需要管理员权限,则检查用户是否是渠道管理员
3. 如果都不是系统管理员,检查用户ID是否等于渠道配置的用户ID 3. 如果渠道没有设置管理员名单,则检查用户是否是系统管理员
4. 如果都不是,返回权限拒绝消息 4. 如果都不是系统管理员,检查用户ID是否等于渠道配置的用户ID
5. 如果都不是,返回权限拒绝消息
""" """
if not self._require_admin:
return None
if not self._channel or not self._source: if not self._channel or not self._source:
return None return None
+31 -10
View File
@@ -11,16 +11,29 @@ from app.log import logger
class DeleteDownloadInput(BaseModel): class DeleteDownloadInput(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)") explanation: str = Field(
downloader: Optional[str] = Field(None, description="Name of specific downloader (optional, if not provided will search all downloaders)") ...,
delete_files: Optional[bool] = Field(False, description="Whether to delete downloaded files along with the task (default: False, only removes the task from downloader)") 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)"
)
downloader: Optional[str] = Field(
None,
description="Name of specific downloader (optional, if not provided will search all downloaders)",
)
delete_files: Optional[bool] = Field(
False,
description="Whether to delete downloaded files along with the task (default: False, only removes the task from downloader)",
)
class DeleteDownloadTool(MoviePilotTool): class DeleteDownloadTool(MoviePilotTool):
name: str = "delete_download" name: str = "delete_download"
description: str = "Delete a download task from the downloader by task hash only. Optionally specify the downloader name and whether to delete downloaded files." description: str = "Delete a download task from the downloader by task hash only. Optionally specify the downloader name and whether to delete downloaded files."
args_schema: Type[BaseModel] = DeleteDownloadInput args_schema: Type[BaseModel] = DeleteDownloadInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据删除参数生成友好的提示消息""" """根据删除参数生成友好的提示消息"""
@@ -36,20 +49,29 @@ class DeleteDownloadTool(MoviePilotTool):
return message return message
async def run(self, hash: str, downloader: Optional[str] = None, async def run(
delete_files: Optional[bool] = False, **kwargs) -> str: self,
logger.info(f"执行工具: {self.name}, 参数: hash={hash}, downloader={downloader}, delete_files={delete_files}") hash: str,
downloader: Optional[str] = None,
delete_files: Optional[bool] = False,
**kwargs,
) -> str:
logger.info(
f"执行工具: {self.name}, 参数: hash={hash}, downloader={downloader}, delete_files={delete_files}"
)
try: try:
download_chain = DownloadChain() download_chain = DownloadChain()
# 仅支持通过hash删除任务 # 仅支持通过hash删除任务
if len(hash) != 40 or not all(c in '0123456789abcdefABCDEF' for c in hash): if len(hash) != 40 or not all(c in "0123456789abcdefABCDEF" for c in hash):
return "参数错误:hash 格式无效,请先使用 query_download_tasks 工具获取正确的 hash。" return "参数错误:hash 格式无效,请先使用 query_download_tasks 工具获取正确的 hash。"
# 删除下载任务 # 删除下载任务
# remove_torrents 支持 delete_file 参数,可以控制是否删除文件 # remove_torrents 支持 delete_file 参数,可以控制是否删除文件
result = download_chain.remove_torrents(hashs=[hash], downloader=downloader, delete_file=delete_files) result = download_chain.remove_torrents(
hashs=[hash], downloader=downloader, delete_file=delete_files
)
if result: if result:
files_info = "(包含文件)" if delete_files else "(不包含文件)" files_info = "(包含文件)" if delete_files else "(不包含文件)"
@@ -59,4 +81,3 @@ class DeleteDownloadTool(MoviePilotTool):
except Exception as e: except Exception as e:
logger.error(f"删除下载任务失败: {e}", exc_info=True) logger.error(f"删除下载任务失败: {e}", exc_info=True)
return f"删除下载任务时发生错误: {str(e)}" return f"删除下载任务时发生错误: {str(e)}"
@@ -26,6 +26,7 @@ class DeleteDownloadHistoryTool(MoviePilotTool):
name: str = "delete_download_history" name: str = "delete_download_history"
description: str = "Delete a download history record by ID. This only removes the record from the database, does not delete any actual files." description: str = "Delete a download history record by ID. This only removes the record from the database, does not delete any actual files."
args_schema: Type[BaseModel] = DeleteDownloadHistoryInput args_schema: Type[BaseModel] = DeleteDownloadHistoryInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
history_id = kwargs.get("history_id") history_id = kwargs.get("history_id")
+17 -11
View File
@@ -14,14 +14,22 @@ from app.schemas.types import EventType
class DeleteSubscribeInput(BaseModel): class DeleteSubscribeInput(BaseModel):
"""删除订阅工具的输入参数模型""" """删除订阅工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
subscribe_id: int = Field(..., description="The ID of the subscription to delete (can be obtained from query_subscribes tool)") explanation: str = Field(
...,
description="Clear explanation of why this tool is being used in the current context",
)
subscribe_id: int = Field(
...,
description="The ID of the subscription to delete (can be obtained from query_subscribes tool)",
)
class DeleteSubscribeTool(MoviePilotTool): class DeleteSubscribeTool(MoviePilotTool):
name: str = "delete_subscribe" name: str = "delete_subscribe"
description: str = "Delete a media subscription by its ID. This will remove the subscription and stop automatic downloads for that media." description: str = "Delete a media subscription by its ID. This will remove the subscription and stop automatic downloads for that media."
args_schema: Type[BaseModel] = DeleteSubscribeInput args_schema: Type[BaseModel] = DeleteSubscribeInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据删除参数生成友好的提示消息""" """根据删除参数生成友好的提示消息"""
@@ -45,19 +53,17 @@ class DeleteSubscribeTool(MoviePilotTool):
subscribe_oper.delete(subscribe_id) subscribe_oper.delete(subscribe_id)
# 发送事件 # 发送事件
await eventmanager.async_send_event(EventType.SubscribeDeleted, { await eventmanager.async_send_event(
"subscribe_id": subscribe_id, EventType.SubscribeDeleted,
"subscribe_info": subscribe_info {"subscribe_id": subscribe_id, "subscribe_info": subscribe_info},
}) )
# 统计订阅 # 统计订阅
SubscribeHelper().sub_done_async({ SubscribeHelper().sub_done_async(
"tmdbid": subscribe.tmdbid, {"tmdbid": subscribe.tmdbid, "doubanid": subscribe.doubanid}
"doubanid": subscribe.doubanid )
})
return f"成功删除订阅:{subscribe.name} ({subscribe.year})" return f"成功删除订阅:{subscribe.name} ({subscribe.year})"
except Exception as e: except Exception as e:
logger.error(f"删除订阅失败: {e}", exc_info=True) logger.error(f"删除订阅失败: {e}", exc_info=True)
return f"删除订阅时发生错误: {str(e)}" return f"删除订阅时发生错误: {str(e)}"
+2 -3
View File
@@ -12,6 +12,7 @@ from app.log import logger
class EditFileInput(BaseModel): class EditFileInput(BaseModel):
"""Input parameters for edit file tool""" """Input parameters for edit file tool"""
file_path: str = Field(..., description="The absolute path of the file to edit") file_path: str = Field(..., description="The absolute path of the file to edit")
old_text: str = Field(..., description="The exact old text to be replaced") old_text: str = Field(..., description="The exact old text to be replaced")
new_text: str = Field(..., description="The new text to replace with") new_text: str = Field(..., description="The new text to replace with")
@@ -21,6 +22,7 @@ class EditFileTool(MoviePilotTool):
name: str = "edit_file" name: str = "edit_file"
description: str = "Edit a file by replacing specific old text with new text. Useful for modifying configuration files, code, or scripts." description: str = "Edit a file by replacing specific old text with new text. Useful for modifying configuration files, code, or scripts."
args_schema: Type[BaseModel] = EditFileInput args_schema: Type[BaseModel] = EditFileInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据参数生成友好的提示消息""" """根据参数生成友好的提示消息"""
@@ -63,7 +65,6 @@ class EditFileTool(MoviePilotTool):
logger.info(f"成功编辑文件 {file_path},替换了 {occurrences} 处内容") logger.info(f"成功编辑文件 {file_path},替换了 {occurrences} 处内容")
return f"成功编辑文件 {file_path} (替换了 {occurrences} 处匹配内容)" return f"成功编辑文件 {file_path} (替换了 {occurrences} 处匹配内容)"
except PermissionError: except PermissionError:
return f"错误:没有访问/修改 {file_path} 的权限" return f"错误:没有访问/修改 {file_path} 的权限"
except UnicodeDecodeError: except UnicodeDecodeError:
@@ -71,5 +72,3 @@ class EditFileTool(MoviePilotTool):
except Exception as e: except Exception as e:
logger.error(f"编辑文件 {file_path} 时发生错误: {str(e)}", exc_info=True) logger.error(f"编辑文件 {file_path} 时发生错误: {str(e)}", exc_info=True)
return f"操作失败: {str(e)}" return f"操作失败: {str(e)}"
+25 -10
View File
@@ -11,15 +11,21 @@ from app.log import logger
class ExecuteCommandInput(BaseModel): class ExecuteCommandInput(BaseModel):
"""执行Shell命令工具的输入参数模型""" """执行Shell命令工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this command is being executed")
explanation: str = Field(
..., description="Clear explanation of why this command is being executed"
)
command: str = Field(..., description="The shell command to execute") command: str = Field(..., description="The shell command to execute")
timeout: Optional[int] = Field(60, description="Max execution time in seconds (default: 60)") timeout: Optional[int] = Field(
60, description="Max execution time in seconds (default: 60)"
)
class ExecuteCommandTool(MoviePilotTool): class ExecuteCommandTool(MoviePilotTool):
name: str = "execute_command" name: str = "execute_command"
description: str = "Safely execute shell commands on the server. Useful for system maintenance, checking status, or running custom scripts. Includes timeout and output limits." description: str = "Safely execute shell commands on the server. Useful for system maintenance, checking status, or running custom scripts. Includes timeout and output limits."
args_schema: Type[BaseModel] = ExecuteCommandInput args_schema: Type[BaseModel] = ExecuteCommandInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据命令生成友好的提示消息""" """根据命令生成友好的提示消息"""
@@ -27,10 +33,19 @@ class ExecuteCommandTool(MoviePilotTool):
return f"正在执行系统命令: {command}" return f"正在执行系统命令: {command}"
async def run(self, command: str, timeout: Optional[int] = 60, **kwargs) -> str: async def run(self, command: str, timeout: Optional[int] = 60, **kwargs) -> str:
logger.info(f"执行工具: {self.name}, 参数: command={command}, timeout={timeout}") logger.info(
f"执行工具: {self.name}, 参数: command={command}, timeout={timeout}"
)
# 简单安全过滤 # 简单安全过滤
forbidden_keywords = ["rm -rf /", ":(){ :|:& };:", "dd if=/dev/zero", "mkfs", "reboot", "shutdown"] forbidden_keywords = [
"rm -rf /",
":(){ :|:& };:",
"dd if=/dev/zero",
"mkfs",
"reboot",
"shutdown",
]
for keyword in forbidden_keywords: for keyword in forbidden_keywords:
if keyword in command: if keyword in command:
return f"错误:命令包含禁止使用的关键字 '{keyword}'" return f"错误:命令包含禁止使用的关键字 '{keyword}'"
@@ -38,18 +53,18 @@ class ExecuteCommandTool(MoviePilotTool):
try: try:
# 执行命令 # 执行命令
process = await asyncio.create_subprocess_shell( process = await asyncio.create_subprocess_shell(
command, command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE
) )
try: try:
# 等待完成,带超时 # 等待完成,带超时
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) stdout, stderr = await asyncio.wait_for(
process.communicate(), timeout=timeout
)
# 处理输出 # 处理输出
stdout_str = stdout.decode('utf-8', errors='replace').strip() stdout_str = stdout.decode("utf-8", errors="replace").strip()
stderr_str = stderr.decode('utf-8', errors='replace').strip() stderr_str = stderr.decode("utf-8", errors="replace").strip()
exit_code = process.returncode exit_code = process.returncode
result = f"命令执行完成 (退出码: {exit_code})" result = f"命令执行完成 (退出码: {exit_code})"
+1
View File
@@ -47,6 +47,7 @@ class ModifyDownloadTool(MoviePilotTool):
"Multiple operations can be performed in a single call." "Multiple operations can be performed in a single call."
) )
args_schema: Type[BaseModel] = ModifyDownloadInput args_schema: Type[BaseModel] = ModifyDownloadInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
hash_value = kwargs.get("hash", "") hash_value = kwargs.get("hash", "")
@@ -37,6 +37,7 @@ class RunPluginCommandTool(MoviePilotTool):
"Note: This tool triggers the command execution but the actual processing happens in the background." "Note: This tool triggers the command execution but the actual processing happens in the background."
) )
args_schema: Type[BaseModel] = RunPluginCommandInput args_schema: Type[BaseModel] = RunPluginCommandInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""生成友好的提示消息""" """生成友好的提示消息"""
+10 -3
View File
@@ -11,14 +11,22 @@ from app.scheduler import Scheduler
class RunSchedulerInput(BaseModel): class RunSchedulerInput(BaseModel):
"""运行定时服务工具的输入参数模型""" """运行定时服务工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
job_id: str = Field(..., description="The ID of the scheduled job to run (can be obtained from query_schedulers tool)") explanation: str = Field(
...,
description="Clear explanation of why this tool is being used in the current context",
)
job_id: str = Field(
...,
description="The ID of the scheduled job to run (can be obtained from query_schedulers tool)",
)
class RunSchedulerTool(MoviePilotTool): class RunSchedulerTool(MoviePilotTool):
name: str = "run_scheduler" name: str = "run_scheduler"
description: str = "Manually trigger a scheduled task to run immediately. This will execute the specified scheduler job by its ID." description: str = "Manually trigger a scheduled task to run immediately. This will execute the specified scheduler job by its ID."
args_schema: Type[BaseModel] = RunSchedulerInput args_schema: Type[BaseModel] = RunSchedulerInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据运行参数生成友好的提示消息""" """根据运行参数生成友好的提示消息"""
@@ -50,4 +58,3 @@ class RunSchedulerTool(MoviePilotTool):
except Exception as e: except Exception as e:
logger.error(f"运行定时服务失败: {e}", exc_info=True) logger.error(f"运行定时服务失败: {e}", exc_info=True)
return f"运行定时服务时发生错误: {str(e)}" return f"运行定时服务时发生错误: {str(e)}"
+22 -8
View File
@@ -13,15 +13,25 @@ from app.log import logger
class RunWorkflowInput(BaseModel): class RunWorkflowInput(BaseModel):
"""执行工作流工具的输入参数模型""" """执行工作流工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
workflow_id: int = Field(..., description="Workflow ID (can be obtained from query_workflows tool)") explanation: str = Field(
from_begin: Optional[bool] = Field(True, description="Whether to run workflow from the beginning (default: True, if False will continue from last executed action)") ...,
description="Clear explanation of why this tool is being used in the current context",
)
workflow_id: int = Field(
..., description="Workflow ID (can be obtained from query_workflows tool)"
)
from_begin: Optional[bool] = Field(
True,
description="Whether to run workflow from the beginning (default: True, if False will continue from last executed action)",
)
class RunWorkflowTool(MoviePilotTool): class RunWorkflowTool(MoviePilotTool):
name: str = "run_workflow" name: str = "run_workflow"
description: str = "Execute a specific workflow manually by workflow ID. Supports running from the beginning or continuing from the last executed action." description: str = "Execute a specific workflow manually by workflow ID. Supports running from the beginning or continuing from the last executed action."
args_schema: Type[BaseModel] = RunWorkflowInput args_schema: Type[BaseModel] = RunWorkflowInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据工作流参数生成友好的提示消息""" """根据工作流参数生成友好的提示消息"""
@@ -36,9 +46,12 @@ class RunWorkflowTool(MoviePilotTool):
return message return message
async def run(self, workflow_id: int, async def run(
from_begin: Optional[bool] = True, **kwargs) -> str: self, workflow_id: int, from_begin: Optional[bool] = True, **kwargs
logger.info(f"执行工具: {self.name}, 参数: workflow_id={workflow_id}, from_begin={from_begin}") ) -> str:
logger.info(
f"执行工具: {self.name}, 参数: workflow_id={workflow_id}, from_begin={from_begin}"
)
try: try:
# 获取数据库会话 # 获取数据库会话
@@ -51,7 +64,9 @@ class RunWorkflowTool(MoviePilotTool):
# 执行工作流 # 执行工作流
workflow_chain = WorkflowChain() workflow_chain = WorkflowChain()
state, errmsg = workflow_chain.process(workflow.id, from_begin=from_begin) state, errmsg = workflow_chain.process(
workflow.id, from_begin=from_begin
)
if not state: if not state:
return f"执行工作流失败:{workflow.name} (ID: {workflow.id})\n错误原因:{errmsg}" return f"执行工作流失败:{workflow.name} (ID: {workflow.id})\n错误原因:{errmsg}"
@@ -60,4 +75,3 @@ class RunWorkflowTool(MoviePilotTool):
except Exception as e: except Exception as e:
logger.error(f"执行工作流失败: {e}", exc_info=True) logger.error(f"执行工作流失败: {e}", exc_info=True)
return f"执行工作流时发生错误: {str(e)}" return f"执行工作流时发生错误: {str(e)}"
+26 -7
View File
@@ -10,23 +10,38 @@ from app.log import logger
class SendMessageInput(BaseModel): class SendMessageInput(BaseModel):
"""发送消息工具的输入参数模型""" """发送消息工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
message: str = Field(..., description="The message content to send to the user (should be clear and informative)") explanation: str = Field(
message_type: Optional[str] = Field("info", ...,
description="Type of message: 'info' for general information, 'success' for successful operations, 'warning' for warnings, 'error' for error messages") description="Clear explanation of why this tool is being used in the current context",
)
message: str = Field(
...,
description="The message content to send to the user (should be clear and informative)",
)
message_type: Optional[str] = Field(
"info",
description="Type of message: 'info' for general information, 'success' for successful operations, 'warning' for warnings, 'error' for error messages",
)
class SendMessageTool(MoviePilotTool): class SendMessageTool(MoviePilotTool):
name: str = "send_message" name: str = "send_message"
description: str = "Send notification message to the user through configured notification channels (Telegram, Slack, WeChat, etc.). Used to inform users about operation results, errors, or important updates." description: str = "Send notification message to the user through configured notification channels (Telegram, Slack, WeChat, etc.). Used to inform users about operation results, errors, or important updates."
args_schema: Type[BaseModel] = SendMessageInput args_schema: Type[BaseModel] = SendMessageInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据消息参数生成友好的提示消息""" """根据消息参数生成友好的提示消息"""
message = kwargs.get("message", "") message = kwargs.get("message", "")
message_type = kwargs.get("message_type", "info") message_type = kwargs.get("message_type", "info")
type_map = {"info": "信息", "success": "成功", "warning": "警告", "error": "错误"} type_map = {
"info": "信息",
"success": "成功",
"warning": "警告",
"error": "错误",
}
type_desc = type_map.get(message_type, message_type) type_desc = type_map.get(message_type, message_type)
# 截断过长的消息 # 截断过长的消息
@@ -35,8 +50,12 @@ class SendMessageTool(MoviePilotTool):
return f"正在发送{type_desc}消息: {message}" return f"正在发送{type_desc}消息: {message}"
async def run(self, message: str, message_type: Optional[str] = None, **kwargs) -> str: async def run(
logger.info(f"执行工具: {self.name}, 参数: message={message}, message_type={message_type}") self, message: str, message_type: Optional[str] = None, **kwargs
) -> str:
logger.info(
f"执行工具: {self.name}, 参数: message={message}, message_type={message_type}"
)
try: try:
await self.send_tool_message(message, title=message_type) await self.send_tool_message(message, title=message_type)
return "消息已发送" return "消息已发送"
+70 -26
View File
@@ -13,23 +13,53 @@ from app.schemas import FileItem, MediaType
class TransferFileInput(BaseModel): class TransferFileInput(BaseModel):
"""整理文件或目录工具的输入参数模型""" """整理文件或目录工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
file_path: str = Field(..., description="Path to the file or directory to transfer (e.g., '/path/to/file.mkv' or '/path/to/directory')") explanation: str = Field(
storage: Optional[str] = Field("local", description="Storage type of the source file (default: 'local', can be 'smb', 'alist', etc.)") ...,
target_path: Optional[str] = Field(None, description="Target path for the transferred file/directory (optional, uses default library path if not specified)") description="Clear explanation of why this tool is being used in the current context",
target_storage: Optional[str] = Field(None, description="Target storage type (optional, uses default storage if not specified)") )
file_path: str = Field(
...,
description="Path to the file or directory to transfer (e.g., '/path/to/file.mkv' or '/path/to/directory')",
)
storage: Optional[str] = Field(
"local",
description="Storage type of the source file (default: 'local', can be 'smb', 'alist', etc.)",
)
target_path: Optional[str] = Field(
None,
description="Target path for the transferred file/directory (optional, uses default library path if not specified)",
)
target_storage: Optional[str] = Field(
None,
description="Target storage type (optional, uses default storage if not specified)",
)
media_type: Optional[str] = Field(None, description="Allowed values: movie, tv") media_type: Optional[str] = Field(None, description="Allowed values: movie, tv")
tmdbid: Optional[int] = Field(None, description="TMDB ID for precise media identification (optional but recommended for accuracy)") tmdbid: Optional[int] = Field(
doubanid: Optional[str] = Field(None, description="Douban ID for media identification (optional)") None,
season: Optional[int] = Field(None, description="Season number for TV shows (optional)") description="TMDB ID for precise media identification (optional but recommended for accuracy)",
transfer_type: Optional[str] = Field(None, description="Transfer mode: 'move' to move files, 'copy' to copy files, 'link' for hard link, 'softlink' for symbolic link (optional, uses default mode if not specified)") )
background: Optional[bool] = Field(False, description="Whether to run transfer in background (default: False, runs synchronously)") doubanid: Optional[str] = Field(
None, description="Douban ID for media identification (optional)"
)
season: Optional[int] = Field(
None, description="Season number for TV shows (optional)"
)
transfer_type: Optional[str] = Field(
None,
description="Transfer mode: 'move' to move files, 'copy' to copy files, 'link' for hard link, 'softlink' for symbolic link (optional, uses default mode if not specified)",
)
background: Optional[bool] = Field(
False,
description="Whether to run transfer in background (default: False, runs synchronously)",
)
class TransferFileTool(MoviePilotTool): class TransferFileTool(MoviePilotTool):
name: str = "transfer_file" name: str = "transfer_file"
description: str = "Transfer/organize a file or directory to the media library. Automatically recognizes media information and organizes files according to configured rules. Supports custom target paths, media identification, and transfer modes." description: str = "Transfer/organize a file or directory to the media library. Automatically recognizes media information and organizes files according to configured rules. Supports custom target paths, media identification, and transfer modes."
args_schema: Type[BaseModel] = TransferFileInput args_schema: Type[BaseModel] = TransferFileInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据整理参数生成友好的提示消息""" """根据整理参数生成友好的提示消息"""
@@ -42,26 +72,37 @@ class TransferFileTool(MoviePilotTool):
if media_type: if media_type:
message += f" [{media_type}]" message += f" [{media_type}]"
if transfer_type: if transfer_type:
transfer_map = {"move": "移动", "copy": "复制", "link": "硬链接", "softlink": "软链接"} transfer_map = {
"move": "移动",
"copy": "复制",
"link": "硬链接",
"softlink": "软链接",
}
message += f" 模式: {transfer_map.get(transfer_type, transfer_type)}" message += f" 模式: {transfer_map.get(transfer_type, transfer_type)}"
if background: if background:
message += " [后台运行]" message += " [后台运行]"
return message return message
async def run(self, file_path: str, storage: Optional[str] = "local", async def run(
target_path: Optional[str] = None, self,
target_storage: Optional[str] = None, file_path: str,
media_type: Optional[str] = None, storage: Optional[str] = "local",
tmdbid: Optional[int] = None, target_path: Optional[str] = None,
doubanid: Optional[str] = None, target_storage: Optional[str] = None,
season: Optional[int] = None, media_type: Optional[str] = None,
transfer_type: Optional[str] = None, tmdbid: Optional[int] = None,
background: Optional[bool] = False, **kwargs) -> str: doubanid: Optional[str] = None,
season: Optional[int] = None,
transfer_type: Optional[str] = None,
background: Optional[bool] = False,
**kwargs,
) -> str:
logger.info( logger.info(
f"执行工具: {self.name}, 参数: file_path={file_path}, storage={storage}, target_path={target_path}, " f"执行工具: {self.name}, 参数: file_path={file_path}, storage={storage}, target_path={target_path}, "
f"target_storage={target_storage}, media_type={media_type}, tmdbid={tmdbid}, doubanid={doubanid}, " f"target_storage={target_storage}, media_type={media_type}, tmdbid={tmdbid}, doubanid={doubanid}, "
f"season={season}, transfer_type={transfer_type}, background={background}") f"season={season}, transfer_type={transfer_type}, background={background}"
)
try: try:
if not file_path: if not file_path:
@@ -70,7 +111,9 @@ class TransferFileTool(MoviePilotTool):
# 规范化路径 # 规范化路径
if storage == "local": if storage == "local":
# 本地路径处理 # 本地路径处理
if not file_path.startswith("/") and not (len(file_path) > 1 and file_path[1] == ":"): if not file_path.startswith("/") and not (
len(file_path) > 1 and file_path[1] == ":"
):
# 相对路径,尝试转换为绝对路径 # 相对路径,尝试转换为绝对路径
file_path = str(Path(file_path).resolve()) file_path = str(Path(file_path).resolve())
else: else:
@@ -82,7 +125,7 @@ class TransferFileTool(MoviePilotTool):
fileitem = FileItem( fileitem = FileItem(
storage=storage or "local", storage=storage or "local",
path=file_path, path=file_path,
type="dir" if file_path.endswith("/") else "file" type="dir" if file_path.endswith("/") else "file",
) )
# 处理目标路径 # 处理目标路径
@@ -108,7 +151,7 @@ class TransferFileTool(MoviePilotTool):
mtype=media_type_enum, mtype=media_type_enum,
season=season, season=season,
transfer_type=transfer_type, transfer_type=transfer_type,
background=background background=background,
) )
if not state: if not state:
@@ -116,7 +159,9 @@ class TransferFileTool(MoviePilotTool):
if isinstance(errormsg, list): if isinstance(errormsg, list):
error_text = f"整理完成,{len(errormsg)} 个文件转移失败" error_text = f"整理完成,{len(errormsg)} 个文件转移失败"
if errormsg: if errormsg:
error_text += f"\n" + "\n".join(str(e) for e in errormsg[:5]) # 只显示前5个错误 error_text += f"\n" + "\n".join(
str(e) for e in errormsg[:5]
) # 只显示前5个错误
if len(errormsg) > 5: if len(errormsg) > 5:
error_text += f"\n... 还有 {len(errormsg) - 5} 个错误" error_text += f"\n... 还有 {len(errormsg) - 5} 个错误"
else: else:
@@ -130,4 +175,3 @@ class TransferFileTool(MoviePilotTool):
except Exception as e: except Exception as e:
logger.error(f"整理文件失败: {e}", exc_info=True) logger.error(f"整理文件失败: {e}", exc_info=True)
return f"整理文件时发生错误: {str(e)}" return f"整理文件时发生错误: {str(e)}"
+82 -50
View File
@@ -16,31 +16,61 @@ from app.utils.string import StringUtils
class UpdateSiteInput(BaseModel): class UpdateSiteInput(BaseModel):
"""更新站点工具的输入参数模型""" """更新站点工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
site_id: int = Field(..., description="The ID of the site to update (can be obtained from query_sites tool)") explanation: str = Field(
...,
description="Clear explanation of why this tool is being used in the current context",
)
site_id: int = Field(
...,
description="The ID of the site to update (can be obtained from query_sites tool)",
)
name: Optional[str] = Field(None, description="Site name (optional)") name: Optional[str] = Field(None, description="Site name (optional)")
url: Optional[str] = Field(None, description="Site URL (optional, will be automatically formatted)") url: Optional[str] = Field(
pri: Optional[int] = Field(None, description="Site priority (optional, smaller value = higher priority, e.g., pri=1 has higher priority than pri=10)") None, description="Site URL (optional, will be automatically formatted)"
)
pri: Optional[int] = Field(
None,
description="Site priority (optional, smaller value = higher priority, e.g., pri=1 has higher priority than pri=10)",
)
rss: Optional[str] = Field(None, description="RSS feed URL (optional)") rss: Optional[str] = Field(None, description="RSS feed URL (optional)")
cookie: Optional[str] = Field(None, description="Site cookie (optional)") cookie: Optional[str] = Field(None, description="Site cookie (optional)")
ua: Optional[str] = Field(None, description="User-Agent string (optional)") ua: Optional[str] = Field(None, description="User-Agent string (optional)")
apikey: Optional[str] = Field(None, description="API key (optional)") apikey: Optional[str] = Field(None, description="API key (optional)")
token: Optional[str] = Field(None, description="API token (optional)") token: Optional[str] = Field(None, description="API token (optional)")
proxy: Optional[int] = Field(None, description="Whether to use proxy: 0 for no, 1 for yes (optional)") proxy: Optional[int] = Field(
filter: Optional[str] = Field(None, description="Filter rule as regular expression (optional)") None, description="Whether to use proxy: 0 for no, 1 for yes (optional)"
)
filter: Optional[str] = Field(
None, description="Filter rule as regular expression (optional)"
)
note: Optional[str] = Field(None, description="Site notes/remarks (optional)") note: Optional[str] = Field(None, description="Site notes/remarks (optional)")
timeout: Optional[int] = Field(None, description="Request timeout in seconds (optional, default: 15)") timeout: Optional[int] = Field(
limit_interval: Optional[int] = Field(None, description="Rate limit interval in seconds (optional)") None, description="Request timeout in seconds (optional, default: 15)"
limit_count: Optional[int] = Field(None, description="Rate limit count per interval (optional)") )
limit_seconds: Optional[int] = Field(None, description="Rate limit seconds between requests (optional)") limit_interval: Optional[int] = Field(
is_active: Optional[bool] = Field(None, description="Whether site is active: True for enabled, False for disabled (optional)") None, description="Rate limit interval in seconds (optional)"
downloader: Optional[str] = Field(None, description="Downloader name for this site (optional)") )
limit_count: Optional[int] = Field(
None, description="Rate limit count per interval (optional)"
)
limit_seconds: Optional[int] = Field(
None, description="Rate limit seconds between requests (optional)"
)
is_active: Optional[bool] = Field(
None,
description="Whether site is active: True for enabled, False for disabled (optional)",
)
downloader: Optional[str] = Field(
None, description="Downloader name for this site (optional)"
)
class UpdateSiteTool(MoviePilotTool): class UpdateSiteTool(MoviePilotTool):
name: str = "update_site" name: str = "update_site"
description: str = "Update site configuration including URL, priority, authentication credentials (cookie, UA, API key), proxy settings, rate limits, and other site properties. Supports updating multiple site attributes at once. Site priority (pri): smaller values have higher priority (e.g., pri=1 has higher priority than pri=10)." description: str = "Update site configuration including URL, priority, authentication credentials (cookie, UA, API key), proxy settings, rate limits, and other site properties. Supports updating multiple site attributes at once. Site priority (pri): smaller values have higher priority (e.g., pri=1 has higher priority than pri=10)."
args_schema: Type[BaseModel] = UpdateSiteInput args_schema: Type[BaseModel] = UpdateSiteInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据更新参数生成友好的提示消息""" """根据更新参数生成友好的提示消息"""
@@ -68,25 +98,28 @@ class UpdateSiteTool(MoviePilotTool):
return f"正在更新站点 #{site_id}: {', '.join(fields_updated)}" return f"正在更新站点 #{site_id}: {', '.join(fields_updated)}"
return f"正在更新站点 #{site_id}" return f"正在更新站点 #{site_id}"
async def run(self, site_id: int, async def run(
name: Optional[str] = None, self,
url: Optional[str] = None, site_id: int,
pri: Optional[int] = None, name: Optional[str] = None,
rss: Optional[str] = None, url: Optional[str] = None,
cookie: Optional[str] = None, pri: Optional[int] = None,
ua: Optional[str] = None, rss: Optional[str] = None,
apikey: Optional[str] = None, cookie: Optional[str] = None,
token: Optional[str] = None, ua: Optional[str] = None,
proxy: Optional[int] = None, apikey: Optional[str] = None,
filter: Optional[str] = None, token: Optional[str] = None,
note: Optional[str] = None, proxy: Optional[int] = None,
timeout: Optional[int] = None, filter: Optional[str] = None,
limit_interval: Optional[int] = None, note: Optional[str] = None,
limit_count: Optional[int] = None, timeout: Optional[int] = None,
limit_seconds: Optional[int] = None, limit_interval: Optional[int] = None,
is_active: Optional[bool] = None, limit_count: Optional[int] = None,
downloader: Optional[str] = None, limit_seconds: Optional[int] = None,
**kwargs) -> str: is_active: Optional[bool] = None,
downloader: Optional[str] = None,
**kwargs,
) -> str:
logger.info(f"执行工具: {self.name}, 参数: site_id={site_id}") logger.info(f"执行工具: {self.name}, 参数: site_id={site_id}")
try: try:
@@ -95,10 +128,10 @@ class UpdateSiteTool(MoviePilotTool):
# 获取站点 # 获取站点
site = await Site.async_get(db, site_id) site = await Site.async_get(db, site_id)
if not site: if not site:
return json.dumps({ return json.dumps(
"success": False, {"success": False, "message": f"站点不存在: {site_id}"},
"message": f"站点不存在: {site_id}" ensure_ascii=False,
}, ensure_ascii=False) )
# 构建更新字典 # 构建更新字典
site_dict = {} site_dict = {}
@@ -153,10 +186,10 @@ class UpdateSiteTool(MoviePilotTool):
# 如果没有要更新的字段 # 如果没有要更新的字段
if not site_dict: if not site_dict:
return json.dumps({ return json.dumps(
"success": False, {"success": False, "message": "没有提供要更新的字段"},
"message": "没有提供要更新的字段" ensure_ascii=False,
}, ensure_ascii=False) )
# 更新站点 # 更新站点
await site.async_update(db, site_dict) await site.async_update(db, site_dict)
@@ -165,16 +198,17 @@ class UpdateSiteTool(MoviePilotTool):
updated_site = await Site.async_get(db, site_id) updated_site = await Site.async_get(db, site_id)
# 发送站点更新事件 # 发送站点更新事件
await eventmanager.async_send_event(EventType.SiteUpdated, { await eventmanager.async_send_event(
"domain": updated_site.domain if updated_site else site.domain EventType.SiteUpdated,
}) {"domain": updated_site.domain if updated_site else site.domain},
)
# 构建返回结果 # 构建返回结果
result = { result = {
"success": True, "success": True,
"message": f"站点 #{site_id} 更新成功", "message": f"站点 #{site_id} 更新成功",
"site_id": site_id, "site_id": site_id,
"updated_fields": list(site_dict.keys()) "updated_fields": list(site_dict.keys()),
} }
if updated_site: if updated_site:
@@ -187,7 +221,7 @@ class UpdateSiteTool(MoviePilotTool):
"is_active": updated_site.is_active, "is_active": updated_site.is_active,
"downloader": updated_site.downloader, "downloader": updated_site.downloader,
"proxy": updated_site.proxy, "proxy": updated_site.proxy,
"timeout": updated_site.timeout "timeout": updated_site.timeout,
} }
return json.dumps(result, ensure_ascii=False, indent=2) return json.dumps(result, ensure_ascii=False, indent=2)
@@ -195,9 +229,7 @@ class UpdateSiteTool(MoviePilotTool):
except Exception as e: except Exception as e:
error_message = f"更新站点失败: {str(e)}" error_message = f"更新站点失败: {str(e)}"
logger.error(f"更新站点失败: {e}", exc_info=True) logger.error(f"更新站点失败: {e}", exc_info=True)
return json.dumps({ return json.dumps(
"success": False, {"success": False, "message": error_message, "site_id": site_id},
"message": error_message, ensure_ascii=False,
"site_id": site_id )
}, ensure_ascii=False)
+26 -8
View File
@@ -12,17 +12,28 @@ from app.log import logger
class UpdateSiteCookieInput(BaseModel): class UpdateSiteCookieInput(BaseModel):
"""更新站点Cookie和UA工具的输入参数模型""" """更新站点Cookie和UA工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
site_identifier: int = Field(..., description="Site ID to update Cookie and User-Agent for (can be obtained from query_sites tool)") explanation: str = Field(
...,
description="Clear explanation of why this tool is being used in the current context",
)
site_identifier: int = Field(
...,
description="Site ID to update Cookie and User-Agent for (can be obtained from query_sites tool)",
)
username: str = Field(..., description="Site login username") username: str = Field(..., description="Site login username")
password: str = Field(..., description="Site login password") password: str = Field(..., description="Site login password")
two_step_code: Optional[str] = Field(None, description="Two-step verification code or secret key (optional, required for sites with 2FA enabled)") two_step_code: Optional[str] = Field(
None,
description="Two-step verification code or secret key (optional, required for sites with 2FA enabled)",
)
class UpdateSiteCookieTool(MoviePilotTool): class UpdateSiteCookieTool(MoviePilotTool):
name: str = "update_site_cookie" name: str = "update_site_cookie"
description: str = "Update site Cookie and User-Agent by logging in with username and password. This tool can automatically obtain and update the site's authentication credentials. Supports two-step verification for sites that require it. Accepts site ID only." description: str = "Update site Cookie and User-Agent by logging in with username and password. This tool can automatically obtain and update the site's authentication credentials. Supports two-step verification for sites that require it. Accepts site ID only."
args_schema: Type[BaseModel] = UpdateSiteCookieInput args_schema: Type[BaseModel] = UpdateSiteCookieInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据更新参数生成友好的提示消息""" """根据更新参数生成友好的提示消息"""
@@ -36,9 +47,17 @@ class UpdateSiteCookieTool(MoviePilotTool):
return message return message
async def run(self, site_identifier: int, username: str, password: str, async def run(
two_step_code: Optional[str] = None, **kwargs) -> str: self,
logger.info(f"执行工具: {self.name}, 参数: site_identifier={site_identifier}, username={username}") site_identifier: int,
username: str,
password: str,
two_step_code: Optional[str] = None,
**kwargs,
) -> str:
logger.info(
f"执行工具: {self.name}, 参数: site_identifier={site_identifier}, username={username}"
)
try: try:
site_oper = SiteOper() site_oper = SiteOper()
@@ -53,7 +72,7 @@ class UpdateSiteCookieTool(MoviePilotTool):
site_info=site, site_info=site,
username=username, username=username,
password=password, password=password,
two_step_code=two_step_code two_step_code=two_step_code,
) )
if status: if status:
@@ -63,4 +82,3 @@ class UpdateSiteCookieTool(MoviePilotTool):
except Exception as e: except Exception as e:
logger.error(f"更新站点Cookie和UA失败: {e}", exc_info=True) logger.error(f"更新站点Cookie和UA失败: {e}", exc_info=True)
return f"更新站点Cookie和UA时发生错误: {str(e)}" return f"更新站点Cookie和UA时发生错误: {str(e)}"
+132 -68
View File
@@ -15,34 +15,81 @@ from app.schemas.types import EventType
class UpdateSubscribeInput(BaseModel): class UpdateSubscribeInput(BaseModel):
"""更新订阅工具的输入参数模型""" """更新订阅工具的输入参数模型"""
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
subscribe_id: int = Field(..., description="The ID of the subscription to update (can be obtained from query_subscribes tool)") explanation: str = Field(
...,
description="Clear explanation of why this tool is being used in the current context",
)
subscribe_id: int = Field(
...,
description="The ID of the subscription to update (can be obtained from query_subscribes tool)",
)
name: Optional[str] = Field(None, description="Subscription name/title (optional)") name: Optional[str] = Field(None, description="Subscription name/title (optional)")
year: Optional[str] = Field(None, description="Release year (optional)") year: Optional[str] = Field(None, description="Release year (optional)")
season: Optional[int] = Field(None, description="Season number for TV shows (optional)") season: Optional[int] = Field(
total_episode: Optional[int] = Field(None, description="Total number of episodes (optional)") None, description="Season number for TV shows (optional)"
lack_episode: Optional[int] = Field(None, description="Number of missing episodes (optional)") )
start_episode: Optional[int] = Field(None, description="Starting episode number (optional)") total_episode: Optional[int] = Field(
quality: Optional[str] = Field(None, description="Quality filter as regular expression (optional, e.g., 'BluRay|WEB-DL|HDTV')") None, description="Total number of episodes (optional)"
resolution: Optional[str] = Field(None, description="Resolution filter as regular expression (optional, e.g., '1080p|720p|2160p')") )
effect: Optional[str] = Field(None, description="Effect filter as regular expression (optional, e.g., 'HDR|DV|SDR')") lack_episode: Optional[int] = Field(
include: Optional[str] = Field(None, description="Include filter as regular expression (optional)") None, description="Number of missing episodes (optional)"
exclude: Optional[str] = Field(None, description="Exclude filter as regular expression (optional)") )
filter: Optional[str] = Field(None, description="Filter rule as regular expression (optional)") start_episode: Optional[int] = Field(
state: Optional[str] = Field(None, description="Subscription state: 'R' for enabled, 'P' for pending, 'S' for paused (optional)") None, description="Starting episode number (optional)"
sites: Optional[List[int]] = Field(None, description="List of site IDs to search from (optional)") )
quality: Optional[str] = Field(
None,
description="Quality filter as regular expression (optional, e.g., 'BluRay|WEB-DL|HDTV')",
)
resolution: Optional[str] = Field(
None,
description="Resolution filter as regular expression (optional, e.g., '1080p|720p|2160p')",
)
effect: Optional[str] = Field(
None,
description="Effect filter as regular expression (optional, e.g., 'HDR|DV|SDR')",
)
include: Optional[str] = Field(
None, description="Include filter as regular expression (optional)"
)
exclude: Optional[str] = Field(
None, description="Exclude filter as regular expression (optional)"
)
filter: Optional[str] = Field(
None, description="Filter rule as regular expression (optional)"
)
state: Optional[str] = Field(
None,
description="Subscription state: 'R' for enabled, 'P' for pending, 'S' for paused (optional)",
)
sites: Optional[List[int]] = Field(
None, description="List of site IDs to search from (optional)"
)
downloader: Optional[str] = Field(None, description="Downloader name (optional)") downloader: Optional[str] = Field(None, description="Downloader name (optional)")
save_path: Optional[str] = Field(None, description="Save path for downloaded files (optional)") save_path: Optional[str] = Field(
best_version: Optional[int] = Field(None, description="Whether to upgrade to best version: 0 for no, 1 for yes (optional)") None, description="Save path for downloaded files (optional)"
custom_words: Optional[str] = Field(None, description="Custom recognition words (optional)") )
media_category: Optional[str] = Field(None, description="Custom media category (optional)") best_version: Optional[int] = Field(
episode_group: Optional[str] = Field(None, description="Episode group ID (optional)") None,
description="Whether to upgrade to best version: 0 for no, 1 for yes (optional)",
)
custom_words: Optional[str] = Field(
None, description="Custom recognition words (optional)"
)
media_category: Optional[str] = Field(
None, description="Custom media category (optional)"
)
episode_group: Optional[str] = Field(
None, description="Episode group ID (optional)"
)
class UpdateSubscribeTool(MoviePilotTool): class UpdateSubscribeTool(MoviePilotTool):
name: str = "update_subscribe" name: str = "update_subscribe"
description: str = "Update subscription properties including filters, episode counts, state, and other settings. Supports updating quality/resolution filters, episode tracking, subscription state, and download configuration." description: str = "Update subscription properties including filters, episode counts, state, and other settings. Supports updating quality/resolution filters, episode tracking, subscription state, and download configuration."
args_schema: Type[BaseModel] = UpdateSubscribeInput args_schema: Type[BaseModel] = UpdateSubscribeInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据更新参数生成友好的提示消息""" """根据更新参数生成友好的提示消息"""
@@ -61,7 +108,9 @@ class UpdateSubscribeTool(MoviePilotTool):
fields_updated.append("分辨率过滤") fields_updated.append("分辨率过滤")
if kwargs.get("state"): if kwargs.get("state"):
state_map = {"R": "启用", "P": "禁用", "S": "暂停"} state_map = {"R": "启用", "P": "禁用", "S": "暂停"}
fields_updated.append(f"状态({state_map.get(kwargs.get('state'), kwargs.get('state'))})") fields_updated.append(
f"状态({state_map.get(kwargs.get('state'), kwargs.get('state'))})"
)
if kwargs.get("sites"): if kwargs.get("sites"):
fields_updated.append("站点") fields_updated.append("站点")
if kwargs.get("downloader"): if kwargs.get("downloader"):
@@ -71,28 +120,31 @@ class UpdateSubscribeTool(MoviePilotTool):
return f"正在更新订阅 #{subscribe_id}: {', '.join(fields_updated)}" return f"正在更新订阅 #{subscribe_id}: {', '.join(fields_updated)}"
return f"正在更新订阅 #{subscribe_id}" return f"正在更新订阅 #{subscribe_id}"
async def run(self, subscribe_id: int, async def run(
name: Optional[str] = None, self,
year: Optional[str] = None, subscribe_id: int,
season: Optional[int] = None, name: Optional[str] = None,
total_episode: Optional[int] = None, year: Optional[str] = None,
lack_episode: Optional[int] = None, season: Optional[int] = None,
start_episode: Optional[int] = None, total_episode: Optional[int] = None,
quality: Optional[str] = None, lack_episode: Optional[int] = None,
resolution: Optional[str] = None, start_episode: Optional[int] = None,
effect: Optional[str] = None, quality: Optional[str] = None,
include: Optional[str] = None, resolution: Optional[str] = None,
exclude: Optional[str] = None, effect: Optional[str] = None,
filter: Optional[str] = None, include: Optional[str] = None,
state: Optional[str] = None, exclude: Optional[str] = None,
sites: Optional[List[int]] = None, filter: Optional[str] = None,
downloader: Optional[str] = None, state: Optional[str] = None,
save_path: Optional[str] = None, sites: Optional[List[int]] = None,
best_version: Optional[int] = None, downloader: Optional[str] = None,
custom_words: Optional[str] = None, save_path: Optional[str] = None,
media_category: Optional[str] = None, best_version: Optional[int] = None,
episode_group: Optional[str] = None, custom_words: Optional[str] = None,
**kwargs) -> str: media_category: Optional[str] = None,
episode_group: Optional[str] = None,
**kwargs,
) -> str:
logger.info(f"执行工具: {self.name}, 参数: subscribe_id={subscribe_id}") logger.info(f"执行工具: {self.name}, 参数: subscribe_id={subscribe_id}")
try: try:
@@ -101,10 +153,10 @@ class UpdateSubscribeTool(MoviePilotTool):
# 获取订阅 # 获取订阅
subscribe = await Subscribe.async_get(db, subscribe_id) subscribe = await Subscribe.async_get(db, subscribe_id)
if not subscribe: if not subscribe:
return json.dumps({ return json.dumps(
"success": False, {"success": False, "message": f"订阅不存在: {subscribe_id}"},
"message": f"订阅不存在: {subscribe_id}" ensure_ascii=False,
}, ensure_ascii=False) )
# 保存旧数据用于事件 # 保存旧数据用于事件
old_subscribe_dict = subscribe.to_dict() old_subscribe_dict = subscribe.to_dict()
@@ -126,7 +178,9 @@ class UpdateSubscribeTool(MoviePilotTool):
# 如果总集数增加,缺失集数也要相应增加 # 如果总集数增加,缺失集数也要相应增加
if total_episode > (subscribe.total_episode or 0): if total_episode > (subscribe.total_episode or 0):
old_lack = subscribe.lack_episode or 0 old_lack = subscribe.lack_episode or 0
subscribe_dict["lack_episode"] = old_lack + (total_episode - (subscribe.total_episode or 0)) subscribe_dict["lack_episode"] = old_lack + (
total_episode - (subscribe.total_episode or 0)
)
# 标记为手动修改过总集数 # 标记为手动修改过总集数
subscribe_dict["manual_total_episode"] = 1 subscribe_dict["manual_total_episode"] = 1
@@ -158,10 +212,13 @@ class UpdateSubscribeTool(MoviePilotTool):
if state is not None: if state is not None:
valid_states = ["R", "P", "S", "N"] valid_states = ["R", "P", "S", "N"]
if state not in valid_states: if state not in valid_states:
return json.dumps({ return json.dumps(
"success": False, {
"message": f"无效的订阅状态: {state},有效状态: {', '.join(valid_states)}" "success": False,
}, ensure_ascii=False) "message": f"无效的订阅状态: {state},有效状态: {', '.join(valid_states)}",
},
ensure_ascii=False,
)
subscribe_dict["state"] = state subscribe_dict["state"] = state
# 下载配置 # 下载配置
@@ -184,10 +241,10 @@ class UpdateSubscribeTool(MoviePilotTool):
# 如果没有要更新的字段 # 如果没有要更新的字段
if not subscribe_dict: if not subscribe_dict:
return json.dumps({ return json.dumps(
"success": False, {"success": False, "message": "没有提供要更新的字段"},
"message": "没有提供要更新的字段" ensure_ascii=False,
}, ensure_ascii=False) )
# 更新订阅 # 更新订阅
await subscribe.async_update(db, subscribe_dict) await subscribe.async_update(db, subscribe_dict)
@@ -196,18 +253,23 @@ class UpdateSubscribeTool(MoviePilotTool):
updated_subscribe = await Subscribe.async_get(db, subscribe_id) updated_subscribe = await Subscribe.async_get(db, subscribe_id)
# 发送订阅调整事件 # 发送订阅调整事件
await eventmanager.async_send_event(EventType.SubscribeModified, { await eventmanager.async_send_event(
"subscribe_id": subscribe_id, EventType.SubscribeModified,
"old_subscribe_info": old_subscribe_dict, {
"subscribe_info": updated_subscribe.to_dict() if updated_subscribe else {}, "subscribe_id": subscribe_id,
}) "old_subscribe_info": old_subscribe_dict,
"subscribe_info": updated_subscribe.to_dict()
if updated_subscribe
else {},
},
)
# 构建返回结果 # 构建返回结果
result = { result = {
"success": True, "success": True,
"message": f"订阅 #{subscribe_id} 更新成功", "message": f"订阅 #{subscribe_id} 更新成功",
"subscribe_id": subscribe_id, "subscribe_id": subscribe_id,
"updated_fields": list(subscribe_dict.keys()) "updated_fields": list(subscribe_dict.keys()),
} }
if updated_subscribe: if updated_subscribe:
@@ -223,7 +285,7 @@ class UpdateSubscribeTool(MoviePilotTool):
"start_episode": updated_subscribe.start_episode, "start_episode": updated_subscribe.start_episode,
"quality": updated_subscribe.quality, "quality": updated_subscribe.quality,
"resolution": updated_subscribe.resolution, "resolution": updated_subscribe.resolution,
"effect": updated_subscribe.effect "effect": updated_subscribe.effect,
} }
return json.dumps(result, ensure_ascii=False, indent=2) return json.dumps(result, ensure_ascii=False, indent=2)
@@ -231,9 +293,11 @@ class UpdateSubscribeTool(MoviePilotTool):
except Exception as e: except Exception as e:
error_message = f"更新订阅失败: {str(e)}" error_message = f"更新订阅失败: {str(e)}"
logger.error(f"更新订阅失败: {e}", exc_info=True) logger.error(f"更新订阅失败: {e}", exc_info=True)
return json.dumps({ return json.dumps(
"success": False, {
"message": error_message, "success": False,
"subscribe_id": subscribe_id "message": error_message,
}, ensure_ascii=False) "subscribe_id": subscribe_id,
},
ensure_ascii=False,
)
+2
View File
@@ -12,6 +12,7 @@ from app.log import logger
class WriteFileInput(BaseModel): class WriteFileInput(BaseModel):
"""Input parameters for write file tool""" """Input parameters for write file tool"""
file_path: str = Field(..., description="The absolute path of the file to write") file_path: str = Field(..., description="The absolute path of the file to write")
content: str = Field(..., description="The content to write into the file") content: str = Field(..., description="The content to write into the file")
@@ -20,6 +21,7 @@ class WriteFileTool(MoviePilotTool):
name: str = "write_file" name: str = "write_file"
description: str = "Write full content to a file. If the file already exists, it will be overwritten. Automatically creates parent directories if they don't exist." description: str = "Write full content to a file. If the file already exists, it will be overwritten. Automatically creates parent directories if they don't exist."
args_schema: Type[BaseModel] = WriteFileInput args_schema: Type[BaseModel] = WriteFileInput
require_admin: bool = True
def get_tool_message(self, **kwargs) -> Optional[str]: def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据参数生成友好的提示消息""" """根据参数生成友好的提示消息"""