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

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

src/crewai/crew.py

@@ -263,7 +263,7 @@ class Crew(BaseModel):
         """Set private attributes."""
         if self.memory:
             self._long_term_memory = (
-                self.long_term_memory if self.long_term_memory else LongTermMemory()
+                self.long_term_memory if self.long_term_memory else LongTermMemory(crew=self, embedder_config=self.embedder)
             )
             self._short_term_memory = (
                 self.short_term_memory

src/crewai/memory/entity/entity_memory.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from pydantic import PrivateAttr
 
@@ -17,47 +17,73 @@ class EntityMemory(Memory):
     _memory_provider: Optional[str] = PrivateAttr()
 
     def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
+        memory_provider = None
+        entity_storage = None
+        
         if crew and hasattr(crew, "memory_config") and crew.memory_config is not None:
             memory_provider = crew.memory_config.get("provider")
-        else:
-            memory_provider = None
+            storage_config = crew.memory_config.get("storage", {})
+            entity_storage = storage_config.get("entity")
+            
+        super().__init__(
+            storage=storage,
+            embedder_config=embedder_config,
+            memory_provider=memory_provider
+        )
 
-        if memory_provider == "mem0":
+        if storage:
+            # Use the provided storage
+            super().__init__(storage=storage, embedder_config=embedder_config)
+        elif entity_storage:
+            # Use the storage from memory_config
+            super().__init__(storage=entity_storage, embedder_config=embedder_config)
+        elif memory_provider == "mem0":
             try:
                 from crewai.memory.storage.mem0_storage import Mem0Storage
             except ImportError:
                 raise ImportError(
                     "Mem0 is not installed. Please install it with `pip install mem0ai`."
                 )
-            storage = Mem0Storage(type="entities", crew=crew)
+            super().__init__(
+                storage=Mem0Storage(type="entities", crew=crew),
+                embedder_config=embedder_config,
+            )
         else:
-            storage = (
-                storage
-                if storage
-                else RAGStorage(
+            # Use RAGStorage (default)
+            super().__init__(
+                storage=RAGStorage(
                     type="entities",
                     allow_reset=True,
-                    embedder_config=embedder_config,
                     crew=crew,
+                    embedder_config=embedder_config,
                     path=path,
-                )
+                ),
+                embedder_config=embedder_config,
             )
+        
 
-        super().__init__(storage=storage)
-        self._memory_provider = memory_provider
-
-    def save(self, item: EntityMemoryItem) -> None:  # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
-        """Saves an entity item into the SQLite storage."""
-        if self._memory_provider == "mem0":
-            data = f"""
-            Remember details about the following entity:
-            Name: {item.name}
-            Type: {item.type}
-            Entity Description: {item.description}
-            """
+    def save(
+        self,
+        value: Any,
+        metadata: Dict[str, Any] = None,
+        agent: str = None,
+    ) -> None:
+        """Saves an entity item or value into the storage."""
+        if isinstance(value, EntityMemoryItem):
+            item = value
+            if self.memory_provider == "mem0":
+                data = f"""
+                Remember details about the following entity:
+                Name: {item.name}
+                Type: {item.type}
+                Entity Description: {item.description}
+                """
+            else:
+                data = f"{item.name}({item.type}): {item.description}"
+            super().save(data, item.metadata)
         else:
-            data = f"{item.name}({item.type}): {item.description}"
-        super().save(data, item.metadata)
+            # Handle regular value and metadata
+            super().save(value, metadata, agent)
 
     def reset(self) -> None:
         try:

src/crewai/memory/long_term/long_term_memory.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from crewai.memory.long_term.long_term_memory_item import LongTermMemoryItem
 from crewai.memory.memory import Memory
@@ -14,23 +14,85 @@ class LongTermMemory(Memory):
     LongTermMemoryItem instances.
     """
 
-    def __init__(self, storage=None, path=None):
-        if not storage:
-            storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
-        super().__init__(storage=storage)
-
-    def save(self, item: LongTermMemoryItem) -> None:  # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"
-        metadata = item.metadata
-        metadata.update({"agent": item.agent, "expected_output": item.expected_output})
-        self.storage.save(  # type: ignore # BUG?: Unexpected keyword argument "task_description","score","datetime" for "save" of "Storage"
-            task_description=item.task,
-            score=metadata["quality"],
-            metadata=metadata,
-            datetime=item.datetime,
+    def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
+        memory_provider = None
+        long_term_storage = None
+        
+        if crew and hasattr(crew, "memory_config") and crew.memory_config is not None:
+            memory_provider = crew.memory_config.get("provider")
+            storage_config = crew.memory_config.get("storage", {})
+            long_term_storage = storage_config.get("long_term")
+            
+        super().__init__(
+            storage=storage,
+            embedder_config=embedder_config,
+            memory_provider=memory_provider
         )
 
-    def search(self, task: str, latest_n: int = 3) -> List[Dict[str, Any]]:  # type: ignore # signature of "search" incompatible with supertype "Memory"
-        return self.storage.load(task, latest_n)  # type: ignore # BUG?: "Storage" has no attribute "load"
+        if storage:
+            # Use the provided storage
+            super().__init__(storage=storage, embedder_config=embedder_config)
+        elif long_term_storage:
+            # Use the storage from memory_config
+            super().__init__(storage=long_term_storage, embedder_config=embedder_config)
+        elif memory_provider == "mem0":
+            try:
+                from crewai.memory.storage.mem0_storage import Mem0Storage
+            except ImportError:
+                raise ImportError(
+                    "Mem0 is not installed. Please install it with `pip install mem0ai`."
+                )
+            super().__init__(
+                storage=Mem0Storage(type="long_term", crew=crew),
+                embedder_config=embedder_config,
+            )
+        else:
+            # Use LTMSQLiteStorage (default)
+            storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
+            super().__init__(storage=storage, embedder_config=embedder_config)
+
+    def save(
+        self,
+        value: Any,
+        metadata: Dict[str, Any] = None,
+        agent: str = None,
+    ) -> None:
+        """Saves a value into the memory."""
+        if isinstance(value, LongTermMemoryItem):
+            item = value
+            item_metadata = item.metadata or {}
+            item_metadata.update({"agent": item.agent, "expected_output": item.expected_output})
+            
+            # Handle special storage types like Mem0Storage
+            if hasattr(self.storage, "save") and callable(getattr(self.storage, "save")) and hasattr(self.storage.save, "__code__") and "task_description" in self.storage.save.__code__.co_varnames:
+                self.storage.save(
+                    task_description=item.task,
+                    score=item_metadata.get("quality", 0),
+                    metadata=item_metadata,
+                    datetime=item.datetime,
+                )
+            else:
+                # Use standard storage interface
+                self.storage.save(item.task, item_metadata)
+        else:
+            # Handle regular value and metadata
+            super().save(value, metadata, agent)
+
+    def search(
+        self,
+        query: str,
+        limit: int = 3,
+        score_threshold: float = 0.35,
+    ) -> List[Any]:
+        """Search for values in the memory."""
+        # Try to use the standard storage interface first
+        if hasattr(self.storage, "search") and callable(getattr(self.storage, "search")):
+            return self.storage.search(query=query, limit=limit, score_threshold=score_threshold)
+        # Fall back to load method for backward compatibility
+        elif hasattr(self.storage, "load") and callable(getattr(self.storage, "load")):
+            return self.storage.load(query, limit)
+        else:
+            raise AttributeError("Storage does not implement search or load method")
 
     def reset(self) -> None:
         self.storage.reset()

src/crewai/memory/memory.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class Memory(BaseModel):
@@ -9,8 +9,8 @@ class Memory(BaseModel):
     """
 
     embedder_config: Optional[Dict[str, Any]] = None
-
     storage: Any
+    memory_provider: Optional[str] = Field(default=None, exclude=True)
 
     def __init__(self, storage: Any, **data: Any):
         super().__init__(storage=storage, **data)

src/crewai/memory/short_term/short_term_memory.py

@@ -19,32 +19,48 @@ class ShortTermMemory(Memory):
     _memory_provider: Optional[str] = PrivateAttr()
 
     def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
+        memory_provider = None
+        short_term_storage = None
+        
         if crew and hasattr(crew, "memory_config") and crew.memory_config is not None:
             memory_provider = crew.memory_config.get("provider")
-        else:
-            memory_provider = None
+            storage_config = crew.memory_config.get("storage", {})
+            short_term_storage = storage_config.get("short_term")
+            
+        super().__init__(
+            storage=storage,
+            embedder_config=embedder_config,
+            memory_provider=memory_provider
+        )
 
-        if memory_provider == "mem0":
+        if storage:
+            # Use the provided storage
+            super().__init__(storage=storage, embedder_config=embedder_config)
+        elif short_term_storage:
+            # Use the storage from memory_config
+            super().__init__(storage=short_term_storage, embedder_config=embedder_config)
+        elif memory_provider == "mem0":
             try:
                 from crewai.memory.storage.mem0_storage import Mem0Storage
             except ImportError:
                 raise ImportError(
                     "Mem0 is not installed. Please install it with `pip install mem0ai`."
                 )
-            storage = Mem0Storage(type="short_term", crew=crew)
-        else:
-            storage = (
-                storage
-                if storage
-                else RAGStorage(
-                    type="short_term",
-                    embedder_config=embedder_config,
-                    crew=crew,
-                    path=path,
-                )
+            super().__init__(
+                storage=Mem0Storage(type="short_term", crew=crew),
+                embedder_config=embedder_config,
+            )
+        else:
+            # Use RAGStorage (default)
+            super().__init__(
+                storage=RAGStorage(
+                    type="short_term",
+                    crew=crew,
+                    embedder_config=embedder_config,
+                    path=path,
+                ),
+                embedder_config=embedder_config,
             )
-        super().__init__(storage=storage)
-        self._memory_provider = memory_provider
 
     def save(
         self,
@@ -53,7 +69,7 @@ class ShortTermMemory(Memory):
         agent: Optional[str] = None,
     ) -> None:
         item = ShortTermMemoryItem(data=value, metadata=metadata, agent=agent)
-        if self._memory_provider == "mem0":
+        if self.memory_provider == "mem0":
             item.data = f"Remember the following insights from Agent run: {item.data}"
 
         super().save(value=item.data, metadata=item.metadata, agent=item.agent)

src/crewai/memory/user/user_memory.py

@@ -11,15 +11,50 @@ class UserMemory(Memory):
     MemoryItem instances.
     """
 
-    def __init__(self, crew=None):
-        try:
-            from crewai.memory.storage.mem0_storage import Mem0Storage
-        except ImportError:
-            raise ImportError(
-                "Mem0 is not installed. Please install it with `pip install mem0ai`."
+    def __init__(self, crew=None, embedder_config=None, storage=None, path=None, **kwargs):
+        memory_provider = None
+        user_storage = None
+        
+        if crew and hasattr(crew, "memory_config") and crew.memory_config is not None:
+            memory_provider = crew.memory_config.get("provider")
+            storage_config = crew.memory_config.get("storage", {})
+            user_storage = storage_config.get("user")
+            
+        super().__init__(
+            storage=storage,
+            embedder_config=embedder_config,
+            memory_provider=memory_provider
+        )
+
+        if storage:
+            # Use the provided storage
+            super().__init__(storage=storage, embedder_config=embedder_config)
+        elif user_storage:
+            # Use the storage from memory_config
+            super().__init__(storage=user_storage, embedder_config=embedder_config)
+        elif memory_provider == "mem0":
+            try:
+                from crewai.memory.storage.mem0_storage import Mem0Storage
+            except ImportError:
+                raise ImportError(
+                    "Mem0 is not installed. Please install it with `pip install mem0ai`."
+                )
+            super().__init__(
+                storage=Mem0Storage(type="user", crew=crew),
+                embedder_config=embedder_config,
+            )
+        else:
+            # Use RAGStorage (default)
+            from crewai.memory.storage.rag_storage import RAGStorage
+            super().__init__(
+                storage=RAGStorage(
+                    type="user",
+                    crew=crew,
+                    embedder_config=embedder_config,
+                    path=path,
+                ),
+                embedder_config=embedder_config,
             )
-        storage = Mem0Storage(type="user", crew=crew)
-        super().__init__(storage)
 
     def save(
         self,
@@ -43,3 +78,9 @@ class UserMemory(Memory):
             score_threshold=score_threshold,
         )
         return results
+        
+    def reset(self) -> None:
+        try:
+            self.storage.reset()
+        except Exception as e:
+            raise Exception(f"An error occurred while resetting the user memory: {e}")