From 4e64519a263c97fa696c10d1b8c27a5243547040 Mon Sep 17 00:00:00 2001 From: Reinier van der Leer Date: Tue, 22 Aug 2023 22:52:26 +0200 Subject: [PATCH] Migrate to mixins for Context and Workspace features --- autogpt/agents/agent.py | 9 ++--- autogpt/agents/{utils => features}/context.py | 3 +- autogpt/agents/features/workspace.py | 34 +++++++++++++++++++ autogpt/agents/planning_agent.py | 12 ++----- autogpt/commands/file_context.py | 13 +++++-- 5 files changed, 53 insertions(+), 18 deletions(-) rename autogpt/agents/{utils => features}/context.py (99%) create mode 100644 autogpt/agents/features/workspace.py diff --git a/autogpt/agents/agent.py b/autogpt/agents/agent.py index 09e2d775..2ce2fc76 100644 --- a/autogpt/agents/agent.py +++ b/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.""" diff --git a/autogpt/agents/utils/context.py b/autogpt/agents/features/context.py similarity index 99% rename from autogpt/agents/utils/context.py rename to autogpt/agents/features/context.py index ea3ef158..3a2f81e3 100644 --- a/autogpt/agents/utils/context.py +++ b/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): diff --git a/autogpt/agents/features/workspace.py b/autogpt/agents/features/workspace.py new file mode 100644 index 00000000..41fffaea --- /dev/null +++ b/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 diff --git a/autogpt/agents/planning_agent.py b/autogpt/agents/planning_agent.py index 22187bde..366d3885 100644 --- a/autogpt/agents/planning_agent.py +++ b/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""" diff --git a/autogpt/commands/file_context.py b/autogpt/commands/file_context.py index e0b1d06a..bfcaba5e 100644 --- a/autogpt/commands/file_context.py +++ b/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]: