Merge branch 'jxxghp:v2' into v2

Author: Seed680

app/agent/__init__.py

@@ -7,7 +7,7 @@
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_community.callbacks import get_openai_callback
 from langchain_core.chat_history import InMemoryChatMessageHistory
-from langchain_core.messages import HumanMessage, AIMessage, ToolCall
+from langchain_core.messages import HumanMessage, AIMessage, ToolCall, ToolMessage, SystemMessage
 from langchain_core.runnables.history import RunnableWithMessageHistory
 
 from app.agent.callback import StreamingCallbackHandler
@@ -56,9 +56,6 @@ def __init__(self, session_id: str, user_id: str = None,
         # 工具
         self.tools = self._initialize_tools()
 
-        # 会话存储
-        self.session_store = self._initialize_session_store()
-
         # 提示词模板
         self.prompt = self._initialize_prompt()
 
@@ -127,7 +124,8 @@ def _initialize_tools(self) -> List:
             channel=self.channel,
             source=self.source,
             username=self.username,
-            callback_handler=self.callback_handler
+            callback_handler=self.callback_handler,
+            memory_mananger=self.memory_manager
         )
 
     @staticmethod
@@ -137,34 +135,36 @@ def _initialize_session_store() -> Dict[str, InMemoryChatMessageHistory]:
 
     def get_session_history(self, session_id: str) -> InMemoryChatMessageHistory:
         """获取会话历史"""
-        if session_id not in self.session_store:
-            chat_history = InMemoryChatMessageHistory()
-            messages: List[dict] = self.memory_manager.get_recent_messages_for_agent(
-                session_id=session_id,
-                user_id=self.user_id
-            )
-            if messages:
-                for msg in messages:
-                    if msg.get("role") == "user":
-                        chat_history.add_user_message(HumanMessage(content=msg.get("content", "")))
-                    elif msg.get("role") == "agent":
-                        chat_history.add_ai_message(AIMessage(content=msg.get("content", "")))
-                    elif msg.get("role") == "tool_call":
-                        metadata = msg.get("metadata", {})
-                        chat_history.add_ai_message(AIMessage(
+        chat_history = InMemoryChatMessageHistory()
+        messages: List[dict] = self.memory_manager.get_recent_messages_for_agent(
+            session_id=session_id,
+            user_id=self.user_id
+        )
+        if messages:
+            for msg in messages:
+                if msg.get("role") == "user":
+                    chat_history.add_message(HumanMessage(content=msg.get("content", "")))
+                elif msg.get("role") == "agent":
+                    chat_history.add_message(AIMessage(content=msg.get("content", "")))
+                elif msg.get("role") == "tool_call":
+                    metadata = msg.get("metadata", {})
+                    chat_history.add_message(
+                        AIMessage(
                             content=msg.get("content", ""),
-                            tool_calls=[ToolCall(
-                                id=metadata.get("call_id"),
-                                name=metadata.get("tool_name"),
-                                args=metadata.get("parameters"),
-                            )]
-                        ))
-                    elif msg.get("role") == "tool_result":
-                        chat_history.add_ai_message(AIMessage(content=msg.get("content", "")))
-                    elif msg.get("role") == "system":
-                        chat_history.add_ai_message(AIMessage(content=msg.get("content", "")))
-            self.session_store[session_id] = chat_history
-        return self.session_store[session_id]
+                            tool_calls=[
+                                ToolCall(
+                                    id=metadata.get("call_id"),
+                                    name=metadata.get("tool_name"),
+                                    args=metadata.get("parameters"),
+                                )
+                            ]
+                        )
+                    )
+                elif msg.get("role") == "tool_result":
+                    chat_history.add_message(ToolMessage(content=msg.get("content", "")))
+                elif msg.get("role") == "system":
+                    chat_history.add_message(SystemMessage(content=msg.get("content", "")))
+        return chat_history
 
     @staticmethod
     def _initialize_prompt() -> ChatPromptTemplate:
