[Agent][Feat] Add remote file caching

erniebot-agent/erniebot_agent/agents/base.py

@@ -14,7 +14,7 @@
 
 import abc
 import json
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union, final
 
 from erniebot_agent import file_io
 from erniebot_agent.agents.callback.callback_manager import CallbackManager
@@ -27,9 +27,9 @@
     ToolResponse,
 )
 from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.file_io import protocol
 from erniebot_agent.file_io.base import File
 from erniebot_agent.file_io.file_manager import FileManager
-from erniebot_agent.file_io.protocol import is_local_file_id, is_remote_file_id
 from erniebot_agent.memory.base import Memory
 from erniebot_agent.messages import Message, SystemMessage
 from erniebot_agent.tools.base import BaseTool
@@ -74,13 +74,14 @@ def __init__(
         else:
             self._callback_manager = CallbackManager(callbacks)
         if file_manager is None:
-            file_manager = file_io.get_file_manager()
+            file_manager = file_io.get_global_file_manager()
         self._file_manager = file_manager
 
     @property
     def tools(self) -> List[BaseTool]:
         return self._tool_manager.get_tools()
 
+    @final
     async def async_run(self, prompt: str, files: Optional[List[File]] = None) -> AgentResponse:
         await self._callback_manager.on_run_start(agent=self, prompt=prompt)
         agent_resp = await self._async_run(prompt, files)
@@ -205,6 +206,7 @@ def _messages_to_dicts(messages):
     async def _async_run(self, prompt: str, files: Optional[List[File]] = None) -> AgentResponse:
         raise NotImplementedError
 
+    @final
     async def _async_run_tool(self, tool_name: str, tool_args: str) -> ToolResponse:
         tool = self._tool_manager.get_tool(tool_name)
         await self._callback_manager.on_tool_start(agent=self, tool=tool, input_args=tool_args)
@@ -216,6 +218,7 @@ async def _async_run_tool(self, tool_name: str, tool_args: str) -> ToolResponse:
         await self._callback_manager.on_tool_end(agent=self, tool=tool, response=tool_resp)
         return tool_resp
 
+    @final
     async def _async_run_llm(self, messages: List[Message], **opts: Any) -> LLMResponse:
         await self._callback_manager.on_llm_start(agent=self, llm=self.llm, messages=messages)
         try:
@@ -262,7 +265,7 @@ async def _sniff_and_extract_files_from_args(
         agent_files: List[AgentFile] = []
         for val in args.values():
             if isinstance(val, str):
-                if is_local_file_id(val):
+                if protocol.is_local_file_id(val):
                     if self._file_manager is None:
                         logger.warning(
                             f"A file is used by {repr(tool)}, but the agent has no file manager to fetch it."
@@ -271,7 +274,7 @@ async def _sniff_and_extract_files_from_args(
                     file = self._file_manager.look_up_file_by_id(val)
                     if file is None:
                         raise RuntimeError(f"Unregistered ID {repr(val)} is used by {repr(tool)}.")
-                elif is_remote_file_id(val):
+                elif protocol.is_remote_file_id(val):
                     if self._file_manager is None:
                         logger.warning(
                             f"A file is used by {repr(tool)}, but the agent has no file manager to fetch it."

erniebot-agent/erniebot_agent/agents/schema.py

@@ -18,8 +18,8 @@
 from dataclasses import dataclass
 from typing import Dict, List, Literal, Optional, Tuple, Union
 
+from erniebot_agent.file_io import protocol
 from erniebot_agent.file_io.base import File
-from erniebot_agent.file_io.protocol import extract_file_ids
 from erniebot_agent.messages import AIMessage, Message
 
 
@@ -80,6 +80,8 @@ def get_output_files(self) -> List[File]:
         return [agent_file.file for agent_file in self.files if agent_file.type == "output"]
 
     def get_tool_input_output_files(self, tool_name: str) -> Tuple[List[File], List[File]]:
+        # XXX: If a tool is used mutliple times, all related files will be
+        # returned in flattened lists.
         input_files: List[File] = []
         output_files: List[File] = []
         for agent_file in self.files:
@@ -94,7 +96,7 @@ def get_tool_input_output_files(self, tool_name: str) -> Tuple[List[File], List[
 
     def output_dict(self) -> Dict[str, List]:
         # 1. split the text into parts and add file id to each part
-        file_ids = extract_file_ids(self.text)
+        file_ids = protocol.extract_file_ids(self.text)
 
         places = []
         for file_id in file_ids:

erniebot-agent/erniebot_agent/file_io/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from erniebot_agent.file_io.factory import get_file_manager
+from erniebot_agent.file_io.factory import get_global_file_manager

erniebot-agent/erniebot_agent/file_io/base.py

@@ -23,7 +23,7 @@ def __init__(
         id: str,
         filename: str,
         byte_size: int,
-        created_at: int,
+        created_at: str,
         purpose: str,
         metadata: Dict[str, Any],
     ) -> None:
@@ -40,7 +40,7 @@ def __eq__(self, other: object) -> bool:
         if isinstance(other, File):
             return self.id == other.id
         else:
-            return False
+            return NotImplemented
 
     def __repr__(self) -> str:
         attrs_str = self._get_attrs_str()