feat: Dynamic template system with editable system prompt fields.

memgpt/agent.py

@@ -6,7 +6,7 @@
 import json
 from pathlib import Path
 import traceback
-from typing import List, Tuple, Optional, cast
+from typing import List, Tuple, Optional, cast, Union
 
 from box import Box
 
@@ -16,7 +16,7 @@
 from memgpt.persistence_manager import PersistenceManager, LocalStateManager
 from memgpt.config import MemGPTConfig
 from memgpt.system import get_login_event, package_function_response, package_summarize_message, get_initial_boot_messages
-from memgpt.memory import CoreMemory as InContextMemory, summarize_messages
+from memgpt.memory import CoreMemory as InContextMemory, summarize_messages, CustomizableCoreMemory as CustomizableInContextMemory
 from memgpt.llm_api_tools import create, is_context_overflow_error
 from memgpt.utils import (
     get_tool_call_id,
@@ -42,6 +42,8 @@
 )
 from .errors import LLMError
 from .functions.functions import USER_FUNCTIONS_DIR, load_all_function_sets
+from .presets.default_templates import default_system_message_layout_template, default_core_memory_section_template
+from .prompts.prompt_template import PromptTemplate
 
 
 def link_functions(function_schemas):
@@ -92,6 +94,10 @@
     return linked_function_set
 
 
+def initialize_custom_memory(core_memory: dict, core_memory_limits: dict):
+    return CustomizableInContextMemory(core_memory, core_memory_limits)
+
+
 def initialize_memory(ai_notes, human_notes):
     if ai_notes is None:
         raise ValueError(ai_notes)
@@ -103,30 +109,69 @@
     return memory
 
 
