Add support for custom memory storage implementations (fixes #2278)

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

tests/memory/test_custom_storage.py

@@ -0,0 +1,140 @@
+import pytest
+from typing import Any, Dict, List
+
+from crewai.crew import Crew
+from crewai.agent import Agent
+from crewai.memory.storage.interface import Storage
+from crewai.memory.short_term.short_term_memory import ShortTermMemory
+from crewai.memory.long_term.long_term_memory import LongTermMemory
+from crewai.memory.entity.entity_memory import EntityMemory
+from crewai.memory.user.user_memory import UserMemory
+
+
+class CustomStorage(Storage):
+    """Custom storage implementation for testing."""
+
+    def __init__(self):
+        self.data = []
+
+    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+        self.data.append({"value": value, "metadata": metadata})
+
+    def search(
+        self, query: str, limit: int = 3, score_threshold: float = 0.35
+    ) -> List[Any]:
+        return [{"context": item["value"], "metadata": item["metadata"]} for item in self.data]
+
+    def reset(self) -> None:
+        self.data = []
+
+
+def test_custom_storage_with_short_term_memory():
+    """Test that custom storage works with short term memory."""
+    custom_storage = CustomStorage()
+    memory = ShortTermMemory(storage=custom_storage)
+    
+    memory.save("test value", {"key": "value"})
+    results = memory.search("test")
+    
+    assert len(results) > 0
+    assert results[0]["context"] == "test value"
+    assert results[0]["metadata"]["key"] == "value"
+
+
+def test_custom_storage_with_long_term_memory():
+    """Test that custom storage works with long term memory."""
+    custom_storage = CustomStorage()
+    memory = LongTermMemory(storage=custom_storage)
+    
+    memory.save("test value", {"key": "value"})
+    results = memory.search("test")
+    
+    assert len(results) > 0
+    assert results[0]["context"] == "test value"
+    assert results[0]["metadata"]["key"] == "value"
+
+
+def test_custom_storage_with_entity_memory():
+    """Test that custom storage works with entity memory."""
+    custom_storage = CustomStorage()
+    memory = EntityMemory(storage=custom_storage)
+    
+    memory.save("test value", {"key": "value"})
+    results = memory.search("test")
+    
+    assert len(results) > 0
+    assert results[0]["context"] == "test value"
+    assert results[0]["metadata"]["key"] == "value"
+
+
+def test_custom_storage_with_user_memory():
+    """Test that custom storage works with user memory."""
+    custom_storage = CustomStorage()
+    memory = UserMemory(storage=custom_storage)
+    
+    memory.save("test value", {"key": "value"})
+    results = memory.search("test")
+    
+    assert len(results) > 0
+    # UserMemory prepends "Remember the details about the user: " to the value
+    assert "test value" in results[0]["context"]
+    assert results[0]["metadata"]["key"] == "value"
+
+
+def test_custom_storage_with_crew():
+    """Test that custom storage works with crew."""
+    short_term_storage = CustomStorage()
+    long_term_storage = CustomStorage()
+    entity_storage = CustomStorage()
+    user_storage = CustomStorage()
+    
+    # Create memory instances with custom storage
+    short_term_memory = ShortTermMemory(storage=short_term_storage)
+    long_term_memory = LongTermMemory(storage=long_term_storage)
+    entity_memory = EntityMemory(storage=entity_storage)
+    user_memory = UserMemory(storage=user_storage)
+    
+    # Create a crew with custom memory instances
+    crew = Crew(
+        agents=[Agent(role="test", goal="test", backstory="test")],
+        memory=True,
+        short_term_memory=short_term_memory,
+        long_term_memory=long_term_memory,
+        entity_memory=entity_memory,
+        memory_config={"user_memory": user_memory},
+    )
+    
+    # Test that the crew has the custom memory instances
+    assert crew._short_term_memory.storage == short_term_storage
+    assert crew._long_term_memory.storage == long_term_storage
+    assert crew._entity_memory.storage == entity_storage
+    assert crew._user_memory.storage == user_storage
+
+
+def test_custom_storage_with_memory_config():
+    """Test that custom storage works with memory_config."""
+    short_term_storage = CustomStorage()
+    long_term_storage = CustomStorage()
+    entity_storage = CustomStorage()
+    user_storage = CustomStorage()
+    
+    # Create a crew with custom storage in memory_config
+    crew = Crew(
+        agents=[Agent(role="test", goal="test", backstory="test")],
+        memory=True,
+        memory_config={
+            "storage": {
+                "short_term": short_term_storage,
+                "long_term": long_term_storage,
+                "entity": entity_storage,
+                "user": user_storage,
+            },
+            "user_memory": {}  # Enable user memory
+        },
+    )
+    
+    # Test that the crew has the custom storage instances
+    assert crew._short_term_memory.storage == short_term_storage
+    assert crew._long_term_memory.storage == long_term_storage
+    assert crew._entity_memory.storage == entity_storage
+    assert crew._user_memory.storage == user_storage