add ai agent

app/agent/tools/impl/add_download.py

@@ -0,0 +1,58 @@
+"""添加下载工具"""
+
+from typing import Optional
+
+from app.chain.download import DownloadChain
+from app.core.context import Context
+from app.core.metainfo import MetaInfo
+from app.log import logger
+from app.schemas import TorrentInfo
+from app.agent.tools.base import MoviePilotTool
+
+
+class AddDownloadTool(MoviePilotTool):
+    name: str = "add_download"
+    description: str = "添加下载任务，将搜索到的种子资源添加到下载器。"
+
+    async def _arun(self, torrent_title: str, torrent_url: str, explanation: str, 
+                    downloader: Optional[str] = None, save_path: Optional[str] = None, 
+                    labels: Optional[str] = None) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: torrent_title={torrent_title}, torrent_url={torrent_url}, downloader={downloader}, save_path={save_path}, labels={labels}")
+        
+        # 发送工具执行说明
+        self._send_tool_message(f"正在添加下载任务: {torrent_title}", "info")
+        
+        try:
+            if not torrent_title or not torrent_url:
+                error_message = "错误：必须提供种子标题和下载链接"
+                self._send_tool_message(error_message, "error")
+                return error_message
+
+            # 使用DownloadChain添加下载
+            download_chain = DownloadChain()
+            
+            # 创建下载上下文
+            torrent_info = TorrentInfo(
+                title=torrent_title,
+                download_url=torrent_url
+            )
+            meta_info = MetaInfo(title=torrent_title)
+            context = Context(
+                torrent_info=torrent_info,
+                meta_info=meta_info
+            )
+
+            did = download_chain.download_single(context=context, downloader=downloader, 
+                                               save_path=save_path, label=labels)
+            if did:
+                success_message = f"成功添加下载任务：{torrent_title}"
+                self._send_tool_message(success_message, "success")
+                return success_message
+            else:
+                error_message = "添加下载任务失败"
+                self._send_tool_message(error_message, "error")
+                return error_message
+        except Exception as e:
+            error_message = f"添加下载任务时发生错误: {str(e)}"
+            self._send_tool_message(error_message, "error")
+            return error_message

app/agent/tools/impl/add_subscribe.py

@@ -0,0 +1,37 @@
+"""添加订阅工具"""
+
+from typing import Optional
+
+from app.chain.subscribe import SubscribeChain
+from app.log import logger
+from app.schemas.types import MediaType
+from app.agent.tools.base import MoviePilotTool
+
+
+class AddSubscribeTool(MoviePilotTool):
+    name: str = "add_subscribe"
+    description: str = "添加媒体订阅，为用户感兴趣的媒体内容创建订阅规则。"
+
+    async def _arun(self, title: str, year: str, media_type: str, explanation: str, 
+                    season: Optional[int] = None, tmdb_id: Optional[str] = None) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, tmdb_id={tmdb_id}")
+        
+        # 发送工具执行说明
+        self._send_tool_message(f"正在添加订阅: {title} ({year}) - {media_type}", "info")
+        
+        try:
+            subscribe_chain = SubscribeChain()
+            sid, message = subscribe_chain.add(mtype=MediaType(media_type), title=title, year=year, 
+                                             tmdbid=tmdb_id, season=season, username=self._user_id)
+            if sid:
+                success_message = f"成功添加订阅：{title} ({year})"
+                self._send_tool_message(success_message, "success")
+                return success_message
+            else:
+                error_message = f"添加订阅失败：{message}"
+                self._send_tool_message(error_message, "error")
+                return error_message
+        except Exception as e:
+            error_message = f"添加订阅时发生错误: {str(e)}"
+            self._send_tool_message(error_message, "error")
+            return error_message

app/agent/tools/impl/get_recommendations.py

@@ -0,0 +1,39 @@
+"""获取推荐工具"""
+
+import json
+from typing import Optional
+
+from app.chain.recommend import RecommendChain
+from app.log import logger
+from app.agent.tools.base import MoviePilotTool
+
+
+class GetRecommendationsTool(MoviePilotTool):
+    name: str = "get_recommendations"
+    description: str = "获取热门媒体推荐，包括电影、电视剧等热门内容。"
+
+    async def _arun(self, explanation: str, source: Optional[str] = "tmdb_trending", 
+                    media_type: Optional[str] = "all", limit: Optional[int] = 20) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: source={source}, media_type={media_type}, limit={limit}")
+        try:
+            recommend_chain = RecommendChain()
+            results = []
+            if source == "tmdb_trending":
+                results = recommend_chain.tmdb_trending(limit=limit)
+            elif source == "douban_hot":
+                if media_type == "movie":
+                    results = recommend_chain.douban_movie_hot(limit=limit)
+                elif media_type == "tv":
+                    results = recommend_chain.douban_tv_hot(limit=limit)
+                else: # all
+                    results.extend(recommend_chain.douban_movie_hot(limit=limit))
+                    results.extend(recommend_chain.douban_tv_hot(limit=limit))
+            elif source == "bangumi_calendar":
+                results = recommend_chain.bangumi_calendar(limit=limit)
+            
+            if results:
+                return json.dumps([r.dict() for r in results], ensure_ascii=False, indent=2)
+            return "未找到推荐内容。"
+        except Exception as e:
+            logger.error(f"获取推荐失败: {e}")
+            return f"获取推荐时发生错误: {str(e)}"

app/agent/tools/impl/query_downloaders.py

@@ -0,0 +1,25 @@
+"""查询下载器工具"""
+
+import json
+
+from app.db.systemconfig_oper import SystemConfigOper
+from app.log import logger
+from app.schemas.types import SystemConfigKey
+from app.agent.tools.base import MoviePilotTool
+
+
+class QueryDownloadersTool(MoviePilotTool):
+    name: str = "query_downloaders"
+    description: str = "查询下载器配置，查看可用的下载器列表和配置信息。"
+
+    async def _arun(self, explanation: str) -> str:
+        logger.info(f"执行工具: {self.name}")
+        try:
+            system_config_oper = SystemConfigOper()
+            downloaders_config = system_config_oper.get(SystemConfigKey.Downloaders)
+            if downloaders_config:
+                return json.dumps(downloaders_config, ensure_ascii=False, indent=2)
+            return "未配置下载器。"
+        except Exception as e:
+            logger.error(f"查询下载器失败: {e}")
+            return f"查询下载器时发生错误: {str(e)}"

app/agent/tools/impl/query_downloads.py

@@ -0,0 +1,33 @@
+"""查询下载工具"""
+
+import json
+from typing import Optional
+
+from app.db.download_oper import DownloadOper
+from app.log import logger
+from app.agent.tools.base import MoviePilotTool
+
+
+class QueryDownloadsTool(MoviePilotTool):
+    name: str = "query_downloads"
+    description: str = "查询下载状态，查看下载器的任务列表和进度。"
+
+    async def _arun(self, explanation: str, downloader: Optional[str] = None, 
+                    status: Optional[str] = "all") -> str:
+        logger.info(f"执行工具: {self.name}, 参数: downloader={downloader}, status={status}")
+        try:
+            download_oper = DownloadOper()
+            downloads = download_oper.list()
+            filtered_downloads = []
+            for dl in downloads:
+                if downloader and dl.downloader != downloader:
+                    continue
+                if status != "all" and dl.status != status:
+                    continue
+                filtered_downloads.append(dl)
+            if filtered_downloads:
+                return json.dumps([d.dict() for d in filtered_downloads], ensure_ascii=False, indent=2)
+            return "未找到相关下载任务。"
+        except Exception as e:
+            logger.error(f"查询下载失败: {e}")
+            return f"查询下载时发生错误: {str(e)}"

app/agent/tools/impl/query_media_library.py

@@ -0,0 +1,33 @@
+"""查询媒体库工具"""
+
+import json
+from typing import Optional
+
+from app.db.media_oper import MediaOper
+from app.log import logger
+from app.agent.tools.base import MoviePilotTool
+
+
+class QueryMediaLibraryTool(MoviePilotTool):
+    name: str = "query_media_library"
+    description: str = "查询媒体库状态，查看已入库的媒体文件情况。"
+
+    async def _arun(self, explanation: str, media_type: Optional[str] = "all", 
+                    title: Optional[str] = None) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: media_type={media_type}, title={title}")
+        try:
+            media_oper = MediaOper()
+            medias = media_oper.list()
+            filtered_medias = []
+            for media in medias:
+                if media_type != "all" and media.type != media_type:
+                    continue
+                if title and title.lower() not in media.title.lower():
+                    continue
+                filtered_medias.append(media)
+            if filtered_medias:
+                return json.dumps([m.dict() for m in filtered_medias], ensure_ascii=False, indent=2)
+            return "媒体库中未找到相关媒体。"
+        except Exception as e:
+            logger.error(f"查询媒体库失败: {e}")
+            return f"查询媒体库时发生错误: {str(e)}"

app/agent/tools/impl/query_subscribes.py

@@ -0,0 +1,33 @@
+"""查询订阅工具"""
+
+import json
+from typing import Optional
+
+from app.db.subscribe_oper import SubscribeOper
+from app.log import logger
+from app.agent.tools.base import MoviePilotTool
+
+
+class QuerySubscribesTool(MoviePilotTool):
+    name: str = "query_subscribes"
+    description: str = "查询订阅状态，查看用户的订阅列表和状态。"
+
+    async def _arun(self, explanation: str, status: Optional[str] = "all", 
+                    media_type: Optional[str] = "all") -> str:
+        logger.info(f"执行工具: {self.name}, 参数: status={status}, media_type={media_type}")
+        try:
+            subscribe_oper = SubscribeOper()
+            subscribes = subscribe_oper.list()
+            filtered_subscribes = []
+            for sub in subscribes:
+                if status != "all" and sub.status != status:
+                    continue
+                if media_type != "all" and sub.type != media_type:
+                    continue
+                filtered_subscribes.append(sub)
+            if filtered_subscribes:
+                return json.dumps([s.dict() for s in filtered_subscribes], ensure_ascii=False, indent=2)
+            return "未找到相关订阅。"
+        except Exception as e:
+            logger.error(f"查询订阅失败: {e}")
+            return f"查询订阅时发生错误: {str(e)}"

