andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-10 00:28:31 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Improve type system and test coverage for custom memory storage

Browse Source 
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-03-04 20:23:51 +00:00
parent 541fa13df7
commit b13590a359
14 changed files with 318 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/crewai/crew.py
View File

@@ -262,8 +262,19 @@ class Crew(BaseModel):
def create_crew_memory(self) -> "Crew":
"""Set private attributes."""
if self.memory:
from crewai.memory.storage.rag_storage import RAGStorage
# Create default storage instances for each memory type if needed
long_term_storage = RAGStorage(type="long_term", crew=self, embedder_config=self.embedder)
short_term_storage = RAGStorage(type="short_term", crew=self, embedder_config=self.embedder)
entity_storage = RAGStorage(type="entity", crew=self, embedder_config=self.embedder)
self._long_term_memory = (
self.long_term_memory if self.long_term_memory else LongTermMemory(crew=self, embedder_config=self.embedder)
self.long_term_memory if self.long_term_memory else LongTermMemory(
crew=self,
embedder_config=self.embedder,
storage=long_term_storage
)
)
self._short_term_memory = (
self.short_term_memory
@@ -271,12 +282,17 @@ class Crew(BaseModel):
else ShortTermMemory(
crew=self,
embedder_config=self.embedder,
storage=short_term_storage
)
)
self._entity_memory = (
self.entity_memory
if self.entity_memory
else EntityMemory(crew=self, embedder_config=self.embedder)
else EntityMemory(
crew=self,
embedder_config=self.embedder,
storage=entity_storage
)
)
if (
self.memory_config and "user_memory" in self.memory_config

8
src/crewai/memory/contextual/contextual_memory.py
View File

@@ -47,7 +47,7 @@ class ContextualMemory:
stm_results = self.stm.search(query)
formatted_results = "\n".join(
[
f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
f"- {result.get('memory', result.get('context', ''))}"
for result in stm_results
]
)
@@ -80,9 +80,9 @@ class ContextualMemory:
em_results = self.em.search(query)
formatted_results = "\n".join(
[
f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
f"- {result.get('memory', result.get('context', ''))}"
for result in em_results
] # type: ignore # Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
]
)
return f"Entities:\n{formatted_results}" if em_results else ""
@@ -99,6 +99,6 @@ class ContextualMemory:
return ""
formatted_memories = "\n".join(
f"- {result['memory']}" for result in user_memories
f"- {result.get('memory', result.get('context', ''))}" for result in user_memories
)
return f"User memories/preferences:\n{formatted_memories}"

50
src/crewai/memory/entity/entity_memory.py
View File

@@ -18,47 +18,43 @@ class EntityMemory(Memory):
def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
memory_provider = None
entity_storage = None
memory_config = 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", {})
entity_storage = storage_config.get("entity")
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
)
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`."
)
super().__init__(
storage=Mem0Storage(type="entities", crew=crew),
embedder_config=embedder_config,
)
else:
# Use RAGStorage (default)
super().__init__(
storage=RAGStorage(
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(
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,
)

42
src/crewai/memory/long_term/long_term_memory.py
View File

@@ -16,40 +16,32 @@ class LongTermMemory(Memory):
def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
memory_provider = None
long_term_storage = None
memory_config = 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")
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
)
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,
try:
# Try to select storage using helper method
self.storage = self._select_storage(
storage=storage,
memory_config=memory_config,
storage_type="long_term",
crew=crew,
path=path,
default_storage_factory=lambda path, crew: LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
)
else:
# Use LTMSQLiteStorage (default)
storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
super().__init__(storage=storage, embedder_config=embedder_config)
except ValueError:
# Fallback to default storage
self.storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
def save(
self,

61
src/crewai/memory/memory.py
View File

@@ -1,20 +1,62 @@
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, TypeVar, Generic, Callable, cast
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, ConfigDict
from crewai.memory.storage.interface import Storage, SearchResult
class Memory(BaseModel):
T = TypeVar('T', bound=Storage)
class Memory(BaseModel, Generic[T]):
"""
Base class for memory, now supporting agent tags and generic metadata.
"""
model_config = ConfigDict(arbitrary_types_allowed=True)
embedder_config: Optional[Dict[str, Any]] = None
storage: Any
storage: T
memory_provider: Optional[str] = Field(default=None, exclude=True)
def __init__(self, storage: Any, **data: Any):
def __init__(self, storage: T, **data: Any):
super().__init__(storage=storage, **data)
def _select_storage(
self,
storage: Optional[T] = None,
memory_config: Optional[Dict[str, Any]] = None,
storage_type: str = "",
crew=None,
path: Optional[str] = None,
default_storage_factory: Optional[Callable] = None,
) -> T:
"""Helper method to select the appropriate storage based on configuration"""
# Use the provided storage if available
if storage:
return storage
# Use storage from memory_config if available
if memory_config and "storage" in memory_config:
storage_config = memory_config.get("storage", {})
if storage_type in storage_config and storage_config[storage_type]:
return cast(T, storage_config[storage_type])
# Use Mem0Storage if specified in memory_config
if memory_config and memory_config.get("provider") == "mem0":
try:
from crewai.memory.storage.mem0_storage import Mem0Storage
return cast(T, Mem0Storage(type=storage_type, crew=crew))
except ImportError:
raise ImportError(
"Mem0 is not installed. Please install it with `pip install mem0ai`."
)
# Use default storage if provided
if default_storage_factory:
return cast(T, default_storage_factory(path=path, crew=crew))
# Fallback to empty storage
raise ValueError(f"No storage available for {storage_type}")
def save(
self,
value: Any,
@@ -25,14 +67,19 @@ class Memory(BaseModel):
if agent:
metadata["agent"] = agent
self.storage.save(value, metadata)
if self.storage:
self.storage.save(value, metadata)
else:
raise ValueError("Storage is not initialized")
def search(
self,
query: str,
limit: int = 3,
score_threshold: float = 0.35,
) -> List[Any]:
) -> List[SearchResult]:
if not self.storage:
raise ValueError("Storage is not initialized")
return self.storage.search(
query=query, limit=limit, score_threshold=score_threshold
)

49
src/crewai/memory/short_term/short_term_memory.py
View File

@@ -20,46 +20,41 @@ class ShortTermMemory(Memory):
def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
memory_provider = None
short_term_storage = None
memory_config = 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", {})
short_term_storage = storage_config.get("short_term")
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
)
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`."
)
super().__init__(
storage=Mem0Storage(type="short_term", crew=crew),
embedder_config=embedder_config,
)
else:
# Use RAGStorage (default)
super().__init__(
storage=RAGStorage(
try:
# Try to select storage using helper method
self.storage = self._select_storage(
storage=storage,
memory_config=memory_config,
storage_type="short_term",
crew=crew,
path=path,
default_storage_factory=lambda path, crew: RAGStorage(
type="short_term",
crew=crew,
embedder_config=embedder_config,
path=path,
),
)
)
except ValueError:
# Fallback to default storage
self.storage = RAGStorage(
type="short_term",
crew=crew,
embedder_config=embedder_config,
path=path,
)
def save(

7
src/crewai/memory/storage/base_rag_storage.py
View File

@@ -1,8 +1,10 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from crewai.memory.storage.interface import Storage, SearchResult
class BaseRAGStorage(ABC):
class BaseRAGStorage(Storage, ABC):
"""
Base class for RAG-based Storage implementations.
"""
@@ -44,9 +46,8 @@ class BaseRAGStorage(ABC):
self,
query: str,
limit: int = 3,
filter: Optional[dict] = None,
score_threshold: float = 0.35,
) -> List[Any]:
) -> List[SearchResult]:
"""Search for entries in the storage."""
pass

