first memory prototype

Author: colinfrisch

mesa_llm/memory.py

@@ -0,0 +1,112 @@
+import os
+from collections import deque
+from dataclasses import dataclass
+
+from mesa_llm.llm_agent import LLMAgent
+from mesa_llm.module_llm import ModuleLLM
+
+
+@dataclass
+class MemoryEntry:
+    type: str
+    content: str
+    step: int
+    metadata: dict
+
+
+class Memory:
+    """
+    Create a memory object that stores the agent's short and long term memory
+
+    Attributes:
+        agent : the agent that the memory belongs to
+
+    Memory is composed of
+        - A short term memory who stores the n (int) most recent interactions (observations, planning, discussions)
+        - A long term memory that is a summary of the memories that are removed from short term memory (summary
+        completed/refactored as it goes)
+
+    """
+
+    def __init__(
+        self,
+        agent: LLMAgent,
+        short_term_capacity: int = 5,
+        consolidation_capacity: int = 2,
+        api_key: str = os.getenv("OPENAI_API_KEY"),
+        llm_model: str = "openai/gpt-4o-mini",
+    ):
+        """
+        Initialize the memory
+
+        Args:
+            short_term_capacity : the number of interactions to store in the short term memory
+            api_key : the API key to use for the LLM
+            llm_model : the model to use for the summarization
+            agent : the agent that the memory belongs to
+        """
+        self.agent = agent
+        self.llm = ModuleLLM(api_key=api_key, model=llm_model)
+        self.capacity = short_term_capacity
+        self.consolidation_capacity = consolidation_capacity
+        self.short_term_memory = deque()
+        self.long_term_memory = ""
+
+        self.system_prompt = """
+        You are a helpful assistant that summarizes the short term memory into a long term memory.
+        The long term memory should be a summary of the short term memory that is concise and informative.
+        If the short term memory is empty, return the long term memory unchanged.
+        If the long term memory is not empty, update it to include the new information from the short term memory.
+        """
+
+        self.llm.set_system_prompt(self.system_prompt)
+
+    def add_to_memory(
+        self, type: str, content: str, step: int, metadata: dict | None = None
+    ):
+        """
+        Add a new entry to the memory
+        """
+        metadata = metadata or {}
+        new_entry = MemoryEntry(type, content, step, metadata)
+        self.short_term_memory.append(new_entry)
+
+        # Consolidate memory if the short term memory is over capacity
+        if len(self.short_term_memory) > self.capacity + self.consolidation_capacity:
+            memories_to_consolidate = [
+                self.short_term_memory.popleft()
+                for _ in range(self.consolidation_capacity)
+            ]
+            self.update_long_term_memory(memories_to_consolidate)
+
+    def get_short_term_memory(self) -> list[MemoryEntry]:
+        """
+        Get the short term memory
+        """
+        return list(self.short_term_memory)
+
+    def update_long_term_memory(self, memories_to_consolidate: list[MemoryEntry]):
+        """
+        Update the long term memory by summarizing the short term memory with a LLM
+        """
+        entries = [self.convert_entry_to_dict(m) for m in memories_to_consolidate]
+
+        prompt = f"""
+            Short term memory:
+                {entries}
+            Long term memory:
+                {self.long_term_memory}
+            """
+
+        self.long_term_memory = self.llm.generate(prompt, self.system_prompt)
+
+    def convert_entry_to_dict(self, entry: MemoryEntry) -> dict:
+        """
+        Convert a memory entry to a dictionary
+        """
+        return {
+            "type": entry.type,
+            "content": entry.content,
+            "step": entry.step,
+            "metadata": entry.metadata,
+        }