@@ -306,8 +306,6 @@ async def send_agent_message(self, message: str, title: str = "MoviePilot助手"
 
     async def cleanup(self):
         """清理智能体资源"""
-        if self.session_id in self.session_store:
-            del self.session_store[self.session_id]
         logger.info(f"MoviePilot智能体已清理: session_id={self.session_id}")
 
 

app/agent/memory/__init__.py

@@ -45,17 +45,27 @@ async def close(self):
 
         logger.info("对话记忆管理器已关闭")
 
+    @staticmethod
+    def get_memory_key(session_id: str, user_id: str):
+        """计算内存Key"""
+        return f"{user_id}:{session_id}" if user_id else session_id
+
+    @staticmethod
+    def get_redis_key(session_id: str, user_id: str):
+        """计算Redis Key"""
+        return f"agent_memory:{user_id}:{session_id}" if user_id else f"agent_memory:{session_id}"
+
     async def get_memory(self, session_id: str, user_id: str) -> ConversationMemory:
         """获取会话记忆"""
         # 首先检查缓存
-        cache_key = f"{user_id}:{session_id}" if user_id else session_id
+        cache_key = self.get_memory_key(session_id, user_id)
         if cache_key in self.memory_cache:
             return self.memory_cache[cache_key]
 
         # 尝试从Redis加载
         if settings.CACHE_BACKEND_TYPE == "redis":
             try:
-                redis_key = f"agent_memory:{user_id}:{session_id}" if user_id else f"agent_memory:{session_id}"
+                redis_key = self.get_redis_key(session_id, user_id)
                 memory_data = await self.redis_helper.get(redis_key, region="AI_AGENT")
                 if memory_data:
                     memory_dict = json.loads(memory_data) if isinstance(memory_data, str) else memory_data
@@ -180,15 +190,13 @@ def get_recent_messages_for_agent(
 
         如果消息Token数量超过模型最大上下文长度的阀值，会自动进行摘要裁剪
         """
-        cache_key = f"{user_id}:{session_id}" if user_id else session_id
+        cache_key = self.get_memory_key(session_id, user_id)
         memory = self.memory_cache.get(cache_key)
         if not memory:
             return []
 
         # 获取所有消息
-        messages = memory.messages
-
-        return messages
+        return memory.messages
 
     async def get_recent_messages(
             self,
@@ -218,7 +226,7 @@ async def clear_memory(self, session_id: str, user_id: str):
             del self.memory_cache[cache_key]
 
         if settings.CACHE_BACKEND_TYPE == "redis":
-            redis_key = f"agent_memory:{user_id}:{session_id}" if user_id else f"agent_memory:{session_id}"
+            redis_key = self.get_redis_key(session_id, user_id)
             await self.redis_helper.delete(redis_key, region="AI_AGENT")
 
         logger.info(f"会话记忆已清空: session_id={session_id}, user_id={user_id}")
@@ -229,14 +237,14 @@ async def _save_memory(self, memory: ConversationMemory):
         Redis中的记忆会自动通过TTL机制过期，无需手动清理
         """
         # 更新内存缓存
-        cache_key = f"{memory.user_id}:{memory.session_id}" if memory.user_id else memory.session_id
+        cache_key = self.get_memory_key(memory.session_id, memory.user_id)
         self.memory_cache[cache_key] = memory
 
         # 保存到Redis，设置TTL自动过期
         if settings.CACHE_BACKEND_TYPE == "redis":
             try:
                 memory_dict = memory.model_dump()
-                redis_key = f"agent_memory:{memory.user_id}:{memory.session_id}" if memory.user_id else f"agent_memory:{memory.session_id}"
+                redis_key = self.get_redis_key(memory.session_id, memory.user_id)
                 ttl = int(timedelta(days=settings.LLM_REDIS_MEMORY_RETENTION_DAYS).total_seconds())
                 await self.redis_helper.set(
                     redis_key,

app/agent/tools/base.py

@@ -1,11 +1,12 @@
 """MoviePilot工具基类"""
+import json
 from abc import ABCMeta, abstractmethod
-from typing import Callable, Any, Optional
+from typing import Any, Optional
 
 from langchain.tools import BaseTool
 from pydantic import PrivateAttr
 
-from app.agent import StreamingCallbackHandler
+from app.agent import StreamingCallbackHandler, ConversationMemoryManager
 from app.chain import ChainBase
 from app.log import logger
 from app.schemas import Notification
@@ -24,6 +25,7 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
     _source: str = PrivateAttr(default=None)
     _username: str = PrivateAttr(default=None)
     _callback_handler: StreamingCallbackHandler = PrivateAttr(default=None)
+    _memory_manager: ConversationMemoryManager = PrivateAttr(default=None)
 
     def __init__(self, session_id: str, user_id: str, **kwargs):
         super().__init__(**kwargs)
@@ -35,24 +37,53 @@ def _run(self, *args: Any, **kwargs: Any) -> Any:
 
     async def _arun(self, **kwargs) -> str:
         """异步运行工具"""
-        # 发送运行工具前的消息
+        # 发送和记忆工具调用前的信息
         agent_message = await self._callback_handler.get_message()
         if agent_message:
+            # 发送消息
             await self.send_tool_message(agent_message, title="MoviePilot助手")
-        # 发送执行工具说明
-        # 优先使用工具自定义的提示消息，如果没有则使用 explanation
+
+        # 记忆工具调用
+        await self._memory_manager.add_memory(
+            session_id=self._session_id,
+            user_id=self._user_id,
+            role="tool_call",
+            content=agent_message,
+            metadata={
+                "call_id": self.__class__.__name__,
+                "tool_name": self.__class__.__name__,
+                "parameters": kwargs
+            }
+        )
+
+        # 发送执行工具说明,优先使用工具自定义的提示消息，如果没有则使用 explanation
         tool_message = self.get_tool_message(**kwargs)
         if not tool_message:
             explanation = kwargs.get("explanation")
             if explanation:
                 tool_message = explanation
-        
         if tool_message:
             formatted_message = f"⚙️ => {tool_message}"
             await self.send_tool_message(formatted_message)
+
         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}')
+
+        # 记忆工具调用结果
+        if isinstance(result, str):
+            formated_result = result
+        elif isinstance(result, int, float):
+            formated_result = str(result)
+        else:
+            formated_result = json.dumps(result, ensure_ascii=False, indent=2)
+        await self._memory_manager.add_memory(
+            session_id=self._session_id,
+            user_id=self._user_id,
+            role="tool_result",
+            content=formated_result
+        )
+
         return result
 
     def get_tool_message(self, **kwargs) -> Optional[str]:
@@ -84,6 +115,10 @@ def set_callback_handler(self, callback_handler: StreamingCallbackHandler):
         """设置回调处理器"""
         self._callback_handler = callback_handler
 
+    def set_memory_manager(self, memory_manager: ConversationMemoryManager):
+        """设置记忆客理器"""
+        self._memory_manager = memory_manager
+
     async def send_tool_message(self, message: str, title: str = ""):
         """发送工具消息"""
         await ToolChain().async_post_message(

app/agent/tools/factory.py

@@ -51,7 +51,7 @@ class MoviePilotToolFactory:
     @staticmethod
     def create_tools(session_id: str, user_id: str,
                      channel: str = None, source: str = None, username: str = None,
-                     callback_handler: Callable = None) -> List[MoviePilotTool]:
+                     callback_handler: Callable = None, memory_mananger: Callable = None) -> List[MoviePilotTool]:
         """创建MoviePilot工具列表"""
         tools = []
         tool_definitions = [
@@ -102,6 +102,7 @@ def create_tools(session_id: str, user_id: str,
             )
             tool.set_message_attr(channel=channel, source=source, username=username)
             tool.set_callback_handler(callback_handler=callback_handler)
+            tool.set_memory_manager(memory_manager=memory_mananger)
             tools.append(tool)
 
         # 加载插件提供的工具
@@ -124,6 +125,7 @@ def create_tools(session_id: str, user_id: str,
                     )
                     tool.set_message_attr(channel=channel, source=source, username=username)
                     tool.set_callback_handler(callback_handler=callback_handler)
+                    tool.set_memory_manager(memory_manager=memory_mananger)
                     tools.append(tool)
                     plugin_tools_count += 1
                     logger.debug(f"成功加载插件 {plugin_name}({plugin_id}) 的工具: {ToolClass.__name__}")

app/agent/tools/impl/query_download_tasks.py

@@ -1,14 +1,16 @@
 """查询下载工具"""
 
 import json
-from typing import Optional, Type
+from typing import Optional, Type, List, Union
 
 from pydantic import BaseModel, Field
 
 from app.agent.tools.base import MoviePilotTool
 from app.chain.download import DownloadChain
 from app.db.downloadhistory_oper import DownloadHistoryOper
 from app.log import logger
+from app.schemas import TransferTorrent, DownloadingTorrent
+from app.schemas.types import TorrentStatus
 
 
 class QueryDownloadTasksInput(BaseModel):
@@ -27,6 +29,27 @@ class QueryDownloadTasksTool(MoviePilotTool):
     description: str = "Query download status and list download tasks. Can query all active downloads, or search for specific tasks by hash or title. Shows download progress, completion status, and task details from configured downloaders."
     args_schema: Type[BaseModel] = QueryDownloadTasksInput
 
+    def _get_all_torrents(self, download_chain: DownloadChain, downloader: Optional[str] = None) -> List[Union[TransferTorrent, DownloadingTorrent]]:
+        """
+        查询所有状态的任务（包括下载中和已完成的任务）
+        """
+        all_torrents = []
+        # 查询正在下载的任务
+        downloading_torrents = download_chain.list_torrents(
+            downloader=downloader, 
+            status=TorrentStatus.DOWNLOADING
+        ) or []
+        all_torrents.extend(downloading_torrents)
+        
+        # 查询已完成的任务（可转移状态）
+        transfer_torrents = download_chain.list_torrents(
+            downloader=downloader,
+            status=TorrentStatus.TRANSFER
+        ) or []
+        all_torrents.extend(transfer_torrents)
+        
+        return all_torrents
+
     def get_tool_message(self, **kwargs) -> Optional[str]:
         """根据查询参数生成友好的提示消息"""
         downloader = kwargs.get("downloader")
@@ -60,7 +83,7 @@ async def run(self, downloader: Optional[str] = None,
 
             # 如果提供了hash，直接查询该hash的任务（不限制状态）
             if hash:
-                torrents = download_chain.list_torrents(downloader=downloader, hashs=[hash])
+                torrents = download_chain.list_torrents(downloader=downloader, hashs=[hash]) or []
                 if not torrents:
                     return f"未找到hash为 {hash} 的下载任务（该任务可能已完成、已删除或不存在）"
                 # 转换为DownloadingTorrent格式
@@ -84,14 +107,25 @@ async def run(self, downloader: Optional[str] = None,
             elif title:
                 # 如果提供了title，查询所有任务并搜索匹配的标题
                 # 查询所有状态的任务
-                all_torrents = download_chain.list_torrents(downloader=downloader) or []
+                all_torrents = self._get_all_torrents(download_chain, downloader)
                 filtered_downloads = []
+                title_lower = title.lower()
                 for torrent in all_torrents:
-                    # 检查标题或名称是否匹配
-                    if (title.lower() in (torrent.title or "").lower()) or \
-                       (title.lower() in (torrent.name or "").lower()):
-                        # 获取下载历史信息
-                        history = DownloadHistoryOper().get_by_hash(torrent.hash)
+                    # 获取下载历史信息
+                    history = DownloadHistoryOper().get_by_hash(torrent.hash)
+                    
+                    # 检查标题或名称是否匹配（包括下载历史中的标题）
+                    matched = False
+                    # 检查torrent的title和name字段
+                    if (title_lower in (torrent.title or "").lower()) or \
+                       (title_lower in (torrent.name or "").lower()):
+                        matched = True
+                    # 检查下载历史中的标题
+                    if history and history.title:
+                        if title_lower in history.title.lower():
+                            matched = True
+                    
+                    if matched:
                         if history:
                             torrent.media = {
                                 "tmdbid": history.tmdbid,
@@ -110,7 +144,7 @@ async def run(self, downloader: Optional[str] = None,
                 # 根据status决定查询方式
                 if status == "downloading":
                     # 如果status为下载中，使用downloading方法
-                    downloads = download_chain.downloading(name=downloader)
+                    downloads = download_chain.downloading(name=downloader) or []
                     filtered_downloads = []
                     for dl in downloads:
                         if downloader and dl.downloader != downloader:
@@ -119,7 +153,7 @@ async def run(self, downloader: Optional[str] = None,
                 else:
                     # 其他状态（completed、paused、all），使用list_torrents查询所有任务
                     # 查询所有状态的任务
-                    all_torrents = download_chain.list_torrents(downloader=downloader) or []
+                    all_torrents = self._get_all_torrents(download_chain, downloader)
                     filtered_downloads = []
                     for torrent in all_torrents:
                         if downloader and torrent.downloader != downloader: