From fac2546a92a2205e7df025f32f3246cbe53c2ed5 Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Thu, 20 Nov 2025 13:02:23 +0800
Subject: [PATCH] Enhance media library query tool with detailed logging and
 improved error handling. Refactor to use MediaServerChain for media existence
 checks and item details retrieval.

---
 app/agent/tools/base.py                     |  6 ++-
 app/agent/tools/impl/query_media_library.py | 52 ++++++++++++++++++---
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/app/agent/tools/base.py b/app/agent/tools/base.py
index 3ba5a178..8e1262aa 100644
--- a/app/agent/tools/base.py
+++ b/app/agent/tools/base.py
@@ -7,6 +7,7 @@ from pydantic import PrivateAttr
 
 from app.agent import StreamingCallbackHandler
 from app.chain import ChainBase
+from app.log import logger
 from app.schemas import Notification
 
 
@@ -49,7 +50,10 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
         if tool_message:
             formatted_message = f"⚙️ => {tool_message}"
             await self.send_tool_message(formatted_message)
-        return await self.run(**kwargs)
+        logger.debug(f'Executing tool {self.name} with args: {kwargs}')
+        result = await self.run(**kwargs)
+        logger.debug(f'Tool {self.name} executed with result: {result}')
+        return result
 
     def get_tool_message(self, **kwargs) -> Optional[str]:
         """
diff --git a/app/agent/tools/impl/query_media_library.py b/app/agent/tools/impl/query_media_library.py
index 23474620..04e34369 100644
--- a/app/agent/tools/impl/query_media_library.py
+++ b/app/agent/tools/impl/query_media_library.py
@@ -6,9 +6,11 @@ from typing import Optional, List, Type
 from pydantic import BaseModel, Field
 
 from app.agent.tools.base import MoviePilotTool
-from app.db.mediaserver_oper import MediaServerOper
+from app.chain.mediaserver import MediaServerChain
+from app.core.context import MediaInfo
 from app.log import logger
 from app.schemas import MediaServerItem
+from app.schemas.types import MediaType
 
 
 class QueryMediaLibraryInput(BaseModel):
@@ -48,11 +50,49 @@ class QueryMediaLibraryTool(MoviePilotTool):
                   title: Optional[str] = None, year: Optional[str] = None, **kwargs) -> str:
         logger.info(f"执行工具: {self.name}, 参数: media_type={media_type}, title={title}")
         try:
-            media_server_oper = MediaServerOper()
-            filtered_medias: List[MediaServerItem] = await media_server_oper.async_exists(title=title, year=year, mtype=media_type)
-            if filtered_medias:
-                return json.dumps([m.to_dict() for m in filtered_medias])
-            return "媒体库中未找到相关媒体"
+            if not title:
+                return "请提供媒体标题进行查询"
+            
+            # 创建 MediaInfo 对象
+            mediainfo = MediaInfo()
+            mediainfo.title = title
+            mediainfo.year = year
+            
+            # 转换媒体类型
+            if media_type == "电影":
+                mediainfo.type = MediaType.MOVIE
+            elif media_type == "电视剧":
+                mediainfo.type = MediaType.TV
+            # media_type == "all" 时不设置类型，让媒体服务器自动判断
+            
+            # 调用媒体服务器接口实时查询
+            media_chain = MediaServerChain()
+            existsinfo = media_chain.media_exists(mediainfo=mediainfo)
+            
+            if not existsinfo:
+                return "媒体库中未找到相关媒体"
+            
+            # 如果找到了，获取详细信息
+            result_items = []
+            if existsinfo.itemid and existsinfo.server:
+                iteminfo = media_chain.iteminfo(server=existsinfo.server, item_id=existsinfo.itemid)
+                if iteminfo:
+                    # 使用 model_dump() 转换为字典格式
+                    item_dict = iteminfo.model_dump(exclude_none=True)
+                    result_items.append(item_dict)
+            
+            if result_items:
+                return json.dumps(result_items, ensure_ascii=False)
+            
+            # 如果找到了但没有详细信息，返回基本信息
+            result_dict = {
+                "type": existsinfo.type.value if existsinfo.type else None,
+                "server": existsinfo.server,
+                "server_type": existsinfo.server_type,
+                "itemid": existsinfo.itemid,
+                "seasons": existsinfo.seasons if existsinfo.seasons else {}
+            }
+            return json.dumps([result_dict], ensure_ascii=False)
         except Exception as e:
             logger.error(f"查询媒体库失败: {e}", exc_info=True)
             return f"查询媒体库时发生错误: {str(e)}"