fix agent

app/agent/tools/base.py

@@ -1,8 +1,11 @@
 """MoviePilot工具基类"""
+from abc import ABCMeta, abstractmethod
+from typing import Callable, Any
 
 from langchain.tools import BaseTool
 from pydantic import PrivateAttr
 
+from app.agent import StreamingCallbackHandler
 from app.chain import ChainBase
 from app.schemas import Notification
 
@@ -11,7 +14,7 @@ class ToolChain(ChainBase):
     pass
 
 
-class MoviePilotTool(BaseTool):
+class MoviePilotTool(BaseTool, metaclass=ABCMeta):
     """MoviePilot专用工具基类"""
 
     _session_id: str = PrivateAttr()
@@ -19,16 +22,26 @@ class MoviePilotTool(BaseTool):
     _channel: str = PrivateAttr(default=None)
     _source: str = PrivateAttr(default=None)
     _username: str = PrivateAttr(default=None)
+    _callback_handler: StreamingCallbackHandler = PrivateAttr(default=None)
 
     def __init__(self, session_id: str, user_id: str, **kwargs):
         super().__init__(**kwargs)
         self._session_id = session_id
         self._user_id = user_id
 
-    def _run(self, **kwargs) -> str:
-        raise NotImplementedError
+    def _run(self, *args: Any, **kwargs: Any) -> Any:
+        pass
 
     async def _arun(self, **kwargs) -> str:
+        """异步运行工具"""
+        # 发送运行工具前的消息
+        agent_message = await self._callback_handler.get_message()
+        if agent_message:
+            await self.send_tool_message(agent_message)
+        return await self.run(**kwargs)
+
+    @abstractmethod
+    async def run(self, **kwargs) -> str:
         raise NotImplementedError
 
     def set_message_attr(self, channel: str, source: str, username: str):
@@ -37,9 +50,13 @@ class MoviePilotTool(BaseTool):
         self._source = source
         self._username = username
 
-    def send_tool_message(self, message: str, title: str = "执行工具"):
+    def set_callback_handler(self, callback_handler: StreamingCallbackHandler):
+        """设置回调处理器"""
+        self._callback_handler = callback_handler
+
+    async def send_tool_message(self, message: str, title: str = "执行工具"):
         """发送工具消息"""
-        ToolChain().post_message(
+        await ToolChain().async_post_message(
             Notification(
                 channel=self._channel,
                 source=self._source,