From 91ff331fecd893677aaae7ac6a3f6c6cd936ac3e Mon Sep 17 00:00:00 2001 From: Ayo Ayibiowu Date: Sun, 22 Sep 2024 22:03:23 +0200 Subject: [PATCH] feat(memory): adds support for customizable memory interface (#1339) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(memory): adds support for customizing crew storage * chore: allow overwriting the crew memory configuration * docs: update custom storage usage * fix(lint): use correct syntax * fix: type check warning * fix: type check warnings * fix(test): address agent default failing test * fix(lint). address type checker error * Update crew.py --------- Co-authored-by: João Moura --- docs/core-concepts/Memory.md | 41 ++++++++++++++++++- src/crewai/crew.py | 18 ++++++-- src/crewai/memory/entity/entity_memory.py | 13 +++--- .../memory/long_term/long_term_memory.py | 4 +- .../memory/short_term/short_term_memory.py | 10 +++-- 5 files changed, 71 insertions(+), 15 deletions(-) diff --git a/docs/core-concepts/Memory.md b/docs/core-concepts/Memory.md index a43fcedeb..099c25229 100644 --- a/docs/core-concepts/Memory.md +++ b/docs/core-concepts/Memory.md @@ -28,7 +28,7 @@ description: Leveraging memory systems in the crewAI framework to enhance agent ## Implementing Memory in Your Crew When configuring a crew, you can enable and customize each memory component to suit the crew's objectives and the nature of tasks it will perform. -By default, the memory system is disabled, and you can ensure it is active by setting `memory=True` in the crew configuration. The memory will use OpenAI embeddings by default, but you can change it by setting `embedder` to a different model. +By default, the memory system is disabled, and you can ensure it is active by setting `memory=True` in the crew configuration. The memory will use OpenAI embeddings by default, but you can change it by setting `embedder` to a different model. It's also possible to initialize the memory instance with your own instance. The 'embedder' only applies to **Short-Term Memory** which uses Chroma for RAG using the EmbedChain package. The **Long-Term Memory** uses SQLite3 to store task results. Currently, there is no way to override these storage implementations. @@ -50,6 +50,45 @@ my_crew = Crew( ) ``` +### Example: Use Custom Memory Instances e.g FAISS as the VectorDB + +```python +from crewai import Crew, Agent, Task, Process + +# Assemble your crew with memory capabilities +my_crew = Crew( + agents=[...], + tasks=[...], + process="Process.sequential", + memory=True, + long_term_memory=EnhanceLongTermMemory( + storage=LTMSQLiteStorage( + db_path="/my_data_dir/my_crew1/long_term_memory_storage.db" + ) + ), + short_term_memory=EnhanceShortTermMemory( + storage=CustomRAGStorage( + crew_name="my_crew", + storage_type="short_term", + data_dir="//my_data_dir", + model=embedder["model"], + dimension=embedder["dimension"], + ), + ), + entity_memory=EnhanceEntityMemory( + storage=CustomRAGStorage( + crew_name="my_crew", + storage_type="entities", + data_dir="//my_data_dir", + model=embedder["model"], + dimension=embedder["dimension"], + ), + ), + verbose=True, +) +``` + + ## Additional Embedding Providers ### Using OpenAI embeddings (already default) diff --git a/src/crewai/crew.py b/src/crewai/crew.py index 1e8d83ab4..6c6992e07 100644 --- a/src/crewai/crew.py +++ b/src/crewai/crew.py @@ -110,6 +110,18 @@ class Crew(BaseModel): default=False, description="Whether the crew should use memory to store memories of it's execution", ) + short_term_memory: Optional[InstanceOf[ShortTermMemory]] = Field( + default=None, + description="An Instance of the ShortTermMemory to be used by the Crew", + ) + long_term_memory: Optional[InstanceOf[LongTermMemory]] = Field( + default=None, + description="An Instance of the LongTermMemory to be used by the Crew", + ) + entity_memory: Optional[InstanceOf[EntityMemory]] = Field( + default=None, + description="An Instance of the EntityMemory to be used by the Crew", + ) embedder: Optional[dict] = Field( default={"provider": "openai"}, description="Configuration for the embedder to be used for the crew.", @@ -212,11 +224,11 @@ class Crew(BaseModel): def create_crew_memory(self) -> "Crew": """Set private attributes.""" if self.memory: - self._long_term_memory = LongTermMemory() - self._short_term_memory = ShortTermMemory( + self._long_term_memory = self.long_term_memory if self.long_term_memory else LongTermMemory() + self._short_term_memory = self.short_term_memory if self.short_term_memory else ShortTermMemory( crew=self, embedder_config=self.embedder ) - self._entity_memory = EntityMemory(crew=self, embedder_config=self.embedder) + self._entity_memory = self.entity_memory if self.entity_memory else EntityMemory(crew=self, embedder_config=self.embedder) return self @model_validator(mode="after") diff --git a/src/crewai/memory/entity/entity_memory.py b/src/crewai/memory/entity/entity_memory.py index 50aaeeaab..2a199349a 100644 --- a/src/crewai/memory/entity/entity_memory.py +++ b/src/crewai/memory/entity/entity_memory.py @@ -10,12 +10,13 @@ class EntityMemory(Memory): Inherits from the Memory class. """ - def __init__(self, crew=None, embedder_config=None): - storage = RAGStorage( - type="entities", - allow_reset=False, - embedder_config=embedder_config, - crew=crew, + def __init__(self, crew=None, embedder_config=None, storage=None): + storage = ( + storage + if storage + else RAGStorage( + type="entities", allow_reset=False, embedder_config=embedder_config, crew=crew + ) ) super().__init__(storage) diff --git a/src/crewai/memory/long_term/long_term_memory.py b/src/crewai/memory/long_term/long_term_memory.py index 041268107..ab225e406 100644 --- a/src/crewai/memory/long_term/long_term_memory.py +++ b/src/crewai/memory/long_term/long_term_memory.py @@ -14,8 +14,8 @@ class LongTermMemory(Memory): LongTermMemoryItem instances. """ - def __init__(self): - storage = LTMSQLiteStorage() + def __init__(self, storage=None): + storage = storage if storage else LTMSQLiteStorage() super().__init__(storage) def save(self, item: LongTermMemoryItem) -> None: # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory" diff --git a/src/crewai/memory/short_term/short_term_memory.py b/src/crewai/memory/short_term/short_term_memory.py index ea62f87f6..919fb6115 100644 --- a/src/crewai/memory/short_term/short_term_memory.py +++ b/src/crewai/memory/short_term/short_term_memory.py @@ -13,9 +13,13 @@ class ShortTermMemory(Memory): MemoryItem instances. """ - def __init__(self, crew=None, embedder_config=None): - storage = RAGStorage( - type="short_term", embedder_config=embedder_config, crew=crew + def __init__(self, crew=None, embedder_config=None, storage=None): + storage = ( + storage + if storage + else RAGStorage( + type="short_term", embedder_config=embedder_config, crew=crew + ) ) super().__init__(storage)