Fix #2426: Add support for Redis as vector store in mem0 for user_memory

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-03-20 13:28:02 +00:00
parent fe0813e831
commit 05c66405cf
3 changed files with 150 additions and 7 deletions

View File

@@ -9,6 +9,29 @@ from crewai.memory.storage.interface import Storage
class Mem0Storage(Storage): class Mem0Storage(Storage):
""" """
Extends Storage to handle embedding and searching across entities using Mem0. 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): def __init__(self, type, crew=None):
@@ -26,19 +49,49 @@ class Mem0Storage(Storage):
if type == "user" and not user_id: if type == "user" and not user_id:
raise ValueError("User ID is required for user memory type") 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", {}) config = self.memory_config.get("config", {})
mem0_api_key = config.get("api_key") or os.getenv("MEM0_API_KEY") mem0_api_key = config.get("api_key") or os.getenv("MEM0_API_KEY")
mem0_org_id = config.get("org_id") mem0_org_id = config.get("org_id")
mem0_project_id = config.get("project_id") mem0_project_id = config.get("project_id")
vector_store_config = config.get("vector_store")
# Initialize MemoryClient with available parameters # If vector store configuration is provided, use Memory.from_config
if mem0_org_id and mem0_project_id: if vector_store_config:
self.memory = MemoryClient( try:
api_key=mem0_api_key, org_id=mem0_org_id, project_id=mem0_project_id 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: 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: def _sanitize_role(self, role: str) -> str:
""" """

View File

@@ -9,6 +9,29 @@ class UserMemory(Memory):
Inherits from the Memory class and utilizes an instance of a class that Inherits from the Memory class and utilizes an instance of a class that
adheres to the Storage for data storage, specifically working with adheres to the Storage for data storage, specifically working with
MemoryItem instances. 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): def __init__(self, crew=None):

View File

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