nas/MoviePilot
1
0
Fork 0
mirror of https://github.com/jxxghp/MoviePilot.git synced 2026-03-20 03:57:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: Improve query download tasks logic and add status filtering

Browse Source 
Co-authored-by: jxxghp <jxxghp@qq.com>
This commit is contained in:
Cursor Agent
2025-12-13 07:01:24 +00:00
parent bda19df87f
commit 07f72b0cdc
1 changed files with 42 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
app/agent/tools/impl/query_download_tasks.py
View File

@@ -1,7 +1,7 @@
"""查询下载工具"""
import json
from typing import Optional, Type
from typing import Optional, Type, List, Union
from pydantic import BaseModel, Field
@@ -9,6 +9,8 @@ from app.agent.tools.base import MoviePilotTool
from app.chain.download import DownloadChain
from app.db.downloadhistory_oper import DownloadHistoryOper
from app.log import logger
from app.schemas import TransferTorrent, DownloadingTorrent
from app.schemas.types import TorrentStatus
class QueryDownloadTasksInput(BaseModel):
@@ -27,6 +29,27 @@ class QueryDownloadTasksTool(MoviePilotTool):
description: str = "Query download status and list download tasks. Can query all active downloads, or search for specific tasks by hash or title. Shows download progress, completion status, and task details from configured downloaders."
args_schema: Type[BaseModel] = QueryDownloadTasksInput
def _get_all_torrents(self, download_chain: DownloadChain, downloader: Optional[str] = None) -> List[Union[TransferTorrent, DownloadingTorrent]]:
"""
查询所有状态的任务（包括下载中和已完成的任务）
"""
all_torrents = []
# 查询正在下载的任务
downloading_torrents = download_chain.list_torrents(
downloader=downloader,
status=TorrentStatus.DOWNLOADING
) or []
all_torrents.extend(downloading_torrents)
# 查询已完成的任务（可转移状态）
transfer_torrents = download_chain.list_torrents(
downloader=downloader,
status=TorrentStatus.TRANSFER
) or []
all_torrents.extend(transfer_torrents)
return all_torrents
def get_tool_message(self, **kwargs) -> Optional[str]:
"""根据查询参数生成友好的提示消息"""
downloader = kwargs.get("downloader")
@@ -84,14 +107,25 @@ class QueryDownloadTasksTool(MoviePilotTool):
elif title:
# 如果提供了title，查询所有任务并搜索匹配的标题
# 查询所有状态的任务
all_torrents = download_chain.list_torrents(downloader=downloader) or []
all_torrents = self._get_all_torrents(download_chain, downloader)
filtered_downloads = []
title_lower = title.lower()
for torrent in all_torrents:
# 检查标题或名称是否匹配
if (title.lower() in (torrent.title or "").lower()) or \
(title.lower() in (torrent.name or "").lower()):
# 获取下载历史信息
history = DownloadHistoryOper().get_by_hash(torrent.hash)
# 获取下载历史信息
history = DownloadHistoryOper().get_by_hash(torrent.hash)
# 检查标题或名称是否匹配（包括下载历史中的标题）
matched = False
# 检查torrent的title和name字段
if (title_lower in (torrent.title or "").lower()) or \
(title_lower in (torrent.name or "").lower()):
matched = True
# 检查下载历史中的标题
if history and history.title:
if title_lower in history.title.lower():
matched = True
if matched:
if history:
torrent.media = {
"tmdbid": history.tmdbid,
@@ -119,7 +153,7 @@ class QueryDownloadTasksTool(MoviePilotTool):
else:
# 其他状态（completed、paused、all），使用list_torrents查询所有任务
# 查询所有状态的任务
all_torrents = download_chain.list_torrents(downloader=downloader) or []
all_torrents = self._get_all_torrents(download_chain, downloader)
filtered_downloads = []
for torrent in all_torrents:
if downloader and torrent.downloader != downloader: