From 247208b8a92af7a618899f3ce08aa8883fe887bc Mon Sep 17 00:00:00 2001 From: jxxghp Date: Sat, 1 Nov 2025 11:41:22 +0800 Subject: [PATCH] fix agent --- app/agent/__init__.py | 30 +++++++------- app/agent/tools/Agent Tools v1.0.json | 4 ++ app/agent/tools/base.py | 23 +++++------ app/agent/tools/factory.py | 11 +++--- app/agent/tools/impl/add_download.py | 31 ++++++++------- app/agent/tools/impl/add_subscribe.py | 31 ++++++++------- app/agent/tools/impl/query_media_library.py | 18 +++------ app/agent/tools/impl/search_media.py | 43 ++++++++++----------- app/agent/tools/impl/search_torrents.py | 38 +++++++++--------- app/agent/tools/impl/send_message.py | 6 +-- 10 files changed, 114 insertions(+), 121 deletions(-) diff --git a/app/agent/__init__.py b/app/agent/__init__.py index 82469cb6..9e749906 100644 --- a/app/agent/__init__.py +++ b/app/agent/__init__.py @@ -249,7 +249,7 @@ class MoviePilotAgent: agent_message = await self.callback_handler.get_message() # 发送Agent回复给用户(通过原渠道) - self._send_message_to_channel(agent_message) + self.send_agent_message(agent_message) # 添加Agent回复到记忆 await self.memory_manager.add_memory( @@ -265,22 +265,9 @@ class MoviePilotAgent: error_message = f"处理消息时发生错误: {str(e)}" logger.error(error_message) # 发送错误消息给用户(通过原渠道) - self._send_message_to_channel(error_message) + self.send_agent_message(error_message) return error_message - def _send_message_to_channel(self, message: str, title: str = "MoviePilot助手"): - """通过原渠道发送消息给用户""" - AgentChain().post_message( - Notification( - channel=self.channel, - source=self.source, - userid=self.user_id, - username=self.username, - title=title, - text=message - ) - ) - async def _execute_agent(self, input_context: Dict[str, Any]) -> Dict[str, Any]: """执行LangChain Agent""" try: @@ -314,6 +301,19 @@ class MoviePilotAgent: "token_usage": {} } + def send_agent_message(self, message: str, title: str = "MoviePilot助手"): + """通过原渠道发送消息给用户""" + AgentChain().post_message( + Notification( + channel=self.channel, + source=self.source, + userid=self.user_id, + username=self.username, + title=title, + text=message + ) + ) + async def cleanup(self): """清理智能体资源""" if self.session_id in self.session_store: diff --git a/app/agent/tools/Agent Tools v1.0.json b/app/agent/tools/Agent Tools v1.0.json index e29cf580..6059a3a9 100644 --- a/app/agent/tools/Agent Tools v1.0.json +++ b/app/agent/tools/Agent Tools v1.0.json @@ -257,6 +257,10 @@ "title": { "description": "Specific media title to search for (optional, if provided returns detailed info for that specific media)", "type": "string" + }, + "year": { + "description": "Release year of the media (optional, helps narrow down search results)", + "type": "string" } }, "required": [ diff --git a/app/agent/tools/base.py b/app/agent/tools/base.py index eddaf132..ac3b1dfc 100644 --- a/app/agent/tools/base.py +++ b/app/agent/tools/base.py @@ -4,8 +4,6 @@ from langchain.tools import BaseTool from pydantic import PrivateAttr from app.chain import ChainBase -from app.helper.message import MessageHelper -from app.log import logger from app.schemas import Notification @@ -18,17 +16,14 @@ class MoviePilotTool(BaseTool): _session_id: str = PrivateAttr() _user_id: str = PrivateAttr() - _message_helper: MessageHelper = PrivateAttr() + _channel: str = PrivateAttr(default=None) + _source: str = PrivateAttr(default=None) + _username: str = PrivateAttr(default=None) - def __init__(self, session_id: str, user_id: str, - channel: str = None, source: str = None, username: str = None, **kwargs): + def __init__(self, session_id: str, user_id: str, **kwargs): super().__init__(**kwargs) self._session_id = session_id self._user_id = user_id - self.channel = channel - self.source = source - self.username = username - self._message_helper = MessageHelper() def _run(self, **kwargs) -> str: raise NotImplementedError @@ -36,8 +31,14 @@ class MoviePilotTool(BaseTool): async def _arun(self, **kwargs) -> str: raise NotImplementedError - def _send_tool_message(self, message: str, title: str = None, **kwargs): - """发送工具执行消息""" + def set_message_attr(self, channel: str, source: str, username: str): + """设置消息属性""" + self._channel = channel + self._source = source + self._username = username + + def send_tool_message(self, message: str, title: str = "执行工具"): + """发送工具消息""" ToolChain().post_message( Notification( channel=self.channel, diff --git a/app/agent/tools/factory.py b/app/agent/tools/factory.py index 30e0b87c..63062b5e 100644 --- a/app/agent/tools/factory.py +++ b/app/agent/tools/factory.py @@ -37,12 +37,11 @@ class MoviePilotToolFactory: SendMessageTool ] for ToolClass in tool_definitions: - tools.append(ToolClass( + tool = ToolClass( session_id=session_id, - user_id=user_id, - channel=channel, - source=source, - username=username - )) + user_id=user_id + ) + tool.set_message_attr(channel=channel, source=source, username=username) + tools.append(tool) logger.info(f"成功创建 {len(tools)} 个MoviePilot工具") return tools diff --git a/app/agent/tools/impl/add_download.py b/app/agent/tools/impl/add_download.py index b05c4cbc..ee1c2910 100644 --- a/app/agent/tools/impl/add_download.py +++ b/app/agent/tools/impl/add_download.py @@ -2,35 +2,36 @@ from typing import Optional +from app.agent.tools.base import MoviePilotTool 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, + 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, **kwargs) -> str: - logger.info(f"执行工具: {self.name}, 参数: torrent_title={torrent_title}, torrent_url={torrent_url}, downloader={downloader}, save_path={save_path}, labels={labels}") - + 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}", title="添加下载") - + self.send_tool_message(f"正在添加下载任务: {torrent_title}", title="添加下载") + try: if not torrent_title or not torrent_url: error_message = "错误:必须提供种子标题和下载链接" - self._send_tool_message(error_message, title="下载失败") + self.send_tool_message(error_message, title="下载失败") return error_message # 使用DownloadChain添加下载 download_chain = DownloadChain() - + # 创建下载上下文 torrent_info = TorrentInfo( title=torrent_title, @@ -43,21 +44,21 @@ class AddDownloadTool(MoviePilotTool): ) did = download_chain.download_single( - context=context, - downloader=downloader, - save_path=save_path, + context=context, + downloader=downloader, + save_path=save_path, label=labels ) if did: success_message = f"成功添加下载任务:{torrent_title}" - self._send_tool_message(success_message, title="下载成功") + self.send_tool_message(success_message, title="下载成功") return success_message else: error_message = "添加下载任务失败" - self._send_tool_message(error_message, title="下载失败") + self.send_tool_message(error_message, title="下载失败") return error_message except Exception as e: error_message = f"添加下载任务时发生错误: {str(e)}" logger.error(f"添加下载任务失败: {e}", exc_info=True) - self._send_tool_message(error_message, title="下载失败") + self.send_tool_message(error_message, title="下载失败") return error_message diff --git a/app/agent/tools/impl/add_subscribe.py b/app/agent/tools/impl/add_subscribe.py index c60fd326..791f95f2 100644 --- a/app/agent/tools/impl/add_subscribe.py +++ b/app/agent/tools/impl/add_subscribe.py @@ -2,23 +2,24 @@ from typing import Optional +from app.agent.tools.base import MoviePilotTool 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, + async def _arun(self, title: str, year: str, media_type: str, explanation: str, season: Optional[int] = None, tmdb_id: Optional[str] = None, **kwargs) -> str: - logger.info(f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, tmdb_id={tmdb_id}") - + 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}", title="添加订阅") - + self.send_tool_message(f"正在添加订阅: {title} ({year}) - {media_type}", title="添加订阅") + try: subscribe_chain = SubscribeChain() # 转换 tmdb_id 为整数 @@ -28,25 +29,25 @@ class AddSubscribeTool(MoviePilotTool): tmdbid_int = int(tmdb_id) except (ValueError, TypeError): logger.warning(f"无效的 tmdb_id: {tmdb_id},将忽略") - + sid, message = subscribe_chain.add( - mtype=MediaType(media_type), - title=title, - year=year, - tmdbid=tmdbid_int, - season=season, + mtype=MediaType(media_type), + title=title, + year=year, + tmdbid=tmdbid_int, + season=season, username=self._user_id ) if sid: success_message = f"成功添加订阅:{title} ({year})" - self._send_tool_message(success_message, title="订阅成功") + self.send_tool_message(success_message, title="订阅成功") return success_message else: error_message = f"添加订阅失败:{message}" - self._send_tool_message(error_message, title="订阅失败") + self.send_tool_message(error_message, title="订阅失败") return error_message except Exception as e: error_message = f"添加订阅时发生错误: {str(e)}" logger.error(f"添加订阅失败: {e}", exc_info=True) - self._send_tool_message(error_message, title="订阅失败") + self.send_tool_message(error_message, title="订阅失败") return error_message diff --git a/app/agent/tools/impl/query_media_library.py b/app/agent/tools/impl/query_media_library.py index a409df6b..a73e26e6 100644 --- a/app/agent/tools/impl/query_media_library.py +++ b/app/agent/tools/impl/query_media_library.py @@ -1,30 +1,24 @@ """查询媒体库工具""" import json -from typing import Optional +from typing import Optional, List +from app.agent.tools.base import MoviePilotTool from app.db.mediaserver_oper import MediaServerOper from app.log import logger -from app.agent.tools.base import MoviePilotTool +from app.schemas import MediaServerItem 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, **kwargs) -> str: + async def _arun(self, explanation: str, media_type: Optional[str] = "all", + 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() - medias = media_server_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) + filtered_medias: List[MediaServerItem] = media_server_oper.exists(title=title, year=year, mtype=media_type) if filtered_medias: return json.dumps([m.to_dict() for m in filtered_medias]) return "媒体库中未找到相关媒体。" diff --git a/app/agent/tools/impl/search_media.py b/app/agent/tools/impl/search_media.py index f786bb40..ca6bb281 100644 --- a/app/agent/tools/impl/search_media.py +++ b/app/agent/tools/impl/search_media.py @@ -3,24 +3,24 @@ import json from typing import Optional +from app.agent.tools.base import MoviePilotTool from app.chain.media import MediaChain -from app.core.metainfo import MetaInfo from app.log import logger from app.schemas.types import MediaType -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, + async def _arun(self, title: str, explanation: str, year: Optional[str] = None, media_type: Optional[str] = None, season: Optional[int] = None, **kwargs) -> str: - logger.info(f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}") - + 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 ""), title="搜索中") - + self.send_tool_message(f"正在搜索媒体资源: {title}" + (f" ({year})" if year else ""), title="搜索中") + try: media_chain = MediaChain() # 构建搜索标题 @@ -31,10 +31,10 @@ class SearchMediaTool(MoviePilotTool): search_title = f"{search_title} {media_type}" if season: search_title = f"{search_title} S{season:02d}" - + # 使用 MediaChain.search 方法 meta, results = media_chain.search(title=search_title) - + # 过滤结果 if results: filtered_results = [] @@ -42,35 +42,32 @@ class SearchMediaTool(MoviePilotTool): if year and result.year != year: continue if media_type: - try: - if result.type != MediaType(media_type): - continue - except: - pass + if result.type != MediaType(media_type): + continue if season and result.season != season: continue filtered_results.append(result) - + if filtered_results: result_message = f"找到 {len(filtered_results)} 个相关媒体资源" - self._send_tool_message(result_message, title="搜索成功") - + self.send_tool_message(result_message, title="搜索成功") + # 发送详细结果 for i, result in enumerate(filtered_results[:5]): # 只显示前5个结果 - media_info = f"{i+1}. {result.title} ({result.year}) - {result.type.value if result.type else '未知'}" - self._send_tool_message(media_info, title="搜索结果") - + media_info = f"{i + 1}. {result.title} ({result.year}) - {result.type.value if result.type else '未知'}" + self.send_tool_message(media_info, title="搜索结果") + return json.dumps([r.to_dict() for r in filtered_results], ensure_ascii=False, indent=2) else: error_message = f"未找到符合条件的媒体资源: {title}" - self._send_tool_message(error_message, title="搜索完成") + self.send_tool_message(error_message, title="搜索完成") return error_message else: error_message = f"未找到相关媒体资源: {title}" - self._send_tool_message(error_message, title="搜索完成") + self.send_tool_message(error_message, title="搜索完成") return error_message except Exception as e: error_message = f"搜索媒体失败: {str(e)}" logger.error(f"搜索媒体失败: {e}", exc_info=True) - self._send_tool_message(error_message, title="搜索失败") + self.send_tool_message(error_message, title="搜索失败") return error_message diff --git a/app/agent/tools/impl/search_torrents.py b/app/agent/tools/impl/search_torrents.py index 7174ba73..1ede8ea4 100644 --- a/app/agent/tools/impl/search_torrents.py +++ b/app/agent/tools/impl/search_torrents.py @@ -3,24 +3,25 @@ import json from typing import List, Optional +from app.agent.tools.base import MoviePilotTool 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, + 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, **kwargs) -> str: - logger.info(f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}") - + 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 ""), title="搜索种子") - + self.send_tool_message(f"正在搜索种子资源: {title}" + (f" ({year})" if year else ""), title="搜索种子") + try: search_chain = SearchChain() torrents = search_chain.search_by_title(title=title, sites=sites) @@ -30,33 +31,30 @@ class SearchTorrentsTool(MoviePilotTool): if year and torrent.meta_info and torrent.meta_info.year != year: continue if media_type and torrent.media_info: - try: - if torrent.media_info.type != MediaType(media_type): - continue - except: - pass + if torrent.media_info.type != MediaType(media_type): + continue if season and torrent.meta_info 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, title="搜索成功") - + self.send_tool_message(result_message, title="搜索成功") + # 发送详细结果 for i, torrent in enumerate(filtered_torrents[:5]): # 只显示前5个结果 torrent_title = torrent.torrent_info.title if torrent.torrent_info else torrent.meta_info.title if torrent.meta_info else "未知" site_name = torrent.torrent_info.site_name if torrent.torrent_info else "未知站点" - torrent_info = f"{i+1}. {torrent_title} - {site_name}" - self._send_tool_message(torrent_info, title="搜索结果") - + torrent_info = f"{i + 1}. {torrent_title} - {site_name}" + self.send_tool_message(torrent_info, title="搜索结果") + return json.dumps([t.to_dict() for t in filtered_torrents], ensure_ascii=False, indent=2) else: error_message = f"未找到相关种子资源: {title}" - self._send_tool_message(error_message, title="搜索完成") + self.send_tool_message(error_message, title="搜索完成") return error_message except Exception as e: error_message = f"搜索种子时发生错误: {str(e)}" logger.error(f"搜索种子失败: {e}", exc_info=True) - self._send_tool_message(error_message, title="搜索失败") + self.send_tool_message(error_message, title="搜索失败") return error_message diff --git a/app/agent/tools/impl/send_message.py b/app/agent/tools/impl/send_message.py index 48554540..1402bffb 100644 --- a/app/agent/tools/impl/send_message.py +++ b/app/agent/tools/impl/send_message.py @@ -2,9 +2,8 @@ from typing import Optional -from app.helper.message import MessageHelper -from app.log import logger from app.agent.tools.base import MoviePilotTool +from app.log import logger class SendMessageTool(MoviePilotTool): @@ -14,8 +13,7 @@ class SendMessageTool(MoviePilotTool): async def _arun(self, message: str, explanation: str, message_type: Optional[str] = "info", **kwargs) -> 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"MoviePilot助手通知 ({message_type})") + self.send_tool_message(message, title=message_type) return "消息已发送。" except Exception as e: logger.error(f"发送消息失败: {e}")