mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-07 07:08:31 +00:00
6f4ad532e6
* clean up. fix type safety. address memory config docs * improve manager * Include fix for o1 models not supporting system messages * more broad with o1 * address fix: Typo in expected_output string #2045 * drop prints * drop prints * wip * wip * fix failing memory tests * Fix memory provider issue * clean up short term memory * revert ltm * drop * clean up linting issues * more linting
67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
from typing import Optional
|
|
|
|
from pydantic import PrivateAttr
|
|
|
|
from crewai.memory.entity.entity_memory_item import EntityMemoryItem
|
|
from crewai.memory.memory import Memory
|
|
from crewai.memory.storage.rag_storage import RAGStorage
|
|
|
|
|
|
class EntityMemory(Memory):
|
|
"""
|
|
EntityMemory class for managing structured information about entities
|
|
and their relationships using SQLite storage.
|
|
Inherits from the Memory class.
|
|
"""
|
|
|
|
_memory_provider: Optional[str] = PrivateAttr()
|
|
|
|
def __init__(self, crew=None, embedder_config=None, storage=None, path=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
|
|
|
|
if 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)
|
|
else:
|
|
storage = (
|
|
storage
|
|
if storage
|
|
else RAGStorage(
|
|
type="entities",
|
|
allow_reset=True,
|
|
embedder_config=embedder_config,
|
|
crew=crew,
|
|
path=path,
|
|
)
|
|
)
|
|
|
|
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}
|
|
"""
|
|
else:
|
|
data = f"{item.name}({item.type}): {item.description}"
|
|
super().save(data, item.metadata)
|
|
|
|
def reset(self) -> None:
|
|
try:
|
|
self.storage.reset()
|
|
except Exception as e:
|
|
raise Exception(f"An error occurred while resetting the entity memory: {e}")
|