From b0405ac8ee1903bc0ee531db138401b41498d265 Mon Sep 17 00:00:00 2001
From: Reinier van der Leer <github@pwuts.nl>
Date: Wed, 23 Aug 2023 00:59:17 +0200
Subject: [PATCH] Add Context segment to prompt

---
 autogpt/agents/agent.py              |  2 +-
 autogpt/agents/features/context.py   | 14 +++++++++++++-
 autogpt/agents/features/workspace.py |  5 +++--
 autogpt/agents/planning_agent.py     |  2 +-
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/autogpt/agents/agent.py b/autogpt/agents/agent.py
index 2ce2fc76..e2d8fb27 100644
--- a/autogpt/agents/agent.py
+++ b/autogpt/agents/agent.py
@@ -45,7 +45,7 @@ from .utils.exceptions import (
 logger = logging.getLogger(__name__)
 
 
-class Agent(BaseAgent, ContextMixin, WorkspaceMixin):
+class Agent(ContextMixin, WorkspaceMixin, BaseAgent):
     """Agent class for interacting with Auto-GPT."""
 
     def __init__(
diff --git a/autogpt/agents/features/context.py b/autogpt/agents/features/context.py
index 3a2f81e3..847f318d 100644
--- a/autogpt/agents/features/context.py
+++ b/autogpt/agents/features/context.py
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from ..base import BaseAgent
 
+from autogpt.llm.base import Message
 from autogpt.models.context_item import ContextItem
 
 
@@ -34,7 +35,7 @@ class AgentContext:
 
 
 class ContextMixin:
-    """Mixin that adds context support to a class"""
+    """Mixin that adds context support to a BaseAgent subclass"""
 
     context: AgentContext
 
@@ -43,6 +44,17 @@ class ContextMixin:
 
         super(ContextMixin, self).__init__(**kwargs)
 
+    def construct_base_prompt(self, *args, **kwargs):
+        if kwargs.get("append_messages") is None:
+            kwargs["append_messages"] = []
+
+        if self.context:
+            kwargs["append_messages"].insert(
+                0, Message("system", "# Context\n" + self.context.format_numbered())
+            )
+
+        return super(ContextMixin, self).construct_base_prompt(*args, **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
index 41fffaea..64491057 100644
--- a/autogpt/agents/features/workspace.py
+++ b/autogpt/agents/features/workspace.py
@@ -16,6 +16,9 @@ class WorkspaceMixin:
     """Workspace that the agent has access to, e.g. for reading/writing files."""
 
     def __init__(self, **kwargs):
+        # Initialize other bases first, because we need the config from BaseAgent
+        super(WorkspaceMixin, self).__init__(**kwargs)
+
         config: Config = getattr(self, "config")
         if not isinstance(config, Config):
             raise ValueError(f"Cannot initialize Workspace for Agent without Config")
@@ -24,8 +27,6 @@ class WorkspaceMixin:
 
         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):
diff --git a/autogpt/agents/planning_agent.py b/autogpt/agents/planning_agent.py
index 366d3885..bdc064fc 100644
--- a/autogpt/agents/planning_agent.py
+++ b/autogpt/agents/planning_agent.py
@@ -38,7 +38,7 @@ from .features.workspace import WorkspaceMixin
 logger = logging.getLogger(__name__)
 
 
-class PlanningAgent(BaseAgent, ContextMixin, WorkspaceMixin):
+class PlanningAgent(ContextMixin, WorkspaceMixin, BaseAgent):
     """Agent class for interacting with Auto-GPT."""
 
     ThoughtProcessID = Literal["plan", "action", "evaluate"]