From ad785baa16064b3ce1e3fc5aa90f11514a8ec360 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 4 Mar 2025 20:31:39 +0000
Subject: [PATCH] Fix syntax errors in memory storage implementation

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/memory/entity/entity_memory.py     | 54 ++++++++++---------
 src/crewai/memory/storage/base_rag_storage.py |  2 +-
 src/crewai/memory/storage/rag_storage.py      |  4 +-
 tests/memory/test_custom_storage.py           |  6 +--
 4 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/src/crewai/memory/entity/entity_memory.py b/src/crewai/memory/entity/entity_memory.py
index 698177261..28f9bd4ba 100644
--- a/src/crewai/memory/entity/entity_memory.py
+++ b/src/crewai/memory/entity/entity_memory.py
@@ -23,39 +23,41 @@ class EntityMemory(Memory):
         if crew and hasattr(crew, "memory_config") and crew.memory_config is not None:
             memory_config = crew.memory_config
             memory_provider = memory_config.get("provider")
-            
-        # Initialize with basic parameters
-        super().__init__(
-            storage=storage,
-            embedder_config=embedder_config,
-            memory_provider=memory_provider
-        )
         
-        try:
-            # Try to select storage using helper method
-            self.storage = self._select_storage(
-                storage=storage,
-                memory_config=memory_config,
-                storage_type="entity",
-                crew=crew,
-                path=path,
-                default_storage_factory=lambda path, crew: RAGStorage(
+        # If no storage is provided, try to create one
+        if storage is None:
+            try:
+                # Try to select storage using helper method
+                storage = self._select_storage(
+                    storage=storage,
+                    memory_config=memory_config,
+                    storage_type="entity",
+                    crew=crew,
+                    path=path,
+                    default_storage_factory=lambda path, crew: RAGStorage(
+                        type="entities",
+                        allow_reset=True,
+                        crew=crew,
+                        embedder_config=embedder_config,
+                        path=path,
+                    )
+                )
+            except ValueError:
+                # Fallback to default storage
+                storage = RAGStorage(
                     type="entities",
                     allow_reset=True,
                     crew=crew,
                     embedder_config=embedder_config,
                     path=path,
                 )
-            )
-        except ValueError:
-            # Fallback to default storage
-            self.storage = RAGStorage(
-                type="entities",
-                allow_reset=True,
-                crew=crew,
-                embedder_config=embedder_config,
-                path=path,
-            )
+        
+        # Initialize with parameters
+        super().__init__(
+            storage=storage,
+            embedder_config=embedder_config,
+            memory_provider=memory_provider
+        )
         
 
     def save(
diff --git a/src/crewai/memory/storage/base_rag_storage.py b/src/crewai/memory/storage/base_rag_storage.py
index 3605c04e7..1c5772e3d 100644
--- a/src/crewai/memory/storage/base_rag_storage.py
+++ b/src/crewai/memory/storage/base_rag_storage.py
@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional
 from crewai.memory.storage.interface import SearchResult, Storage
 
 
-class BaseRAGStorage(Storage, ABC):
+class BaseRAGStorage(Storage[Any], ABC):
     """
     Base class for RAG-based Storage implementations.
     """
diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py
index c3c3643ec..aebb3ebf5 100644
--- a/src/crewai/memory/storage/rag_storage.py
+++ b/src/crewai/memory/storage/rag_storage.py
@@ -38,7 +38,7 @@ class RAGStorage(BaseRAGStorage):
     search efficiency.
     """
 
-    app: ClientAPI | None = None
+    app: Optional[ClientAPI] = None
 
     def __init__(
         self, type, allow_reset=True, embedder_config=None, crew=None, path=None
@@ -137,7 +137,7 @@ class RAGStorage(BaseRAGStorage):
             logging.error(f"Error during {self.type} search: {str(e)}")
             return []
 
-    def _generate_embedding(self, text: str, metadata: Dict[str, Any]) -> None:  # type: ignore
+    def _generate_embedding(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> Any:
         if not hasattr(self, "app") or not hasattr(self, "collection"):
             self._initialize_app()
 
diff --git a/tests/memory/test_custom_storage.py b/tests/memory/test_custom_storage.py
index 6ade5e42c..dd725b038 100644
--- a/tests/memory/test_custom_storage.py
+++ b/tests/memory/test_custom_storage.py
@@ -11,7 +11,7 @@ from crewai.memory.storage.interface import SearchResult, Storage
 from crewai.memory.user.user_memory import UserMemory
 
 
-class CustomStorage(Storage):
+class CustomStorage(Storage[Any]):
     """Custom storage implementation for testing."""
 
     def __init__(self):
@@ -141,7 +141,7 @@ def test_custom_storage_with_memory_config():
 def test_custom_storage_error_handling():
     """Test error handling with custom storage."""
     # Test exception propagation
-    class ErrorStorage(Storage):
+    class ErrorStorage(Storage[Any]):
         """Storage implementation that raises exceptions."""
         def __init__(self):
             self.data = []
@@ -172,7 +172,7 @@ def test_custom_storage_error_handling():
 
 def test_custom_storage_edge_cases():
     """Test edge cases with custom storage."""
-    class EdgeCaseStorage(Storage):
+    class EdgeCaseStorage(Storage[Any]):
         """Storage implementation for testing edge cases."""
         def __init__(self):
             self.data = []