[ShortTermMemory]

docs/swarms/memory/short_term_memory.md

@@ -0,0 +1,248 @@
+# Short-Term Memory Module Documentation
+
+## Introduction
+The Short-Term Memory module is a component of the SWARMS framework designed for managing short-term and medium-term memory in a multi-agent system. This documentation provides a detailed explanation of the Short-Term Memory module, its purpose, functions, and usage.
+
+### Purpose
+The Short-Term Memory module serves the following purposes:
+1. To store and manage messages in short-term memory.
+2. To provide functions for retrieving, updating, and clearing memory.
+3. To facilitate searching for specific terms within the memory.
+4. To enable saving and loading memory data to/from a file.
+
+### Class Definition
+```python
+class ShortTermMemory(BaseStructure):
+    def __init__(
+        self,
+        return_str: bool = True,
+        autosave: bool = True,
+        *args,
+        **kwargs,
+    ):
+    ...
+```
+
+#### Parameters
+| Parameter           | Type     | Default Value | Description                                                                                                      |
+|---------------------|----------|---------------|------------------------------------------------------------------------------------------------------------------|
+| `return_str`        | bool     | True          | If True, returns memory as a string.                                                                            |
+| `autosave`          | bool     | True          | If True, enables automatic saving of memory data to a file.                                                     |
+| `*args`, `**kwargs` |          |               | Additional arguments and keyword arguments (not used in the constructor but allowed for flexibility).          |
+
+### Functions
+
+#### 1. `add`
+```python
+def add(self, role: str = None, message: str = None, *args, **kwargs):
+```
+
+- Adds a message to the short-term memory.
+- Parameters:
+  - `role` (str, optional): Role associated with the message.
+  - `message` (str, optional): The message to be added.
+- Returns: The added memory.
+
+##### Example 1: Adding a Message to Short-Term Memory
+```python
+memory.add(role="Agent 1", message="Received task assignment.")
+```
+
+##### Example 2: Adding Multiple Messages to Short-Term Memory
+```python
+messages = [("Agent 1", "Received task assignment."), ("Agent 2", "Task completed.")]
+for role, message in messages:
+    memory.add(role=role, message=message)
+```
+
+#### 2. `get_short_term`
+```python
+def get_short_term(self):
+```
+
+- Retrieves the short-term memory.
+- Returns: The contents of the short-term memory.
+
+##### Example: Retrieving Short-Term Memory
+```python
+short_term_memory = memory.get_short_term()
+for entry in short_term_memory:
+    print(entry["role"], ":", entry["message"])
+```
+
+#### 3. `get_medium_term`
+```python
+def get_medium_term(self):
+```
+
+- Retrieves the medium-term memory.
+- Returns: The contents of the medium-term memory.
+
+##### Example: Retrieving Medium-Term Memory
+```python
+medium_term_memory = memory.get_medium_term()
+for entry in medium_term_memory:
+    print(entry["role"], ":", entry["message"])
+```
+
+#### 4. `clear_medium_term`
+```python
+def clear_medium_term(self):
+```
+
+- Clears the medium-term memory.
+
+##### Example: Clearing Medium-Term Memory
+```python
+memory.clear_medium_term()
+```
+
+#### 5. `get_short_term_memory_str`
+```python
+def get_short_term_memory_str(self, *args, **kwargs):
+```
+
+- Retrieves the short-term memory as a string.
+- Returns: A string representation of the short-term memory.
+
+##### Example: Getting Short-Term Memory as a String
+```python
+short_term_memory_str = memory.get_short_term_memory_str()
+print(short_term_memory_str)
+```
+
+#### 6. `update_short_term`
+```python
+def update_short_term(self, index, role: str, message: str, *args, **kwargs):
+```
+
+- Updates a message in the short-term memory.
+- Parameters:
+  - `index` (int): The index of the message to update.
+  - `role` (str): New role for the message.
+  - `message` (str): New message content.
+- Returns: None.
+
+##### Example: Updating a Message in Short-Term Memory
+```python
+memory.update_short_term(index=0, role="Updated Role", message="Updated message content.")
+```
+
+#### 7. `clear`
+```python
+def clear(self):
+```
+
+- Clears the short-term memory.
+
+##### Example: Clearing Short-Term Memory
+```python
+memory.clear()
+```
+
+#### 8. `search_memory`
+```python
+def search_memory(self, term):
+```
+
+- Searches the memory for a specific term.
+- Parameters:
+  - `term` (str): The term to search for.
+- Returns: A dictionary containing search results for short-term and medium-term memory.
+
+##### Example: Searching Memory for a Term
+```python
+search_results = memory.search_memory("task")
+print("Short-Term Memory Results:", search_results["short_term"])
+print("Medium-Term Memory Results:", search_results["medium_term"])
+```
+
+#### 9. `return_shortmemory_as_str`
+```python
+def return_shortmemory_as_str(self):
+```
+
+- Returns the memory as a string.
+
+##### Example: Returning Short-Term Memory as a String
+```python
+short_term_memory_str = memory.return_shortmemory_as_str()
+print(short_term_memory_str)
+```
+
+#### 10. `move_to_medium_term`
+```python
+def move_to_medium_term(self, index):
+```
+
+- Moves a message from the short-term memory to the medium-term memory.
+- Parameters:
+  - `index` (int): The index of the message to move.
+
+##### Example: Moving a Message to Medium-Term Memory
+```python
+memory.move_to_medium_term(index=0)
+```
+
+#### 11. `return_medium_memory_as_str`
+```python
+def return_medium_memory_as_str(self):
+```
+
+- Returns the medium-term memory as a string.
+
+##### Example: Returning Medium-Term Memory as a String
+```python
+medium_term_memory_str = memory.return_medium_memory_as_str()
+print(medium_term_memory_str)
+```
+
+#### 12. `save_to_file`
+```python
+def save_to_file(self, filename: str):
+```
+
+- Saves the memory data to a file.
+- Parameters:
+  - `filename` (str): The name of the file to save the data to.
+
+##### Example: Saving Memory Data to a File
+```python
+memory.save_to_file("memory_data.json")
+```
+
+#### 13. `load_from_file`
+```python
+def load_from_file(self, filename: str, *args, **kwargs):
+```
+
+- Loads memory data from a file.
+- Parameters:
+  - `filename` (str): The name of the file to load data from.
+
+##### Example: Loading Memory Data from a File
+```python
+memory.load_from_file("memory_data.json")
+```
+
+### Additional Information and Tips
+
+- To use the Short-Term Memory module effectively, consider the following tips:
+  - Use the `add` function to store messages in short-term memory.
+  -
+
+ Retrieve memory contents using `get_short_term` and `get_medium_term` functions.
+  - Clear memory as needed using `clear` and `clear_medium_term` functions.
+  - Search for specific terms within the memory using the `search_memory` function.
+  - Save and load memory data to/from files using `save_to_file` and `load_from_file` functions.
+
+- Ensure proper exception handling when using memory functions to handle potential errors gracefully.
+
+- When using the `search_memory` function, iterate through the results dictionary to access search results for short-term and medium-term memory.
+
+### References and Resources
+
+- For more information on multi-agent systems and memory management, refer to the SWARMS framework documentation: [SWARMS Documentation](https://swarms.apac.ai/).
+
+- For advanced memory management and customization, explore the SWARMS framework source code.
+

mkdocs.yml

@@ -107,6 +107,7 @@ nav:
     - Weaviate: "swarms/memory/weaviate.md"
     - PineconDB: "swarms/memory/pinecone.md"
     - PGVectorStore: "swarms/memory/pg.md"
+    - ShortTermMemory: "swarms/memory/short_term_memory.md"
   - swarms.utils:
     - phoenix_trace_decorator: "swarms/utils/phoenix_tracer.md"
 - Guides:

swarms/__init__.py

@@ -9,4 +9,5 @@ from swarms.models import *  # noqa: E402, F403
 from swarms.telemetry import *  # noqa: E402, F403
 from swarms.utils import *  # noqa: E402, F403
 from swarms.prompts import *  # noqa: E402, F403
+
 # from swarms.cli import *  # noqa: E402, F403

swarms/memory/__init__.py

@@ -1,5 +1,7 @@
 from swarms.memory.base_vectordb import VectorDatabase
+from swarms.memory.short_term_memory import ShortTermMemory
 
 __all__ = [
     "VectorDatabase",
+    "ShortTermMemory"
 ]

swarms/memory/short_term_memory.py

@@ -0,0 +1,166 @@
+import logging
+from swarms.structs.base import BaseStructure
+import threading
+import json
+import os
+
+
+class ShortTermMemory(BaseStructure):
+    def __init__(
+        self,
+        return_str: bool = True,
+        autosave: bool = True,
+        *args,
+        **kwargs,
+    ):
+        self.return_str = return_str
+        self.autosave = autosave
+        self.short_term_memory = []
+        self.medium_term_memory = []
+        self.lock = threading.Lock()
+
+    def add(
+        self, role: str = None, message: str = None, *args, **kwargs
+    ):
+        """Add a message to the short term memory.
+
+        Args:
+            role (str, optional): _description_. Defaults to None.
+            message (str, optional): _description_. Defaults to None.
+
+        Returns:
+            _type_: _description_
+        """
+        try:
+            memory = self.short_term_memory.append(
+                {"role": role, "message": message}
+            )
+
+            return memory
+        except Exception as error:
+            print(f"Add to short term memory failed: {error}")
+            raise error
+
+    def get_short_term(self):
+        """Get the short term memory.
+
+        Returns:
+            _type_: _description_
+        """
+        return self.short_term_memory
+
+    def get_medium_term(self):
+        """Get the medium term memory.
+
+        Returns:
+            _type_: _description_
+        """
+        return self.medium_term_memory
+
+    def clear_medium_term(self):
+        """Clear the medium term memory."""
+        self.medium_term_memory = []
+
+    def get_short_term_memory_str(self, *args, **kwargs):
+        """Get the short term memory as a string."""
+        return str(self.short_term_memory)
+
+    def update_short_term(
+        self, index, role: str, message: str, *args, **kwargs
+    ):
+        self.short_term_memory[index] = {
+            "role": role,
+            "message": message,
+        }
+
+    def clear(self):
+        """Clear the short term memory."""
+        self.short_term_memory = []
+
+    def search_memory(self, term):
+        """Search the memory for a term.
+
+        Args:
+            term (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        results = {"short_term": [], "medium_term": []}
+        for i, message in enumerate(self.short_term_memory):
+            if term in message["message"]:
+                results["short_term"].append((i, message))
+        for i, message in enumerate(self.medium_term_memory):
+            if term in message["message"]:
+                results["medium_term"].append((i, message))
+        return results
+
+    def return_shortmemory_as_str(self):
+        """Return the memory as a string.
+
+        Returns:
+            _type_: _description_
+        """
+        return str(self.short_term_memory)
+
+    def move_to_medium_term(self, index):
+        """Move a message from the short term memory to the medium term memory.
+
+        Args:
+            index (_type_): _description_
+        """
+        message = self.short_term_memory.pop(index)
+        self.medium_term_memory.append(message)
+
+    def return_medium_memory_as_str(self):
+        """Return the medium term memory as a string.
+
+        Returns:
+            _type_: _description_
+        """
+        return str(self.medium_term_memory)
+
+    def save_to_file(self, filename: str):
+        """Save the memory to a file.
+
+        Args:
+            filename (str): _description_
+        """
+        try:
+            with self.lock:
+                with open(filename, "w") as f:
+                    json.dump(
+                        {
+                            "short_term_memory": (
+                                self.short_term_memory
+                            ),
+                            "medium_term_memory": (
+                                self.medium_term_memory
+                            ),
+                        },
+                        f,
+                    )
+
+                    logging.info(f"Saved memory to {filename}")
+        except Exception as error:
+            print(f"Error saving memory to {filename}: {error}")
+
+    def load_from_file(self, filename: str, *args, **kwargs):
+        """Load the memory from a file.
+
+        Args:
+            filename (str): _description_
+        """
+        try:
+            with self.lock:
+                with open(filename, "r") as f:
+                    data = json.load(f)
+                self.short_term_memory = data.get(
+                    "short_term_memory", []
+                )
+                self.medium_term_memory = data.get(
+                    "medium_term_memory", []
+                )
+                logging.info(f"Loaded memory from {filename}")
+        except Exception as error:
+            print(f"Erorr loading memory from {filename}: {error}")

swarms/structs/base.py

@@ -79,7 +79,6 @@ class BaseStructure(ABC):
         self.save_metadata_path = save_metadata_path
         self.save_error_path = save_error_path
 
-    @abstractmethod
     def run(self, *args, **kwargs):
         """Run the structure."""
         pass

tests/memory/test_short_term_memory.py

@@ -0,0 +1,128 @@
+import pytest
+from swarms.memory.short_term_memory import ShortTermMemory
+import threading
+
+def test_init():
+    memory = ShortTermMemory()
+    assert memory.short_term_memory == []
+    assert memory.medium_term_memory == []
+
+
+def test_add():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    assert memory.short_term_memory == [
+        {"role": "user", "message": "Hello, world!"}
+    ]
+
+
+def test_get_short_term():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    assert memory.get_short_term() == [
+        {"role": "user", "message": "Hello, world!"}
+    ]
+
+
+def test_get_medium_term():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    memory.move_to_medium_term(0)
+    assert memory.get_medium_term() == [
+        {"role": "user", "message": "Hello, world!"}
+    ]
+
+
+def test_clear_medium_term():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    memory.move_to_medium_term(0)
+    memory.clear_medium_term()
+    assert memory.get_medium_term() == []
+
+
+def test_get_short_term_memory_str():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    assert (
+        memory.get_short_term_memory_str()
+        == "[{'role': 'user', 'message': 'Hello, world!'}]"
+    )
+
+
+def test_update_short_term():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    memory.update_short_term(0, "user", "Goodbye, world!")
+    assert memory.get_short_term() == [
+        {"role": "user", "message": "Goodbye, world!"}
+    ]
+
+
+def test_clear():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    memory.clear()
+    assert memory.get_short_term() == []
+
+
+def test_search_memory():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    assert memory.search_memory("Hello") == {
+        "short_term": [
+            (0, {"role": "user", "message": "Hello, world!"})
+        ],
+        "medium_term": [],
+    }
+
+
+def test_return_shortmemory_as_str():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    assert (
+        memory.return_shortmemory_as_str()
+        == "[{'role': 'user', 'message': 'Hello, world!'}]"
+    )
+
+
+def test_move_to_medium_term():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    memory.move_to_medium_term(0)
+    assert memory.get_medium_term() == [
+        {"role": "user", "message": "Hello, world!"}
+    ]
+    assert memory.get_short_term() == []
+
+
+def test_return_medium_memory_as_str():
+    memory = ShortTermMemory()
+    memory.add("user", "Hello, world!")
+    memory.move_to_medium_term(0)
+    assert (
+        memory.return_medium_memory_as_str()
+        == "[{'role': 'user', 'message': 'Hello, world!'}]"
+    )
+
+
+def test_thread_safety():
+    memory = ShortTermMemory()
+    def add_messages():
+        for _ in range(1000):
+            memory.add("user", "Hello, world!")
+    threads = [threading.Thread(target=add_messages) for _ in range(10)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+    assert len(memory.get_short_term()) == 10000
+
+def test_save_and_load():
+    memory1 = ShortTermMemory()
+    memory1.add("user", "Hello, world!")
+    memory1.save_to_file("memory.json")
+    memory2 = ShortTermMemory()
+    memory2.load_from_file("memory.json")
+    assert memory1.get_short_term() == memory2.get_short_term()
+    assert memory1.get_medium_term() == memory2.get_medium_term()