clean : memory is tested and functional

Author: colinfrisch

mesa_llm/llm_agent.py

@@ -16,18 +16,13 @@ class LLMAgent:
         llm (ModuleLLM): The internal LLM interface used by the agent.
         memory (Memory | None): The memory module attached to this agent, if any.
 
-    Notes:
-        - Each agent can only have one memory instance associated with it.
-        - If no memory is passed at initialization, one can be attached later using `attach_memory()`.
-        - Reassigning or replacing memory after it's been attached is not allowed and will raise a ValueError.
     """
 
     def __init__(
         self,
         api_key: str,
         model: str = "openai/gpt-4o",
         system_prompt: str | None = None,
-        memory: Memory | None = None,
     ):
         self.llm = ModuleLLM(api_key=api_key, model=model, system_prompt=system_prompt)
         self._memory = Memory(

mesa_llm/memory.py

@@ -1,10 +1,13 @@
 import os
 from collections import deque
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
-from mesa_llm.llm_agent import LLMAgent
 from mesa_llm.module_llm import ModuleLLM
 
+if TYPE_CHECKING:
+    from mesa_llm.llm_agent import LLMAgent
+
 
 @dataclass
 class MemoryEntry:
@@ -30,7 +33,7 @@ class Memory:
 
     def __init__(
         self,
-        agent: LLMAgent,
+        agent: "LLMAgent",
         short_term_capacity: int = 5,
         consolidation_capacity: int = 2,
         api_key: str = os.getenv("OPENAI_API_KEY"),
@@ -85,6 +88,12 @@ def get_short_term_memory(self) -> list[MemoryEntry]:
         """
         return list(self.short_term_memory)
 
+    def get_long_term_memory(self) -> str:
+        """
+        Get the long term memory
+        """
+        return self.long_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

mesa_llm/module_llm.py

@@ -20,15 +20,15 @@ def __init__(self, api_key: str, model: str, system_prompt: str | None = None):
             system_prompt: The system prompt to use for the LLM
         """
         self.api_key = api_key
-        self.model=model
+        self.model = model
         self.system_prompt = system_prompt
         provider = self.model.split("/")[0].upper()
         os.environ[f"{provider}_API_KEY"] = self.api_key
 
     def set_system_prompt(self, system_prompt: str):
         """Set or update the system prompt."""
         self.system_prompt = system_prompt
-    
+
     def set_model(self, api_key: str, model: str):
         """Set or update the model and API key."""
         self.api_key = api_key
@@ -53,9 +53,10 @@ def generate(self, prompt: str) -> str:
     from dotenv import load_dotenv
 
     load_dotenv()
+    print("ready to go ------------------------------")
 
-    api_key = os.getenv("OPENAI_API_KEY")  # Or simply your API key
-    llm = ModuleLLM(api_key, "openai/gpt-4o")
+    api_key = os.getenv("GEMINI_API_KEY")  # Or simply your API key
+    llm = ModuleLLM(api_key, "gemini/gemini-2.0-flash")
 
     response = llm.generate("Hello, how are you?")
     print(response.choices[0].message.content)