-def construct_system_with_memory(system, memory, memory_edit_timestamp, archival_memory=None, recall_memory=None, include_char_count=True):
-    full_system_message = "\n".join(
-        [
-            system,
-            "\n",
-            f"### Memory [last modified: {memory_edit_timestamp.strip()}]",
-            f"{len(recall_memory) if recall_memory else 0} previous messages between you and the user are stored in recall memory (use functions to access them)",
-            f"{len(archival_memory) if archival_memory else 0} total memories you created are stored in archival memory (use functions to access them)",
-            "\nCore memory shown below (limited in size, additional information stored in archival / recall memory):",
-            f'<persona characters="{len(memory.persona)}/{memory.persona_char_limit}">' if include_char_count else "<persona>",
-            memory.persona,
-            "</persona>",
-            f'<human characters="{len(memory.human)}/{memory.human_char_limit}">' if include_char_count else "<human>",
-            memory.human,
-            "</human>",
-        ]
-    )
+def construct_system_with_memory(
+    system,
+    memory,
+    memory_edit_timestamp,
+    system_message_layout_template,
+    core_memory_section_template,
+    archival_memory=None,
+    recall_memory=None,
+    include_char_count=True,
+):
+    system_template = PromptTemplate.from_string(system_message_layout_template)
+
+    if isinstance(memory, InContextMemory):
+        core_memory_section_template = PromptTemplate.from_string(core_memory_section_template)
+        core_memory_content = (
+            core_memory_section_template.generate_prompt(
+                {
+                    "memory_key": "persona",
+                    "memory_value": memory.persona,
+                    "memory_value_length": len(memory.persona),
+                    "memory_value_limit": memory.persona_char_limit,
+                }
+            )
+            + "\n"
+        )
+        core_memory_content += (
+            core_memory_section_template.generate_prompt(
+                {
+                    "memory_key": "human",
+                    "memory_value": memory.human,
+                    "memory_value_length": len(memory.human),
+                    "memory_value_limit": memory.human_char_limit,
+                }
+            )
+            + "\n"
+        )
+        template_fields = {
+            "system": system,
+            "len_recall_memory": len(recall_memory) if recall_memory else 0,
+            "len_archival_memory": len(archival_memory) if archival_memory else 0,
+            "core_memory_content": core_memory_content,
+            "memory_edit_timestamp": memory_edit_timestamp.strip(),
+        }
+        full_system_message = system_template.generate_prompt(template_fields)
+    else:
+        core_memory_content = memory.get_memory_view(core_memory_section_template)
+        template_fields = {
+            "system": system,
+            "len_recall_memory": f"{len(recall_memory) if recall_memory else 0}",
+            "len_archival_memory": f"{len(archival_memory) if archival_memory else 0}",
+            "core_memory_content": core_memory_content,
+            "memory_edit_timestamp": memory_edit_timestamp.strip(),
+        }
+        full_system_message = system_template.generate_prompt(template_fields)
     return full_system_message
 
 
 def initialize_message_sequence(
     model,
     system,
     memory,
+    system_message_layout_template,
+    core_memory_section_template,
     archival_memory=None,
     recall_memory=None,
     memory_edit_timestamp=None,
@@ -136,7 +181,13 @@
         memory_edit_timestamp = get_local_time()
 
     full_system_message = construct_system_with_memory(
-        system, memory, memory_edit_timestamp, archival_memory=archival_memory, recall_memory=recall_memory
+        system,
+        memory,
+        memory_edit_timestamp,
+        system_message_layout_template,
+        core_memory_section_template,
+        archival_memory=archival_memory,
+        recall_memory=recall_memory,
     )
     first_user_message = get_login_event()  # event letting MemGPT know the user just logged in
 
@@ -184,21 +235,40 @@
         if "system" not in agent_state.state:
             raise ValueError(f"'system' not found in provided AgentState")
         self.system = agent_state.state["system"]
-
+        if "system_template" not in agent_state.state:
+            self.system_template = ""
+        else:
+            self.system_template = agent_state.state["system_template"]
+        if "system_message_layout_template" not in agent_state.state:
+            self.system_message_layout_template = default_system_message_layout_template
+        else:
+            self.system_message_layout_template = agent_state.state["system_message_layout_template"]
+        if "core_memory_section_template" not in agent_state.state:
+            self.core_memory_section_template = default_core_memory_section_template
+        else:
+            self.core_memory_section_template = agent_state.state["core_memory_section_template"]
+        if "system_template_fields" not in agent_state.state:
+            self.system_template_fields = {}
+        else:
+            self.system_template_fields = agent_state.state["system_template_fields"]
         if "functions" not in agent_state.state:
             raise ValueError(f"'functions' not found in provided AgentState")
         # Store the functions schemas (this is passed as an argument to ChatCompletion)
         self.functions = agent_state.state["functions"]  # these are the schema
         # Link the actual python functions corresponding to the schemas
         self.functions_python = {k: v["python_function"] for k, v in link_functions(function_schemas=self.functions).items()}
         assert all([callable(f) for k, f in self.functions_python.items()]), self.functions_python
-
-        # Initialize the memory object
-        if "persona" not in agent_state.state:
-            raise ValueError(f"'persona' not found in provided AgentState")
-        if "human" not in agent_state.state:
-            raise ValueError(f"'human' not found in provided AgentState")
-        self.memory = initialize_memory(ai_notes=agent_state.state["persona"], human_notes=agent_state.state["human"])
+        if "core_memory_type" in agent_state.state and agent_state.state["core_memory_type"] == "custom":
+            if "core_memory" not in agent_state.state:
+                raise ValueError(f"'core_memory' not found in provided AgentState")
+            self.memory = initialize_custom_memory(agent_state.state["core_memory"], agent_state.state["core_memory_limits"])
+        else:
+            # Initialize the memory object
+            if "persona" not in agent_state.state:
+                raise ValueError(f"'persona' not found in provided AgentState")
+            if "human" not in agent_state.state:
+                raise ValueError(f"'human' not found in provided AgentState")
+            self.memory = initialize_memory(ai_notes=agent_state.state["persona"], human_notes=agent_state.state["human"])
 
         # Interface must implement:
         # - internal_monologue
@@ -252,9 +322,7 @@
         else:
             # print(f"Agent.__init__ :: creating, state={agent_state.state['messages']}")
             init_messages = initialize_message_sequence(
-                self.model,
-                self.system,
-                self.memory,
+                self.model, self.system, self.memory, self.system_message_layout_template, self.core_memory_section_template
             )
             init_messages_objs = []
             for msg in init_messages:
@@ -639,7 +707,7 @@
 
             # Step 4: extend the message history
             if user_message is not None:
                 all_new_messages = [packed_user_message_obj] + all_response_messages
             else:
                 all_new_messages = all_response_messages
 
@@ -669,7 +737,7 @@
 
             self._append_to_messages(all_new_messages)
             all_new_messages_dicts = [msg.to_openai_dict() for msg in all_new_messages]
             return all_new_messages_dicts, heartbeat_request, function_failed, active_memory_warning, response.usage.completion_tokens
 
         except Exception as e:
             printd(f"step() failed\nuser_message = {user_message}\nerror = {e}")
@@ -816,6 +884,43 @@
             self.model,
             self.system,
             self.memory,
+            self.system_message_layout_template,
+            self.core_memory_section_template,
+            archival_memory=self.persistence_manager.archival_memory,
+            recall_memory=self.persistence_manager.recall_memory,
+        )[0]
+
+        diff = united_diff(curr_system_message["content"], new_system_message["content"])
+        printd(f"Rebuilding system with new memory...\nDiff:\n{diff}")
+
+        # Swap the system message out
+        self._swap_system_message(
+            Message.dict_to_message(
+                agent_id=self.agent_state.id, user_id=self.agent_state.user_id, model=self.model, openai_message_dict=new_system_message
+            )
+        )
+
+    def edit_system_template_field(self, field_name: str, field_value: Union[str, float, int], rebuild_system_template: bool = True):
+        """Edits a system template field"""
+        if field_name not in self.system_template_fields:
+            raise ValueError(f"'{field_name}' not found in system template fields")
+
+        self.system_template_fields[field_name] = field_value
+        if rebuild_system_template:
+            self.rebuild_system_template()
+
+    def rebuild_system_template(self):
+        """Rebuilds the system message with the latest template field values"""
+        curr_system_message = self.messages[0]  # this is the system + memory bank, not just the system prompt
+
+        template = PromptTemplate.from_string(self.system_template)
+
+        new_system_message = initialize_message_sequence(
+            self.model,
+            template.generate_prompt(self.system_template_fields),
+            self.memory,
+            self.system_message_layout_template,
+            self.core_memory_section_template,
             archival_memory=self.persistence_manager.archival_memory,
             recall_memory=self.persistence_manager.recall_memory,
         )[0]
