andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-05-03 08:12:39 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Update Mem0Storage to support v2 API with run_id parameter

Browse Source 
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-05-07 21:42:36 +00:00
parent 7ad51d9d05
commit 85ff023b76
3 changed files with 187 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,6 +9,9 @@ from crewai.memory.storage.interface import Storage
class Mem0Storage(Storage):
"""
Extends Storage to handle embedding and searching across entities using Mem0.
Supports Mem0 v2 API with run_id for associating memories with specific conversation
sessions. By default, uses v2 API which is recommended for better context management.
"""
def __init__(self, type, crew=None, config=None):
@@ -25,6 +28,11 @@ class Mem0Storage(Storage):
self.config = config or {}
# TODO: Memory config will be removed in the future the config will be passed as a parameter
self.memory_config = self.config or getattr(crew, "memory_config", {}) or {}
config = self._get_config()
self.version = config.get("version", "v2")
self.run_id = config.get("run_id")
# User ID is required for user memory type "user" since it's used as a unique identifier for the user.
user_id = self._get_user_id()
@@ -89,7 +97,11 @@ class Mem0Storage(Storage):
if params:
if isinstance(self.memory, MemoryClient):
params["output_format"] = "v1.1"
params["version"] = self.version
if self.run_id:
params["run_id"] = self.run_id
self.memory.add(value, **params)
def search(
@@ -98,7 +110,7 @@ class Mem0Storage(Storage):
limit: int = 3,
score_threshold: float = 0.35,
) -> List[Any]:
params = {"query": query, "limit": limit, "output_format": "v1.1"}
params = {"query": query, "limit": limit}
if user_id := self._get_user_id():
params["user_id"] = user_id
@@ -116,10 +128,16 @@ class Mem0Storage(Storage):
params["agent_id"] = agent_name
params["metadata"] = {"type": "external"}
# Discard the filters for now since we create the filters
# automatically when the crew is created.
if isinstance(self.memory, Memory):
del params["metadata"], params["output_format"]
# Add version and run_id for MemoryClient
if isinstance(self.memory, MemoryClient):
params["version"] = self.version
if self.run_id:
params["run_id"] = self.run_id
# Discard the filters for Memory (OSS version)
elif isinstance(self.memory, Memory):
if "metadata" in params:
del params["metadata"]
results = self.memory.search(**params)
return [r for r in results["results"] if r["score"] >= score_threshold]