andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-08 15:48:29 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

[FEATURE] Support for custom path in RAGStorage (#1659)

Browse Source 
* added path to RAGStorage

* added path to short term and entity memory

* add path for long_term_storage for completeness

---------

Co-authored-by: Brandon Hancock (bhancock_ai) <109994880+bhancockio@users.noreply.github.com>
This commit is contained in:
Tom Mahler, PhD
2024-12-03 19:22:29 +02:00
committed by GitHub
parent 79cfa87341
commit d6fdb9980c
4 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -10,7 +10,7 @@ class EntityMemory(Memory):
Inherits from the Memory class.
"""
def __init__(self, crew=None, embedder_config=None, storage=None):
def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
if hasattr(crew, "memory_config") and crew.memory_config is not None:
self.memory_provider = crew.memory_config.get("provider")
else:
@@ -33,6 +33,7 @@ class EntityMemory(Memory):
allow_reset=True,
embedder_config=embedder_config,
crew=crew,
path=path,
)
)
super().__init__(storage)

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

@@ -14,8 +14,9 @@ class LongTermMemory(Memory):
LongTermMemoryItem instances.
"""
def __init__(self, storage=None):
storage = storage if storage else LTMSQLiteStorage()
def __init__(self, storage=None, path=None):
if not storage:
storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
super().__init__(storage)
def save(self, item: LongTermMemoryItem) -> None: # type: ignore # BUG?: Signature of "save" incompatible with supertype "Memory"

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

@@ -13,7 +13,7 @@ class ShortTermMemory(Memory):
MemoryItem instances.
"""
def __init__(self, crew=None, embedder_config=None, storage=None):
def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
if hasattr(crew, "memory_config") and crew.memory_config is not None:
self.memory_provider = crew.memory_config.get("provider")
else:
@@ -32,7 +32,7 @@ class ShortTermMemory(Memory):
storage
if storage
else RAGStorage(
type="short_term", embedder_config=embedder_config, crew=crew
type="short_term", embedder_config=embedder_config, crew=crew, path=path
)
)
super().__init__(storage)

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

@@ -37,7 +37,7 @@ class RAGStorage(BaseRAGStorage):
app: ClientAPI | None = None
def __init__(self, type, allow_reset=True, embedder_config=None, crew=None):
def __init__(self, type, allow_reset=True, embedder_config=None, crew=None, path=None):
super().__init__(type, allow_reset, embedder_config, crew)
agents = crew.agents if crew else []
agents = [self._sanitize_role(agent.role) for agent in agents]
@@ -47,6 +47,7 @@ class RAGStorage(BaseRAGStorage):
self.type = type
self.allow_reset = allow_reset
self.path = path
self._initialize_app()
def _set_embedder_config(self):
@@ -59,7 +60,7 @@ class RAGStorage(BaseRAGStorage):
self._set_embedder_config()
chroma_client = chromadb.PersistentClient(
path=f"{db_storage_path()}/{self.type}/{self.agents}",
path=self.path if self.path else f"{db_storage_path()}/{self.type}/{self.agents}",
settings=Settings(allow_reset=self.allow_reset),
)