Skip to content

Commit 57d4786

Browse files
authored
Merge pull request jxxghp#5332 from PKC278/v2
2 parents 30059ef + df03145 commit 57d4786

File tree

2 files changed

+115
-0
lines changed

2 files changed

+115
-0
lines changed

app/agent/tools/factory.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from app.agent.tools.impl.recognize_media import RecognizeMediaTool
2828
from app.agent.tools.impl.scrape_metadata import ScrapeMetadataTool
2929
from app.agent.tools.impl.query_episode_schedule import QueryEpisodeScheduleTool
30+
from app.agent.tools.impl.query_media_detail import QueryMediaDetailTool
3031
from app.agent.tools.impl.search_torrents import SearchTorrentsTool
3132
from app.agent.tools.impl.search_web import SearchWebTool
3233
from app.agent.tools.impl.send_message import SendMessageTool
@@ -61,6 +62,7 @@ def create_tools(session_id: str, user_id: str,
6162
RecognizeMediaTool,
6263
ScrapeMetadataTool,
6364
QueryEpisodeScheduleTool,
65+
QueryMediaDetailTool,
6466
AddSubscribeTool,
6567
UpdateSubscribeTool,
6668
SearchSubscribeTool,
Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
"""查询媒体详情工具"""
2+
3+
import json
4+
from typing import Optional, Type
5+
6+
from pydantic import BaseModel, Field
7+
8+
from app.agent.tools.base import MoviePilotTool
9+
from app.chain.media import MediaChain
10+
from app.log import logger
11+
from app.schemas import MediaType
12+
13+
14+
class QueryMediaDetailInput(BaseModel):
15+
"""查询媒体详情工具的输入参数模型"""
16+
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
17+
tmdb_id: int = Field(..., description="TMDB ID of the media (movie or TV series)")
18+
19+
20+
class QueryMediaDetailTool(MoviePilotTool):
21+
name: str = "query_media_detail"
22+
description: str = "Query detailed media information from TMDB by ID. Returns core metadata including title, year, overview, status, genres, directors, actors, and season count for TV series."
23+
args_schema: Type[BaseModel] = QueryMediaDetailInput
24+
25+
def get_tool_message(self, **kwargs) -> Optional[str]:
26+
"""根据查询参数生成友好的提示消息"""
27+
tmdb_id = kwargs.get("tmdb_id")
28+
return f"正在查询媒体详情: TMDB ID {tmdb_id}"
29+
30+
async def run(self, tmdb_id: int, **kwargs) -> str:
31+
logger.info(f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}")
32+
33+
try:
34+
media_chain = MediaChain()
35+
36+
# 通过自动媒体类型识别获取媒体信息(电影或电视剧)
37+
mediainfo = await media_chain.async_recognize_media(tmdbid=tmdb_id, mtype=None)
38+
39+
if not mediainfo:
40+
return json.dumps({
41+
"success": False,
42+
"message": f"未找到 TMDB ID {tmdb_id} 的媒体信息"
43+
}, ensure_ascii=False)
44+
45+
# 精简 genres - 只保留名称
46+
genres = [g.get("name") for g in (mediainfo.genres or []) if g.get("name")]
47+
48+
# 精简 directors - 只保留姓名和职位
49+
directors = [
50+
{
51+
"name": d.get("name"),
52+
"job": d.get("job")
53+
}
54+
for d in (mediainfo.directors or [])
55+
if d.get("name")
56+
]
57+
58+
# 精简 actors - 只保留姓名和角色
59+
actors = [
60+
{
61+
"name": a.get("name"),
62+
"character": a.get("character")
63+
}
64+
for a in (mediainfo.actors or [])
65+
if a.get("name")
66+
]
67+
68+
# 构建基础媒体详情信息
69+
result = {
70+
"success": True,
71+
"tmdb_id": tmdb_id,
72+
"type": mediainfo.type.value if mediainfo.type else None,
73+
"title": mediainfo.title,
74+
"year": mediainfo.year,
75+
"overview": mediainfo.overview,
76+
"status": mediainfo.status,
77+
"genres": genres,
78+
"directors": directors,
79+
"actors": actors
80+
}
81+
82+
# 如果是电视剧,添加电视剧特有信息
83+
if mediainfo.type == MediaType.TV:
84+
# 精简 season_info - 只保留基础摘要
85+
season_info = [
86+
{
87+
"season_number": s.get("season_number"),
88+
"name": s.get("name"),
89+
"episode_count": s.get("episode_count"),
90+
"air_date": s.get("air_date")
91+
}
92+
for s in (mediainfo.season_info or [])
93+
if s.get("season_number") is not None
94+
]
95+
96+
result.update({
97+
"number_of_seasons": mediainfo.number_of_seasons,
98+
"number_of_episodes": mediainfo.number_of_episodes,
99+
"first_air_date": mediainfo.first_air_date,
100+
"last_air_date": mediainfo.last_air_date,
101+
"season_info": season_info
102+
})
103+
104+
return json.dumps(result, ensure_ascii=False, indent=2)
105+
106+
except Exception as e:
107+
error_message = f"查询媒体详情失败: {str(e)}"
108+
logger.error(f"查询媒体详情失败: {e}", exc_info=True)
109+
return json.dumps({
110+
"success": False,
111+
"message": error_message,
112+
"tmdb_id": tmdb_id
113+
}, ensure_ascii=False)

0 commit comments

Comments
 (0)