app/agent/tools/impl/search_media.py

@@ -0,0 +1,42 @@
+"""搜索媒体工具"""
+
+import json
+from typing import Optional
+
+from app.chain.media import MediaChain
+from app.log import logger
+from app.agent.tools.base import MoviePilotTool
+
+
+class SearchMediaTool(MoviePilotTool):
+    name: str = "search_media"
+    description: str = "搜索媒体资源，包括电影、电视剧、动漫等。可以根据标题、年份、类型等条件进行搜索。"
+
+    async def _arun(self, title: str, explanation: str, year: Optional[str] = None, 
+                    media_type: Optional[str] = None, season: Optional[int] = None) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}")
+        
+        # 发送工具执行说明
+        self._send_tool_message(f"正在搜索媒体资源: {title}" + (f" ({year})" if year else ""), "info")
+        
+        try:
+            media_chain = MediaChain()
+            results = media_chain.search_media(title=title, year=year, mtype=media_type, season=season)
+            if results:
+                result_message = f"找到 {len(results)} 个相关媒体资源"
+                self._send_tool_message(result_message, "success")
+                
+                # 发送详细结果
+                for i, result in enumerate(results[:5]):  # 只显示前5个结果
+                    media_info = f"{i+1}. {result.title} ({result.year}) - {result.type}"
+                    self._send_tool_message(media_info, "info")
+                
+                return json.dumps([r.dict() for r in results], ensure_ascii=False, indent=2)
+            else:
+                error_message = f"未找到相关媒体资源: {title}"
+                self._send_tool_message(error_message, "warning")
+                return error_message
+        except Exception as e:
+            error_message = f"搜索媒体失败: {str(e)}"
+            self._send_tool_message(error_message, "error")
+            return error_message

app/agent/tools/impl/search_torrents.py

@@ -0,0 +1,54 @@
+"""搜索种子工具"""
+
+import json
+from typing import List, Optional
+
+from app.chain.search import SearchChain
+from app.log import logger
+from app.schemas.types import MediaType
+from app.agent.tools.base import MoviePilotTool
+
+
+class SearchTorrentsTool(MoviePilotTool):
+    name: str = "search_torrents"
+    description: str = "搜索站点种子资源，根据媒体信息搜索可下载的种子文件。"
+
+    async def _arun(self, title: str, explanation: str, year: Optional[str] = None, 
+                    media_type: Optional[str] = None, season: Optional[int] = None, 
+                    sites: Optional[List[int]] = None) -> str:
+        logger.info(f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}")
+        
+        # 发送工具执行说明
+        self._send_tool_message(f"正在搜索种子资源: {title}" + (f" ({year})" if year else ""), "info")
+        
+        try:
+            search_chain = SearchChain()
+            torrents = search_chain.search_by_title(title=title, sites=sites)
+            filtered_torrents = []
+            for torrent in torrents:
+                if year and torrent.meta_info.year != year:
+                    continue
+                if media_type and torrent.media_info and torrent.media_info.type != MediaType(media_type):
+                    continue
+                if season and torrent.meta_info.begin_season != season:
+                    continue
+                filtered_torrents.append(torrent)
+            
+            if filtered_torrents:
+                result_message = f"找到 {len(filtered_torrents)} 个相关种子资源"
+                self._send_tool_message(result_message, "success")
+                
+                # 发送详细结果
+                for i, torrent in enumerate(filtered_torrents[:5]):  # 只显示前5个结果
+                    torrent_info = f"{i+1}. {torrent.title} - {torrent.site_name}"
+                    self._send_tool_message(torrent_info, "info")
+                
+                return json.dumps([t.dict() for t in filtered_torrents], ensure_ascii=False, indent=2)
+            else:
+                error_message = f"未找到相关种子资源: {title}"
+                self._send_tool_message(error_message, "warning")
+                return error_message
+        except Exception as e:
+            error_message = f"搜索种子时发生错误: {str(e)}"
+            self._send_tool_message(error_message, "error")
+            return error_message

app/agent/tools/impl/send_message.py

@@ -0,0 +1,22 @@
+"""发送消息工具"""
+
+from typing import Optional
+
+from app.helper.message import MessageHelper
+from app.log import logger
+from app.agent.tools.base import MoviePilotTool
+
+
+class SendMessageTool(MoviePilotTool):
+    name: str = "send_message"
+    description: str = "发送消息通知，向用户发送操作结果或重要信息。"
+
+    async def _arun(self, message: str, explanation: str, message_type: Optional[str] = "info") -> str:
+        logger.info(f"执行工具: {self.name}, 参数: message={message}, message_type={message_type}")
+        try:
+            message_helper = MessageHelper()
+            message_helper.put(message=message, role="system", title=f"AI助手通知 ({message_type})")
+            return "消息已发送。"
+        except Exception as e:
+            logger.error(f"发送消息失败: {e}")
+            return f"发送消息时发生错误: {str(e)}"