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(

26
tests/memory/long_term_memory_test.py
View File

@@ -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):

7
tests/memory/short_term_memory_test.py
View File

@@ -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):

102
tests/memory/test_custom_storage.py
View File

@@ -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"]