Fix lint issue: Properly organize imports in test_redis_mem0_storage.py

Co-Authored-By: Joe Moura <joao@crewai.com>

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": "your-user-id",
+                "api_key": os.getenv("MEM0_API_KEY"),  # Use environment variable
+                "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:
         """

src/crewai/memory/user/user_memory.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any, Dict, Optional
 
 from crewai.memory.memory import Memory
@@ -9,6 +10,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": "your-user-id",
+                "api_key": os.getenv("MEM0_API_KEY"),  # Use environment variable
+                "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):

tests/memory/test_redis_mem0_storage.py

@@ -0,0 +1,68 @@
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+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"