feat(agent): Fully abstracted file storage access with FileStorage (S…

…ignificant-Gravitas#6931)

* Rename `FileWorkspace` to `FileStorage`
   - `autogpt.file_workspace` -> `autogpt.file_storage`
   - `LocalFileWorkspace` -> `LocalFileStorage`
   - `S3FileWorkspace` -> `S3FileStorage`
   - `GCSFileWorkspace` -> `GCSFileStorage`

* Rename `WORKSPACE_BACKEND` to `FILE_STORAGE_BACKEND`
* Rename `WORKSPACE_STORAGE_BUCKET` to `STORAGE_BUCKET`

* Rewrite `AgentManager` to use `FileStorage` rather than direct local file access
* Rename `AgentManager.retrieve_state(..)` method to `load_agent_state`
* Add docstrings to `AgentManager`

* Create `AgentFileManagerMixin` to replace `AgentFileManager`, `FileWorkspaceMixin`, `BaseAgent.attach_fs(..)`
* Replace `BaseAgentSettings.save_to_json_file(..)` method by `AgentFileManagerMixin.save_state()`
* Replace `BaseAgent.set_id(..)` method by `AgentFileManagerMixin.change_agent_id(..)`
* Remove `BaseAgentSettings.load_from_json_file(..)`
* Remove `AgentSettings.agent_data_dir`

* Update `AgentProtocolServer` to work with the new `FileStorage` system and `AgentFileManagerMixin`

* Make `agent_id` and `file_storage` parameters for creating an Agent:
   - `create_agent`, `configure_agent_with_state`, `_configure_agent`, `create_agent_state` in `autogpt.agent_factory.configurators`
   - `generate_agent_for_task` in `autogpt.agent_factory.generators`
   - `Agent.__init__(..)`
   - `BaseAgent.__init__(..)`
   - Initialize and pass in `file_storage` in `autogpt.app.main.run_auto_gpt(..)` and `autogpt.app.main.run_auto_gpt_server(..)`

* Add `clone_with_subroot` to `FileStorage`
* Add `exists`, `make_dir`, `delete_dir`, `rename`, `list_files`, `list_folders` methods to `FileStorage`

* Update `autogpt.commands.file_operations` to use `FileStorage` and `AgentFileManagerMixin` features

* Update tests for `FileStorage` implementations and usages
* Rename `workspace` fixture to `storage`
   * Update conftest.py

autogpts/autogpt/.env.template

@@ -20,12 +20,12 @@ OPENAI_API_KEY=your-openai-api-key
 ## DISABLED_COMMAND_CATEGORIES - The list of categories of commands that are disabled (Default: None)
 # DISABLED_COMMAND_CATEGORIES=
 
-## WORKSPACE_BACKEND - Choose a storage backend for workspace contents
+## FILE_STORAGE_BACKEND - Choose a storage backend for contents
 ## Options: local, gcs, s3
-# WORKSPACE_BACKEND=local
+# FILE_STORAGE_BACKEND=local
 
-## WORKSPACE_STORAGE_BUCKET - GCS/S3 Bucket to store workspace contents in
-# WORKSPACE_STORAGE_BUCKET=autogpt
+## STORAGE_BUCKET - GCS/S3 Bucket to store contents in
+# STORAGE_BUCKET=autogpt
 
 ## GCS Credentials
 # see https://cloud.google.com/storage/docs/authentication#libauth

autogpts/autogpt/agbenchmark_config/benchmarks.py

@@ -3,10 +3,12 @@
 import sys
 from pathlib import Path
 
+from autogpt.agent_manager.agent_manager import AgentManager
 from autogpt.agents.agent import Agent, AgentConfiguration, AgentSettings
 from autogpt.app.main import _configure_openai_provider, run_interaction_loop
 from autogpt.commands import COMMAND_CATEGORIES
 from autogpt.config import AIProfile, ConfigBuilder
+from autogpt.file_storage import FileStorageBackendName, get_storage
 from autogpt.logs.config import configure_logging
 from autogpt.models.command_registry import CommandRegistry
 
@@ -42,6 +44,7 @@ def bootstrap_agent(task: str, continuous_mode: bool) -> Agent:
     agent_prompt_config.use_functions_api = config.openai_functions
     agent_settings = AgentSettings(
         name=Agent.default_settings.name,
+        agent_id=AgentManager.generate_id("AutoGPT-benchmark"),
         description=Agent.default_settings.description,
         ai_profile=ai_profile,
         config=AgentConfiguration(
@@ -55,13 +58,20 @@ def bootstrap_agent(task: str, continuous_mode: bool) -> Agent:
         history=Agent.default_settings.history.copy(deep=True),
     )
 
+    local = config.file_storage_backend == FileStorageBackendName.LOCAL
+    restrict_to_root = not local or config.restrict_to_workspace
+    file_storage = get_storage(
+        config.file_storage_backend, root_path="data", restrict_to_root=restrict_to_root
+    )
+    file_storage.initialize()
+
     agent = Agent(
         settings=agent_settings,
         llm_provider=_configure_openai_provider(config),
         command_registry=command_registry,
+        file_storage=file_storage,
         legacy_config=config,
     )
-    agent.attach_fs(config.app_data_dir / "agents" / "AutoGPT-benchmark")  # HACK
     return agent
 
 

autogpts/autogpt/autogpt/agent_factory/configurators.py

@@ -1,19 +1,21 @@
 from typing import Optional
 
-from autogpt.agent_manager import AgentManager
 from autogpt.agents.agent import Agent, AgentConfiguration, AgentSettings
 from autogpt.commands import COMMAND_CATEGORIES
 from autogpt.config import AIDirectives, AIProfile, Config
 from autogpt.core.resource.model_providers import ChatModelProvider
+from autogpt.file_storage.base import FileStorage
 from autogpt.logs.config import configure_chat_plugins
 from autogpt.models.command_registry import CommandRegistry
 from autogpt.plugins import scan_plugins
 
 
 def create_agent(
+    agent_id: str,
     task: str,
     ai_profile: AIProfile,
     app_config: Config,
+    file_storage: FileStorage,
     llm_provider: ChatModelProvider,
     directives: Optional[AIDirectives] = None,
 ) -> Agent:
@@ -23,41 +25,46 @@ def create_agent(
         directives = AIDirectives.from_file(app_config.prompt_settings_file)
 
     agent = _configure_agent(
+        agent_id=agent_id,
         task=task,
         ai_profile=ai_profile,
         directives=directives,
         app_config=app_config,
+        file_storage=file_storage,
         llm_provider=llm_provider,
     )
 
-    agent.state.agent_id = AgentManager.generate_id(agent.ai_profile.ai_name)
-
     return agent
 
 
 def configure_agent_with_state(
     state: AgentSettings,
     app_config: Config,
+    file_storage: FileStorage,
     llm_provider: ChatModelProvider,
 ) -> Agent:
     return _configure_agent(
         state=state,
         app_config=app_config,
+        file_storage=file_storage,
         llm_provider=llm_provider,
     )
 
 
 def _configure_agent(
     app_config: Config,
     llm_provider: ChatModelProvider,
+    file_storage: FileStorage,
+    agent_id: str = "",
     task: str = "",
     ai_profile: Optional[AIProfile] = None,
     directives: Optional[AIDirectives] = None,
     state: Optional[AgentSettings] = None,
 ) -> Agent:
-    if not (state or task and ai_profile and directives):
+    if not (state or agent_id and task and ai_profile and directives):
         raise TypeError(
-            "Either (state) or (task, ai_profile, directives) must be specified"
+            "Either (state) or (agent_id, task, ai_profile, directives)"
+            " must be specified"
         )
 
     app_config.plugins = scan_plugins(app_config)
@@ -70,6 +77,7 @@ def _configure_agent(
     )
 
     agent_state = state or create_agent_state(
+        agent_id=agent_id,
         task=task,
         ai_profile=ai_profile,
         directives=directives,
@@ -82,11 +90,13 @@ def _configure_agent(
         settings=agent_state,
         llm_provider=llm_provider,
         command_registry=command_registry,
+        file_storage=file_storage,
         legacy_config=app_config,
     )
 
 
 def create_agent_state(
+    agent_id: str,
     task: str,
     ai_profile: AIProfile,
     directives: AIDirectives,
@@ -96,6 +106,7 @@ def create_agent_state(
     agent_prompt_config.use_functions_api = app_config.openai_functions
 
     return AgentSettings(
+        agent_id=agent_id,
         name=Agent.default_settings.name,
         description=Agent.default_settings.description,
         task=task,

autogpts/autogpt/autogpt/agent_factory/generators.py

@@ -1,31 +1,38 @@
-from typing import TYPE_CHECKING
+from __future__ import annotations
 
-if TYPE_CHECKING:
-    from autogpt.agents.agent import Agent
-    from autogpt.config import Config
-    from autogpt.core.resource.model_providers.schema import ChatModelProvider
+from typing import TYPE_CHECKING
 
 from autogpt.config.ai_directives import AIDirectives
+from autogpt.file_storage.base import FileStorage
 
 from .configurators import _configure_agent
 from .profile_generator import generate_agent_profile_for_task
 
+if TYPE_CHECKING:
+    from autogpt.agents.agent import Agent
+    from autogpt.config import Config
+    from autogpt.core.resource.model_providers.schema import ChatModelProvider
+
 
 async def generate_agent_for_task(
+    agent_id: str,
     task: str,
-    app_config: "Config",
-    llm_provider: "ChatModelProvider",
-) -> "Agent":
+    app_config: Config,
+    file_storage: FileStorage,
+    llm_provider: ChatModelProvider,
+) -> Agent:
     base_directives = AIDirectives.from_file(app_config.prompt_settings_file)
     ai_profile, task_directives = await generate_agent_profile_for_task(
         task=task,
         app_config=app_config,
         llm_provider=llm_provider,
     )
     return _configure_agent(
+        agent_id=agent_id,
         task=task,
         ai_profile=ai_profile,
         directives=base_directives + task_directives,
         app_config=app_config,
+        file_storage=file_storage,
         llm_provider=llm_provider,
     )

autogpts/autogpt/autogpt/agent_manager/agent_manager.py

@@ -2,47 +2,44 @@
 
 import uuid
 from pathlib import Path
-from typing import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    from autogpt.agents.agent import AgentSettings
-
-from autogpt.agents.utils.agent_file_manager import AgentFileManager
+from autogpt.agents.agent import AgentSettings
+from autogpt.file_storage.base import FileStorage
 
 
 class AgentManager:
-    def __init__(self, app_data_dir: Path):
-        self.agents_dir = app_data_dir / "agents"
-        if not self.agents_dir.exists():
-            self.agents_dir.mkdir()
+    def __init__(self, file_storage: FileStorage):
+        self.file_manager = file_storage.clone_with_subroot("agents")
 
     @staticmethod
     def generate_id(agent_name: str) -> str:
+        """Generate a unique ID for an agent given agent name."""
         unique_id = str(uuid.uuid4())[:8]
         return f"{agent_name}-{unique_id}"
 
     def list_agents(self) -> list[str]:
-        return [
-            dir.name
-            for dir in self.agents_dir.iterdir()
-            if dir.is_dir() and AgentFileManager(dir).state_file_path.exists()
-        ]
-
-    def get_agent_dir(self, agent_id: str, must_exist: bool = False) -> Path:
+        """Return all agent directories within storage."""
+        agent_dirs: list[str] = []
+        for dir in self.file_manager.list_folders():
+            if self.file_manager.exists(dir / "state.json"):
+                agent_dirs.append(dir.name)
+        return agent_dirs
+
+    def get_agent_dir(self, agent_id: str) -> Path:
+        """Return the directory of the agent with the given ID."""
         assert len(agent_id) > 0
-        agent_dir = self.agents_dir / agent_id
-        if must_exist and not agent_dir.exists():
+        agent_dir: Path | None = None
+        if self.file_manager.exists(agent_id):
+            agent_dir = self.file_manager.root / agent_id
+        else:
             raise FileNotFoundError(f"No agent with ID '{agent_id}'")
         return agent_dir
 
-    def retrieve_state(self, agent_id: str) -> AgentSettings:
-        from autogpt.agents.agent import AgentSettings
-
-        agent_dir = self.get_agent_dir(agent_id, True)
-        state_file = AgentFileManager(agent_dir).state_file_path
-        if not state_file.exists():
+    def load_agent_state(self, agent_id: str) -> AgentSettings:
+        """Load the state of the agent with the given ID."""
+        state_file_path = Path(agent_id) / "state.json"
+        if not self.file_manager.exists(state_file_path):
             raise FileNotFoundError(f"Agent with ID '{agent_id}' has no state.json")
 
-        state = AgentSettings.load_from_json_file(state_file)
-        state.agent_data_dir = agent_dir
-        return state
+        state = self.file_manager.read_file(state_file_path)
+        return AgentSettings.parse_raw(state)

autogpts/autogpt/autogpt/agents/agent.py

@@ -6,10 +6,6 @@
 from datetime import datetime
 from typing import TYPE_CHECKING, Optional
 
-if TYPE_CHECKING:
-    from autogpt.config import Config
-    from autogpt.models.command_registry import CommandRegistry
-
 import sentry_sdk
 from pydantic import Field
 
@@ -20,6 +16,7 @@
     ChatMessage,
     ChatModelProvider,
 )
+from autogpt.file_storage.base import FileStorage
 from autogpt.llm.api_manager import ApiManager
 from autogpt.logs.log_cycle import (
     CURRENT_CONTEXT_FILE_NAME,
@@ -39,8 +36,8 @@
 from autogpt.models.context_item import ContextItem
 
 from .base import BaseAgent, BaseAgentConfiguration, BaseAgentSettings
+from .features.agent_file_manager import AgentFileManagerMixin
 from .features.context import ContextMixin
-from .features.file_workspace import FileWorkspaceMixin
 from .features.watchdog import WatchdogMixin
 from .prompt_strategies.one_shot import (
     OneShotAgentPromptConfiguration,
@@ -54,6 +51,10 @@
     UnknownCommandError,
 )
 
+if TYPE_CHECKING:
+    from autogpt.config import Config
+    from autogpt.models.command_registry import CommandRegistry
+
 logger = logging.getLogger(__name__)
 
 
@@ -72,7 +73,7 @@ class AgentSettings(BaseAgentSettings):
 
 class Agent(
     ContextMixin,
-    FileWorkspaceMixin,
+    AgentFileManagerMixin,
     WatchdogMixin,
     BaseAgent,
     Configurable[AgentSettings],
@@ -91,6 +92,7 @@ def __init__(
         settings: AgentSettings,
         llm_provider: ChatModelProvider,
         command_registry: CommandRegistry,
+        file_storage: FileStorage,
         legacy_config: Config,
     ):
         prompt_strategy = OneShotAgentPromptStrategy(
@@ -102,6 +104,7 @@ def __init__(
             llm_provider=llm_provider,
             prompt_strategy=prompt_strategy,
             command_registry=command_registry,
+            file_storage=file_storage,
             legacy_config=legacy_config,
         )