31
src/crewai/memory/storage/interface.py
View File

@@ -1,16 +1,39 @@
from typing import Any, Dict, List
from abc import ABC, abstractmethod
from typing import Any, Dict, List, TypeVar, Generic, TypedDict, ClassVar, Protocol, runtime_checkable
from pydantic import BaseModel, ConfigDict
class Storage:
class SearchResult(TypedDict, total=False):
"""Type definition for search results"""
context: str
metadata: Dict[str, Any]
score: float
memory: str # For Mem0Storage compatibility
T = TypeVar('T')
@runtime_checkable
class StorageProtocol(Protocol):
"""Protocol defining the storage interface"""
def save(self, value: Any, metadata: Dict[str, Any]) -> None: ...
def search(self, query: str, limit: int, score_threshold: float) -> List[Any]: ...
def reset(self) -> None: ...
class Storage(ABC, Generic[T]):
"""Abstract base class defining the storage interface"""
model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
@abstractmethod
def save(self, value: Any, metadata: Dict[str, Any]) -> None:
pass
@abstractmethod
def search(
self, query: str, limit: int, score_threshold: float
) -> Dict[str, Any] | List[Any]:
return {}
) -> List[SearchResult]:
pass
@abstractmethod
def reset(self) -> None:
pass

6
src/crewai/memory/storage/mem0_storage.py
View File

@@ -111,3 +111,9 @@ class Mem0Storage(Storage):
agents = [self._sanitize_role(agent.role) for agent in agents]
agents = "_".join(agents)
return agents
def reset(self) -> None:
"""Reset the storage by clearing all memories."""
# Mem0 doesn't have a direct reset method, but we can implement
# this in the future if needed. For now, we'll just pass.
pass

7
src/crewai/memory/storage/rag_storage.py
View File

@@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional
from chromadb.api import ClientAPI
from crewai.memory.storage.base_rag_storage import BaseRAGStorage
from crewai.memory.storage.interface import SearchResult
from crewai.utilities import EmbeddingConfigurator
from crewai.utilities.constants import MAX_FILE_NAME_LENGTH
from crewai.utilities.paths import db_storage_path
@@ -112,9 +113,8 @@ class RAGStorage(BaseRAGStorage):
self,
query: str,
limit: int = 3,
filter: Optional[dict] = None,
score_threshold: float = 0.35,
) -> List[Any]:
) -> List[SearchResult]:
if not hasattr(self, "app"):
self._initialize_app()
@@ -124,8 +124,7 @@ class RAGStorage(BaseRAGStorage):
results = []
for i in range(len(response["ids"][0])):
result = {
"id": response["ids"][0][i],
result: SearchResult = {
"metadata": response["metadatas"][0][i],
"context": response["documents"][0][i],
"score": response["distances"][0][i],

50
src/crewai/memory/user/user_memory.py
View File

@@ -13,47 +13,43 @@ class UserMemory(Memory):
def __init__(self, crew=None, embedder_config=None, storage=None, path=None, **kwargs):
memory_provider = None
user_storage = None
memory_config = 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")
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
)
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)
try:
# Try to select storage using helper method
from crewai.memory.storage.rag_storage import RAGStorage
super().__init__(
storage=RAGStorage(
self.storage = self._select_storage(
storage=storage,
memory_config=memory_config,
storage_type="user",
crew=crew,
path=path,
default_storage_factory=lambda path, crew: RAGStorage(
type="user",
crew=crew,
embedder_config=embedder_config,
path=path,
),
)
)
except ValueError:
# Fallback to default storage
from crewai.memory.storage.rag_storage import RAGStorage
self.storage = RAGStorage(
type="user",
crew=crew,
embedder_config=embedder_config,
path=path,
)
def save(