ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[ShortTermMemory]

Browse Source
pull/299/head
Kye 1 year ago
parent e3f76b998a
commit d48cc86f32
7 changed files with 547 additions and 2 deletions
Show Stats Download Patch File Download Diff File

248
docs/swarms/memory/short_term_memory.md
Unescape View File

@ -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.

1
mkdocs.yml
Unescape View File

@ -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:

3
swarms/__init__.py
Unescape View File

@ -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
# from swarms.cli import * # noqa: E402, F403

2
swarms/memory/__init__.py
Unescape View File

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

166
swarms/memory/short_term_memory.py
Unescape View File

@ -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}")

1
swarms/structs/base.py
Unescape View File

@ -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

128
tests/memory/test_short_term_memory.py
Unescape View File

@ -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()
Write Preview
Loading…
Cancel
Save