Improve type system and test coverage for custom memory storage

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

tests/memory/long_term_memory_test.py

@@ -7,7 +7,31 @@ from crewai.memory.long_term.long_term_memory_item import LongTermMemoryItem
 @pytest.fixture
 def long_term_memory():
     """Fixture to create a LongTermMemory instance"""
-    return LongTermMemory()
+    # Create a mock storage for testing
+    from crewai.memory.storage.interface import Storage
+    
+    class MockStorage(Storage):
+        def __init__(self):
+            self.data = []
+            
+        def save(self, value, metadata):
+            self.data.append({"value": value, "metadata": metadata})
+            
+        def search(self, query, limit=3, score_threshold=0.35):
+            return [
+                {
+                    "context": item["value"],
+                    "metadata": item["metadata"],
+                    "score": 0.5,
+                    "datetime": item["metadata"].get("datetime", "test_datetime")
+                }
+                for item in self.data
+            ]
+            
+        def reset(self):
+            self.data = []
+            
+    return LongTermMemory(storage=MockStorage())
 
 
 def test_save_and_search(long_term_memory):

tests/memory/short_term_memory_test.py

@@ -12,6 +12,8 @@ from crewai.task import Task
 @pytest.fixture
 def short_term_memory():
     """Fixture to create a ShortTermMemory instance"""
+    from crewai.memory.storage.rag_storage import RAGStorage
+    
     agent = Agent(
         role="Researcher",
         goal="Search relevant data and provide results",
@@ -25,7 +27,10 @@ def short_term_memory():
         expected_output="A list of relevant URLs based on the search query.",
         agent=agent,
     )
-    return ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
+    
+    storage = RAGStorage(type="short_term")
+    crew = Crew(agents=[agent], tasks=[task])
+    return ShortTermMemory(storage=storage, crew=crew)
 
 
 def test_save_and_search(short_term_memory):

tests/memory/test_custom_storage.py

@@ -7,7 +7,7 @@ from crewai.crew import Crew
 from crewai.memory.entity.entity_memory import EntityMemory
 from crewai.memory.long_term.long_term_memory import LongTermMemory
 from crewai.memory.short_term.short_term_memory import ShortTermMemory
-from crewai.memory.storage.interface import Storage
+from crewai.memory.storage.interface import Storage, SearchResult
 from crewai.memory.user.user_memory import UserMemory
 
 
@@ -22,8 +22,8 @@ class CustomStorage(Storage):
 
     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]
+    ) -> List[SearchResult]:
+        return [{"context": item["value"], "metadata": item["metadata"], "score": 0.9} for item in self.data]
 
     def reset(self) -> None:
         self.data = []
@@ -115,27 +115,97 @@ def test_custom_storage_with_crew():
 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()
+    long_term_memory = LongTermMemory(storage=CustomStorage())
+    entity_memory = EntityMemory(storage=CustomStorage())
+    user_memory = UserMemory(storage=CustomStorage())
     
     # Create a crew with custom storage in memory_config
     crew = Crew(
         agents=[Agent(role="test", goal="test", backstory="test")],
         memory=True,
+        short_term_memory=ShortTermMemory(storage=short_term_storage),
+        long_term_memory=long_term_memory,
+        entity_memory=entity_memory,
         memory_config={
-            "storage": {
-                "short_term": short_term_storage,
-                "long_term": long_term_storage,
-                "entity": entity_storage,
-                "user": user_storage,
-            },
-            "user_memory": {}  # Enable user memory
+            "user_memory": 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
+    assert crew._long_term_memory == long_term_memory
+    assert crew._entity_memory == entity_memory
+    assert crew._user_memory == user_memory
+
+
+def test_custom_storage_error_handling():
+    """Test error handling with custom storage."""
+    # Test exception propagation
+    class ErrorStorage(Storage):
+        """Storage implementation that raises exceptions."""
+        def __init__(self):
+            self.data = []
+            
+        def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+            raise ValueError("Save error")
+            
+        def search(
+            self, query: str, limit: int = 3, score_threshold: float = 0.35
+        ) -> List[SearchResult]:
+            raise ValueError("Search error")
+            
+        def reset(self) -> None:
+            raise ValueError("Reset error")
+            
+    storage = ErrorStorage()
+    memory = ShortTermMemory(storage=storage)
+    
+    with pytest.raises(ValueError, match="Save error"):
+        memory.save("test", {})
+        
+    with pytest.raises(ValueError, match="Search error"):
+        memory.search("test")
+        
+    with pytest.raises(Exception, match="An error occurred while resetting the short-term memory: Reset error"):
+        memory.reset()
+
+
+def test_custom_storage_edge_cases():
+    """Test edge cases with custom storage."""
+    class EdgeCaseStorage(Storage):
+        """Storage implementation for testing edge cases."""
+        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[SearchResult]:
+            return [{"context": item["value"], "metadata": item["metadata"], "score": 0.5} for item in self.data]
+            
+        def reset(self) -> None:
+            self.data = []
+            
+    storage = EdgeCaseStorage()
+    memory = ShortTermMemory(storage=storage)
+    
+    # Test empty query
+    memory.save("test value", {"key": "value"})
+    results = memory.search("")
+    assert len(results) > 0
+    
+    # Test very large metadata
+    large_metadata = {"key" + str(i): "value" * 100 for i in range(100)}
+    memory.save("test value", large_metadata)
+    results = memory.search("test")
+    assert len(results) > 0
+    assert results[1]["metadata"] == large_metadata
+    
+    # Test unicode and special characters
+    unicode_value = "测试值 with special chars: !@#$%^&*()"
+    memory.save(unicode_value, {"key": "value"})
+    results = memory.search("测试")
+    assert len(results) > 0
+    assert unicode_value in results[2]["context"]