fix(agent): serialize download task media models

This commit is contained in:
jxxghp
2026-08-31 06:52:32 +08:00
parent b63bcbe94b
commit e9812f0e75
2 changed files with 63 additions and 13 deletions

View File

@@ -379,19 +379,22 @@ class QueryDownloadTasksTool(MoviePilotTool):
# 精简 media 字段
media = getattr(d, "media", None)
if media:
media_data = (
media.model_dump() if isinstance(media, BaseModel) else media
)
simplified["media"] = {
"type": media_type_to_agent(media.get("type")),
"title": media.get("title"),
"season": media.get("season"),
"episode": media.get("episode"),
"media_source": media.get("media_source"),
"media_id": media.get("media_id"),
"music_type": media.get("music_type"),
"artists": media.get("artists"),
"album": media.get("album"),
"album_id": media.get("album_id"),
"total_tracks": media.get("total_tracks"),
"track_number": media.get("track_number"),
"type": media_type_to_agent(media_data.get("type")),
"title": media_data.get("title"),
"season": media_data.get("season"),
"episode": media_data.get("episode"),
"media_source": media_data.get("media_source"),
"media_id": media_data.get("media_id"),
"music_type": media_data.get("music_type"),
"artists": media_data.get("artists"),
"album": media_data.get("album"),
"album_id": media_data.get("album_id"),
"total_tracks": media_data.get("total_tracks"),
"track_number": media_data.get("track_number"),
}
simplified_downloads.append(simplified)
result_json = json.dumps(simplified_downloads, ensure_ascii=False, indent=2)

View File

@@ -4,7 +4,8 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from app.agent.tools.impl.query_download_tasks import QueryDownloadTasksTool
from app.schemas.transfer import DownloaderTorrent
from app.schemas.transfer import DownloadTaskMedia, DownloaderTorrent
from app.schemas.types import MediaSource
def _tool() -> QueryDownloadTasksTool:
@@ -109,6 +110,52 @@ def test_run_completed_status_formats_completed_download_tasks():
assert payload[0]["trackers"] == ["https://tracker.example/announce"]
def test_run_formats_pydantic_download_task_media():
"""下载任务媒体为 Pydantic 模型时应正常转换为 Agent 输出字段。"""
torrent = DownloaderTorrent(
downloader="qb",
hash="hash-with-media",
title="Movie With Media",
size=1024,
progress=100,
state="completed",
tags="moviepilot",
media=DownloadTaskMedia(
type="电影",
title="补全标题",
season=1,
episode=[2, 3],
media_source=MediaSource.TMDB,
media_id="123",
),
)
with patch.object(
QueryDownloadTasksTool,
"_query_downloads_sync",
return_value={"downloads": [torrent]},
):
result = asyncio.run(
QueryDownloadTasksTool(session_id="session-1", user_id="10001").run()
)
payload = json.loads(result)
assert payload[0]["media"] == {
"type": "movie",
"title": "补全标题",
"season": 1,
"episode": [2, 3],
"media_source": MediaSource.TMDB,
"media_id": "123",
"music_type": None,
"artists": [],
"album": None,
"album_id": None,
"total_tracks": None,
"track_number": None,
}
def test_hash_query_loads_trackers_for_matching_task():
"""
按 Hash 查询详情时应额外加载下载器支持的 Tracker 列表。