From d48cc86f3294a0c03802e4187329c9e9d9ce586f Mon Sep 17 00:00:00 2001 From: Kye Date: Thu, 14 Dec 2023 14:11:57 -0800 Subject: [PATCH] [ShortTermMemory] --- docs/swarms/memory/short_term_memory.md | 248 ++++++++++++++++++++++++ mkdocs.yml | 1 + swarms/__init__.py | 3 +- swarms/memory/__init__.py | 2 + swarms/memory/short_term_memory.py | 166 ++++++++++++++++ swarms/structs/base.py | 1 - tests/memory/test_short_term_memory.py | 128 ++++++++++++ 7 files changed, 547 insertions(+), 2 deletions(-) create mode 100644 docs/swarms/memory/short_term_memory.md create mode 100644 swarms/memory/short_term_memory.py create mode 100644 tests/memory/test_short_term_memory.py diff --git a/docs/swarms/memory/short_term_memory.md b/docs/swarms/memory/short_term_memory.md new file mode 100644 index 00000000..2aabbd5c --- /dev/null +++ b/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. + diff --git a/mkdocs.yml b/mkdocs.yml index 2d200d91..b0e34ca7 100644 --- a/mkdocs.yml +++ b/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: diff --git a/swarms/__init__.py b/swarms/__init__.py index 5ee363d8..f6f04205 100644 --- a/swarms/__init__.py +++ b/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 \ No newline at end of file + +# from swarms.cli import * # noqa: E402, F403 diff --git a/swarms/memory/__init__.py b/swarms/memory/__init__.py index 66639678..086bd1bd 100644 --- a/swarms/memory/__init__.py +++ b/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" ] diff --git a/swarms/memory/short_term_memory.py b/swarms/memory/short_term_memory.py new file mode 100644 index 00000000..d343a80f --- /dev/null +++ b/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}") diff --git a/swarms/structs/base.py b/swarms/structs/base.py index e32b0b48..7d365b23 100644 --- a/swarms/structs/base.py +++ b/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 diff --git a/tests/memory/test_short_term_memory.py b/tests/memory/test_short_term_memory.py new file mode 100644 index 00000000..32d5d008 --- /dev/null +++ b/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() \ No newline at end of file