Migrate to mixins for Context and Workspace features

autogpt/agents/agent.py

@@ -31,10 +31,10 @@ from autogpt.models.agent_actions import (
 )
 from autogpt.models.command import CommandOutput
 from autogpt.models.context_item import ContextItem
-from autogpt.workspace import Workspace
 
 from .base import BaseAgent
-from .utils.context import ContextMixin
+from .features.context import ContextMixin
+from .features.workspace import WorkspaceMixin
 from .utils.exceptions import (
     AgentException,
     CommandExecutionError,
@@ -45,7 +45,7 @@ from .utils.exceptions import (
 logger = logging.getLogger(__name__)
 
 
-class Agent(BaseAgent, ContextMixin):
+class Agent(BaseAgent, ContextMixin, WorkspaceMixin):
     """Agent class for interacting with Auto-GPT."""
 
     def __init__(
@@ -68,9 +68,6 @@ class Agent(BaseAgent, ContextMixin):
         self.memory = memory
         """VectorMemoryProvider used to manage the agent's context (TODO)"""
 
-        self.workspace = Workspace(config.workspace_path, config.restrict_to_workspace)
-        """Workspace that the agent has access to, e.g. for reading/writing files."""
-
         self.created_at = datetime.now().strftime("%Y%m%d_%H%M%S")
         """Timestamp the agent was created; only used for structured debug logging."""
 

autogpt/agents/features/context.py

@@ -39,9 +39,10 @@ class ContextMixin:
     context: AgentContext
 
     def __init__(self, **kwargs):
-        super(ContextMixin, self).__init__(**kwargs)
         self.context = AgentContext()
 
+        super(ContextMixin, self).__init__(**kwargs)
+
 
 def get_agent_context(agent: BaseAgent) -> AgentContext | None:
     if isinstance(agent, ContextMixin):

autogpt/agents/features/workspace.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..base import BaseAgent
+
+from autogpt.config import Config
+from autogpt.workspace import Workspace
+
+
+class WorkspaceMixin:
+    """Mixin that adds workspace support to a class"""
+
+    workspace: Workspace
+    """Workspace that the agent has access to, e.g. for reading/writing files."""
+
+    def __init__(self, **kwargs):
+        config: Config = getattr(self, "config")
+        if not isinstance(config, Config):
+            raise ValueError(f"Cannot initialize Workspace for Agent without Config")
+        if not config.workspace_path:
+            raise ValueError(f"Cannot set up Workspace: no WORKSPACE_PATH in config")
+
+        self.workspace = Workspace(config.workspace_path, config.restrict_to_workspace)
+
+        super(WorkspaceMixin, self).__init__(**kwargs)
+
+
+def get_agent_workspace(agent: BaseAgent) -> Workspace | None:
+    if isinstance(agent, WorkspaceMixin):
+        return agent.workspace
+
+    return None

autogpt/agents/planning_agent.py

@@ -29,16 +29,16 @@ from autogpt.models.agent_actions import (
     ActionSuccessResult,
 )
 from autogpt.models.context_item import ContextItem
-from autogpt.workspace import Workspace
 
 from .agent import execute_command, extract_command
 from .base import BaseAgent
-from .utils.context import AgentContext
+from .features.context import ContextMixin
+from .features.workspace import WorkspaceMixin
 
 logger = logging.getLogger(__name__)
 
 
-class PlanningAgent(BaseAgent):
+class PlanningAgent(BaseAgent, ContextMixin, WorkspaceMixin):
     """Agent class for interacting with Auto-GPT."""
 
     ThoughtProcessID = Literal["plan", "action", "evaluate"]
@@ -63,9 +63,6 @@ class PlanningAgent(BaseAgent):
         self.memory = memory
         """VectorMemoryProvider used to manage the agent's context (TODO)"""
 
-        self.workspace = Workspace(config.workspace_path, config.restrict_to_workspace)
-        """Workspace that the agent has access to, e.g. for reading/writing files."""
-
         self.created_at = datetime.now().strftime("%Y%m%d_%H%M%S")
         """Timestamp the agent was created; only used for structured debug logging."""
 
@@ -74,9 +71,6 @@ class PlanningAgent(BaseAgent):
 
         self.action_history = ActionHistory()
 
-        self.context = AgentContext()
-        """Dynamic segment of the prompt, to provide the LLM with relevant context"""
-
         self.plan: list[str] = []
         """List of steps that the Agent plans to take"""
 

autogpt/commands/file_context.py

@@ -7,9 +7,12 @@ COMMAND_CATEGORY_TITLE = "File Operations"
 
 import contextlib
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-from autogpt.agents.agent import Agent
-from autogpt.agents.utils.context import get_agent_context
+if TYPE_CHECKING:
+    from autogpt.agents import Agent, BaseAgent
+
+from autogpt.agents.features.context import ContextMixin, get_agent_context
 from autogpt.agents.utils.exceptions import (
     CommandExecutionError,
     DuplicateOperationError,
@@ -20,6 +23,10 @@ from autogpt.models.context_item import FileContextItem, FolderContextItem
 from .decorators import sanitize_path_arg
 
 
+def compatible_with_agent(agent: BaseAgent):
+    return isinstance(agent, ContextMixin)
+
+
 @command(
     "open_file",
     "Open a file for editing, creating it if it does not exist yet",
@@ -30,6 +37,7 @@ from .decorators import sanitize_path_arg
             "required": True,
         }
     },
+    available=compatible_with_agent,
 )
 @sanitize_path_arg("file_path")
 def open_file(file_path: Path, agent: Agent) -> tuple[str, FileContextItem]:
@@ -78,6 +86,7 @@ def open_file(file_path: Path, agent: Agent) -> tuple[str, FileContextItem]:
             "required": True,
         }
     },
+    available=compatible_with_agent,
 )
 @sanitize_path_arg("path")
 def open_folder(path: Path, agent: Agent) -> tuple[str, FolderContextItem]: