diff --git a/docs/concepts/memory.mdx b/docs/concepts/memory.mdx index bb4e885cd..f3f1812c2 100644 --- a/docs/concepts/memory.mdx +++ b/docs/concepts/memory.mdx @@ -164,7 +164,10 @@ crew = Crew( [Mem0](https://mem0.ai/) is a self-improving memory layer for LLM applications, enabling personalized AI experiences. -To include user-specific memory you can get your API key [here](https://app.mem0.ai/dashboard/api-keys) and refer the [docs](https://docs.mem0.ai/platform/quickstart#4-1-create-memories) for adding user preferences. + +### Using Mem0 API platform + +To include user-specific memory you can get your API key [here](https://app.mem0.ai/dashboard/api-keys) and refer the [docs](https://docs.mem0.ai/platform/quickstart#4-1-create-memories) for adding user preferences. In this case `user_memory` is set to `MemoryClient` from mem0. ```python Code @@ -175,18 +178,7 @@ from mem0 import MemoryClient # Set environment variables for Mem0 os.environ["MEM0_API_KEY"] = "m0-xx" -# Step 1: Record preferences based on past conversation or user input -client = MemoryClient() -messages = [ - {"role": "user", "content": "Hi there! I'm planning a vacation and could use some advice."}, - {"role": "assistant", "content": "Hello! I'd be happy to help with your vacation planning. What kind of destination do you prefer?"}, - {"role": "user", "content": "I am more of a beach person than a mountain person."}, - {"role": "assistant", "content": "That's interesting. Do you like hotels or Airbnb?"}, - {"role": "user", "content": "I like Airbnb more."}, -] -client.add(messages, user_id="john") - -# Step 2: Create a Crew with User Memory +# Step 1: Create a Crew with User Memory crew = Crew( agents=[...], @@ -197,11 +189,12 @@ crew = Crew( memory_config={ "provider": "mem0", "config": {"user_id": "john"}, + "user_memory" : {} #Set user_memory explicitly to a dictionary, we are working on this issue. }, ) ``` -## Memory Configuration Options +#### Additional Memory Configuration Options If you want to access a specific organization and project, you can set the `org_id` and `project_id` parameters in the memory configuration. ```python Code @@ -215,10 +208,74 @@ crew = Crew( memory_config={ "provider": "mem0", "config": {"user_id": "john", "org_id": "my_org_id", "project_id": "my_project_id"}, + "user_memory" : {} #Set user_memory explicitly to a dictionary, we are working on this issue. }, ) ``` +### Using Local Mem0 memory +If you want to use local mem0 memory, with a custom configuration, you can set a parameter `local_mem0_config` in the config itself. +If both os environment key is set and local_mem0_config is given, the API platform takes higher priority over the local configuration. +Check [this](https://docs.mem0.ai/open-source/python-quickstart#run-mem0-locally) mem0 local configuration docs for more understanding. +In this case `user_memory` is set to `Memory` from mem0. + + +```python Code +from crewai import Crew + + +#local mem0 config +config = { + "vector_store": { + "provider": "qdrant", + "config": { + "host": "localhost", + "port": 6333 + } + }, + "llm": { + "provider": "openai", + "config": { + "api_key": "your-api-key", + "model": "gpt-4" + } + }, + "embedder": { + "provider": "openai", + "config": { + "api_key": "your-api-key", + "model": "text-embedding-3-small" + } + }, + "graph_store": { + "provider": "neo4j", + "config": { + "url": "neo4j+s://your-instance", + "username": "neo4j", + "password": "password" + } + }, + "history_db_path": "/path/to/history.db", + "version": "v1.1", + "custom_fact_extraction_prompt": "Optional custom prompt for fact extraction for memory", + "custom_update_memory_prompt": "Optional custom prompt for update memory" +} + +crew = Crew( + agents=[...], + tasks=[...], + verbose=True, + memory=True, + memory_config={ + "provider": "mem0", + "config": {"user_id": "john", 'local_mem0_config': config}, + "user_memory" : {} #Set user_memory explicitly to a dictionary, we are working on this issue. + }, +) +``` + + + ## Additional Embedding Providers ### Using OpenAI embeddings (already default) diff --git a/src/crewai/crew.py b/src/crewai/crew.py index 0f6db8c4a..60f7c5677 100644 --- a/src/crewai/crew.py +++ b/src/crewai/crew.py @@ -290,24 +290,25 @@ class Crew(BaseModel): else EntityMemory(crew=self, embedder_config=self.embedder) ) if ( - self.memory_config and "user_memory" in self.memory_config + self.memory_config and "user_memory" in self.memory_config and self.memory_config.get('provider') == 'mem0' ): # Check for user_memory in config user_memory_config = self.memory_config["user_memory"] if isinstance( - user_memory_config, UserMemory - ): # Check if it is already an instance - self._user_memory = user_memory_config - elif isinstance( user_memory_config, dict ): # Check if it's a configuration dict self._user_memory = UserMemory( - crew=self, **user_memory_config - ) # Initialize with config + crew=self + ) else: raise TypeError( - "user_memory must be a UserMemory instance or a configuration dictionary" + "user_memory must be a configuration dictionary" ) else: + self._logger.log( + "warning", + "User memory initialization failed. For setup instructions, please refer to the memory documentation: https://docs.crewai.com/concepts/memory#integrating-mem0-for-enhanced-user-memory", + color="yellow" + ) self._user_memory = None # No user memory if not in config return self diff --git a/src/crewai/memory/contextual/contextual_memory.py b/src/crewai/memory/contextual/contextual_memory.py index cdb9cf836..a9f657d8a 100644 --- a/src/crewai/memory/contextual/contextual_memory.py +++ b/src/crewai/memory/contextual/contextual_memory.py @@ -94,6 +94,10 @@ class ContextualMemory: Returns: str: Formatted user memories as bullet points, or an empty string if none found. """ + + if self.um is None: + return "" + user_memories = self.um.search(query) if not user_memories: return "" diff --git a/src/crewai/memory/storage/mem0_storage.py b/src/crewai/memory/storage/mem0_storage.py index 0319c6a8a..6c9ffc682 100644 --- a/src/crewai/memory/storage/mem0_storage.py +++ b/src/crewai/memory/storage/mem0_storage.py @@ -31,6 +31,7 @@ class Mem0Storage(Storage): mem0_api_key = config.get("api_key") or os.getenv("MEM0_API_KEY") mem0_org_id = config.get("org_id") mem0_project_id = config.get("project_id") + mem0_local_config = config.get("local_mem0_config") # Initialize MemoryClient or Memory based on the presence of the mem0_api_key if mem0_api_key: @@ -41,7 +42,10 @@ class Mem0Storage(Storage): else: self.memory = MemoryClient(api_key=mem0_api_key) else: - self.memory = Memory() # Fallback to Memory if no Mem0 API key is provided + if mem0_local_config and len(mem0_local_config): + self.memory = Memory.from_config(config) + else: + self.memory = Memory() def _sanitize_role(self, role: str) -> str: """ @@ -114,3 +118,7 @@ class Mem0Storage(Storage): agents = [self._sanitize_role(agent.role) for agent in agents] agents = "_".join(agents) return agents + + def reset(self): + if self.memory: + self.memory.reset() diff --git a/src/crewai/memory/user/user_memory.py b/src/crewai/memory/user/user_memory.py index 24e5fe035..1c710bb69 100644 --- a/src/crewai/memory/user/user_memory.py +++ b/src/crewai/memory/user/user_memory.py @@ -43,3 +43,11 @@ class UserMemory(Memory): score_threshold=score_threshold, ) return results + + def reset(self) -> None: + try: + self.storage.reset() + except Exception as e: + raise Exception( + f"An error occurred while resetting the user memory: {e}" + ) diff --git a/tests/memory/user_memory_test.py b/tests/memory/user_memory_test.py new file mode 100644 index 000000000..53514f638 --- /dev/null +++ b/tests/memory/user_memory_test.py @@ -0,0 +1,68 @@ + +from unittest.mock import MagicMock, patch + +import pytest +from mem0.memory.main import Memory + +from crewai.memory.user.user_memory import UserMemory +from crewai.memory.user.user_memory_item import UserMemoryItem + + +class MockCrew: + def __init__(self, memory_config): + self.memory_config = memory_config + +@pytest.fixture +def user_memory(): + """Fixture to create a UserMemory instance""" + crew = MockCrew( + memory_config={ + "provider": "mem0", + "config": {"user_id": "john"}, + "user_memory" : {} + } + ) + + user_memory = MagicMock(spec=UserMemory) + + with patch.object(Memory,'__new__',return_value=user_memory): + user_memory_instance = UserMemory(crew=crew) + + return user_memory_instance + +def test_save_and_search(user_memory): + memory = UserMemoryItem( + data="""test value test value test value test value test value test value + test value test value test value test value test value test value + test value test value test value test value test value test value""", + user="test_user", + metadata={"task": "test_task"}, + ) + + with patch.object(UserMemory, "save") as mock_save: + user_memory.save( + value=memory.data, + metadata=memory.metadata, + user=memory.user + ) + + mock_save.assert_called_once_with( + value=memory.data, + metadata=memory.metadata, + user=memory.user + ) + + expected_result = [ + { + "context": memory.data, + "metadata": {"agent": "test_agent"}, + "score": 0.95, + } + ] + expected_result = ["mocked_result"] + + # Use patch.object to mock UserMemory's search method + with patch.object(UserMemory, 'search', return_value=expected_result) as mock_search: + find = UserMemory.search("test value", score_threshold=0.01)[0] + mock_search.assert_called_once_with("test value", score_threshold=0.01) + assert find == expected_result[0] \ No newline at end of file diff --git a/tests/storage/test_mem0_storage.py b/tests/storage/test_mem0_storage.py new file mode 100644 index 000000000..e3d68092a --- /dev/null +++ b/tests/storage/test_mem0_storage.py @@ -0,0 +1,114 @@ +import os +from unittest.mock import MagicMock, patch + +import pytest +from mem0.client.main import MemoryClient +from mem0.memory.main import Memory + +from crewai.agent import Agent +from crewai.crew import Crew +from crewai.memory.storage.mem0_storage import Mem0Storage +from crewai.task import Task + + +# Define the class (if not already defined) +class MockCrew: + def __init__(self, memory_config): + self.memory_config = memory_config + + +@pytest.fixture +def mock_mem0_memory(): + """Fixture to create a mock Memory instance""" + mock_memory = MagicMock(spec=Memory) + return mock_memory + + +@pytest.fixture +def mem0_storage_with_mocked_config(mock_mem0_memory): + """Fixture to create a Mem0Storage instance with mocked dependencies""" + + # Patch the Memory class to return our mock + with patch('mem0.memory.main.Memory.from_config', return_value=mock_mem0_memory): + config = { + "vector_store": { + "provider": "mock_vector_store", + "config": { + "host": "localhost", + "port": 6333 + } + }, + "llm": { + "provider": "mock_llm", + "config": { + "api_key": "mock-api-key", + "model": "mock-model" + } + }, + "embedder": { + "provider": "mock_embedder", + "config": { + "api_key": "mock-api-key", + "model": "mock-model" + } + }, + "graph_store": { + "provider": "mock_graph_store", + "config": { + "url": "mock-url", + "username": "mock-user", + "password": "mock-password" + } + }, + "history_db_path": "/mock/path", + "version": "test-version", + "custom_fact_extraction_prompt": "mock prompt 1", + "custom_update_memory_prompt": "mock prompt 2" + } + + # Instantiate the class with memory_config + crew = MockCrew( + memory_config={ + "provider": "mem0", + "config": {"user_id": "test_user", "local_mem0_config": config}, + } + ) + + mem0_storage = Mem0Storage(type="short_term", crew=crew) + return mem0_storage + + +def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_memory): + """Test that Mem0Storage initializes correctly with the mocked config""" + assert mem0_storage_with_mocked_config.memory_type == "short_term" + assert mem0_storage_with_mocked_config.memory is mock_mem0_memory + + +@pytest.fixture +def mock_mem0_memory_client(): + """Fixture to create a mock MemoryClient instance""" + mock_memory = MagicMock(spec=MemoryClient) + return mock_memory + + +@pytest.fixture +def mem0_storage_with_memory_client(mock_mem0_memory_client): + """Fixture to create a Mem0Storage instance with mocked dependencies""" + + # We need to patch the MemoryClient before it's instantiated + with patch.object(MemoryClient, '__new__', return_value=mock_mem0_memory_client): + crew = MockCrew( + memory_config={ + "provider": "mem0", + "config": {"user_id": "test_user", "api_key": "ABCDEFGH", "org_id": "my_org_id", "project_id": "my_project_id"}, + } + ) + + mem0_storage = Mem0Storage(type="short_term", crew=crew) + return mem0_storage + + +def test_mem0_storage_with_memory_client_initialization(mem0_storage_with_memory_client, mock_mem0_memory_client): + """Test Mem0Storage initialization with MemoryClient""" + assert mem0_storage_with_memory_client.memory_type == "short_term" + assert mem0_storage_with_memory_client.memory is mock_mem0_memory_client