From 05c66405cfbaa3d14c2dbe250d235e1cc64fac5a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 20 Mar 2025 13:28:02 +0000 Subject: [PATCH] Fix #2426: Add support for Redis as vector store in mem0 for user_memory Co-Authored-By: Joe Moura --- src/crewai/memory/storage/mem0_storage.py | 67 ++++++++++++++++++++--- src/crewai/memory/user/user_memory.py | 23 ++++++++ tests/memory/test_redis_mem0_storage.py | 67 +++++++++++++++++++++++ 3 files changed, 150 insertions(+), 7 deletions(-) create mode 100644 tests/memory/test_redis_mem0_storage.py diff --git a/src/crewai/memory/storage/mem0_storage.py b/src/crewai/memory/storage/mem0_storage.py index be889afff..4a65e12ed 100644 --- a/src/crewai/memory/storage/mem0_storage.py +++ b/src/crewai/memory/storage/mem0_storage.py @@ -9,6 +9,29 @@ from crewai.memory.storage.interface import Storage class Mem0Storage(Storage): """ Extends Storage to handle embedding and searching across entities using Mem0. + + Supports configuring Redis as a vector store through the memory_config: + + ```python + crew = Crew( + memory=True, + memory_config={ + "provider": "mem0", + "config": { + "user_id": "test-user", + "api_key": "mem0-api-key", + "vector_store": { + "provider": "redis", + "config": { + "collection_name": "collection_name", + "embedding_model_dims": 1536, + "redis_url": "redis://redis-host:6379/0" + } + } + } + } + ) + ``` """ def __init__(self, type, crew=None): @@ -26,19 +49,49 @@ class Mem0Storage(Storage): if type == "user" and not user_id: raise ValueError("User ID is required for user memory type") - # API key in memory config overrides the environment variable + # Get configuration from memory_config config = self.memory_config.get("config", {}) 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") + vector_store_config = config.get("vector_store") - # Initialize MemoryClient with available parameters - if mem0_org_id and mem0_project_id: - self.memory = MemoryClient( - api_key=mem0_api_key, org_id=mem0_org_id, project_id=mem0_project_id - ) + # If vector store configuration is provided, use Memory.from_config + if vector_store_config: + try: + from mem0.memory.main import Memory + + # Prepare memory config with vector store configuration + memory_config = { + "vector_store": vector_store_config + } + + # Add API key if provided + if mem0_api_key: + memory_config["api_key"] = mem0_api_key + + # Add org_id and project_id if provided + if mem0_org_id: + memory_config["org_id"] = mem0_org_id + if mem0_project_id: + memory_config["project_id"] = mem0_project_id + + # Initialize Memory with configuration + self.memory = Memory.from_config(memory_config) + except ImportError: + raise ImportError( + "Mem0 is not installed. Please install it with `pip install mem0ai`." + ) + except Exception as e: + raise ValueError(f"Failed to initialize Memory with vector store configuration: {e}") else: - self.memory = MemoryClient(api_key=mem0_api_key) + # Fall back to default MemoryClient initialization + if mem0_org_id and mem0_project_id: + self.memory = MemoryClient( + api_key=mem0_api_key, org_id=mem0_org_id, project_id=mem0_project_id + ) + else: + self.memory = MemoryClient(api_key=mem0_api_key) def _sanitize_role(self, role: str) -> str: """ diff --git a/src/crewai/memory/user/user_memory.py b/src/crewai/memory/user/user_memory.py index 24e5fe035..051199ad3 100644 --- a/src/crewai/memory/user/user_memory.py +++ b/src/crewai/memory/user/user_memory.py @@ -9,6 +9,29 @@ class UserMemory(Memory): Inherits from the Memory class and utilizes an instance of a class that adheres to the Storage for data storage, specifically working with MemoryItem instances. + + To configure with Redis as a vector store, provide a memory_config to the Crew: + + ```python + crew = Crew( + memory=True, + memory_config={ + "provider": "mem0", + "config": { + "user_id": "test-user", + "api_key": "mem0-api-key", + "vector_store": { + "provider": "redis", + "config": { + "collection_name": "collection_name", + "embedding_model_dims": 1536, + "redis_url": "redis://redis-host:6379/0" + } + } + } + } + ) + ``` """ def __init__(self, crew=None): diff --git a/tests/memory/test_redis_mem0_storage.py b/tests/memory/test_redis_mem0_storage.py new file mode 100644 index 000000000..6fef19d1e --- /dev/null +++ b/tests/memory/test_redis_mem0_storage.py @@ -0,0 +1,67 @@ +import os +import pytest +from unittest.mock import patch, MagicMock + +from crewai.memory.storage.mem0_storage import Mem0Storage + + +class TestMem0RedisIntegration: + @pytest.fixture + def mock_memory(self): + with patch("mem0.memory.main.Memory") as mock_memory: + mock_memory_instance = MagicMock() + mock_memory.from_config.return_value = mock_memory_instance + yield mock_memory + + def test_mem0_with_redis_config(self, mock_memory): + # Create a mock crew with Redis vector store configuration + mock_crew = MagicMock() + mock_crew.memory_config = { + "provider": "mem0", + "config": { + "user_id": "test-user", + "api_key": "test-api-key", + "vector_store": { + "provider": "redis", + "config": { + "collection_name": "test_collection", + "embedding_model_dims": 1536, + "redis_url": "redis://localhost:6379/0" + } + } + } + } + + # Create Mem0Storage instance + with patch("crewai.memory.storage.mem0_storage.MemoryClient"): + storage = Mem0Storage(type="user", crew=mock_crew) + + # Check that Memory.from_config was called with correct parameters + mock_memory.from_config.assert_called_once() + config_arg = mock_memory.from_config.call_args[0][0] + assert "vector_store" in config_arg + assert config_arg["vector_store"]["provider"] == "redis" + assert config_arg["vector_store"]["config"]["redis_url"] == "redis://localhost:6379/0" + + def test_fallback_to_memory_client(self): + # Create a mock crew without vector store configuration + mock_crew = MagicMock() + mock_crew.memory_config = { + "provider": "mem0", + "config": { + "user_id": "test-user", + "api_key": "test-api-key" + } + } + + # Mock MemoryClient + with patch("crewai.memory.storage.mem0_storage.MemoryClient") as mock_client: + mock_client_instance = MagicMock() + mock_client.return_value = mock_client_instance + + # Create Mem0Storage instance + storage = Mem0Storage(type="user", crew=mock_crew) + + # Check that MemoryClient was called (fallback path) + mock_client.assert_called_once() + assert mock_client.call_args[1]["api_key"] == "test-api-key"