@@ -916,13 +1021,32 @@
     #        self.ms.update_agent(agent=new_agent_state)
 
     def update_state(self) -> AgentState:
-        updated_state = {
-            "persona": self.memory.persona,
-            "human": self.memory.human,
-            "system": self.system,
-            "functions": self.functions,
-            "messages": [str(msg.id) for msg in self._messages],
-        }
+        if isinstance(self.memory, InContextMemory):
+            updated_state = {
+                "persona": self.memory.persona,
+                "human": self.memory.human,
+                "core_memory_type": "default",
+                "system_message_layout_template": self.system_message_layout_template,
+                "core_memory_section_template": self.core_memory_section_template,
+                "system": self.system,
+                "system_template": self.system_template,
+                "functions": self.functions,
+                "messages": [str(msg.id) for msg in self._messages],
+            }
+        elif isinstance(self.memory, CustomizableInContextMemory):
+            updated_state = {
+                "core_memory": self.memory.core_memory,
+                "core_memory_limits": self.memory.memory_field_limits,
+                "core_memory_type": "custom",
+                "system_message_layout_template": self.system_message_layout_template,
+                "core_memory_section_template": self.core_memory_section_template,
+                "system": self.system,
+                "system_template": self.system_template,
+                "functions": self.functions,
+                "messages": [str(msg.id) for msg in self._messages],
+            }
+        else:
+            updated_state = {}
 
         self.agent_state = AgentState(
             name=self.agent_state.name,

memgpt/main.py

@@ -294,7 +294,34 @@ def run_agent_loop(memgpt_agent, config: MemGPTConfig, first, ms: MetadataStore,
                             bold=True,
                         )
                         continue
-
+                elif user_input.lower().startswith("/edit_template_field"):
+                    try:
+                        field_name = questionary.text(
+                            "Enter the template field name:",
+                            multiline=multiline_input,
+                            qmark=">",
+                        ).ask()
+                        clear_line(strip_ui)
+
+                        field_value = questionary.text(
+                            "Enter the template field value:",
+                            multiline=multiline_input,
+                            qmark=">",
+                        ).ask()
+                        memgpt_agent.edit_system_template_field(field_name, field_value)
+                        clear_line(strip_ui)
+                        typer.secho(
+                            f"/edit_template_field succeeded: {field_value} : {field_value}",
+                            fg=typer.colors.GREEN,
+                            bold=True,
+                        )
+                    except ValueError as e:
+                        typer.secho(
+                            f"/edit_template_field failed:\n{e}",
+                            fg=typer.colors.RED,
+                            bold=True,
+                        )
+                        continue
                 # No skip options
                 elif user_input.lower() == "/wipe":
                     memgpt_agent = agent.Agent(interface)