From 85ff023b767d26c619997ad220d252a19503a855 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 7 May 2025 21:42:36 +0000 Subject: [PATCH] Update Mem0Storage to support v2 API with run_id parameter Co-Authored-By: Joe Moura --- src/crewai/memory/storage/mem0_storage.py | 30 +++- tests/storage/test_mem0_storage.py | 4 +- tests/storage/test_mem0_storage_v2.py | 161 ++++++++++++++++++++++ 3 files changed, 187 insertions(+), 8 deletions(-) create mode 100644 tests/storage/test_mem0_storage_v2.py diff --git a/src/crewai/memory/storage/mem0_storage.py b/src/crewai/memory/storage/mem0_storage.py index 5d601ac1f..03588d2dd 100644 --- a/src/crewai/memory/storage/mem0_storage.py +++ b/src/crewai/memory/storage/mem0_storage.py @@ -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] diff --git a/tests/storage/test_mem0_storage.py b/tests/storage/test_mem0_storage.py index b651521f9..41e006279 100644 --- a/tests/storage/test_mem0_storage.py +++ b/tests/storage/test_mem0_storage.py @@ -195,7 +195,7 @@ def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_co agent_id="Test_Agent", infer=False, metadata={"type": "short_term", "key": "value"}, - output_format="v1.1" + version="v2" ) @@ -232,7 +232,7 @@ def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_ agent_id="Test_Agent", metadata={"type": "short_term"}, user_id="test_user", - output_format='v1.1' + version='v2' ) assert len(results) == 1 diff --git a/tests/storage/test_mem0_storage_v2.py b/tests/storage/test_mem0_storage_v2.py new file mode 100644 index 000000000..9b8ccb458 --- /dev/null +++ b/tests/storage/test_mem0_storage_v2.py @@ -0,0 +1,161 @@ +import os +from unittest.mock import MagicMock, patch + +import pytest +from mem0.client.main import MemoryClient +from mem0.memory.main import Memory + +from crewai.agent import Agent +from crewai.crew import Crew +from crewai.memory.storage.mem0_storage import Mem0Storage +from crewai.task import Task + + +class MockCrew: + def __init__(self, memory_config): + self.memory_config = memory_config + self.agents = [MagicMock(role="Test Agent")] + + +@pytest.fixture +def mock_mem0_memory_client(): + """Fixture to create a mock MemoryClient instance""" + mock_memory = MagicMock(spec=MemoryClient) + return mock_memory + + +@pytest.fixture +def mem0_storage_with_v2_api(mock_mem0_memory_client): + """Fixture to create a Mem0Storage instance with v2 API configuration""" + + with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client): + crew = MockCrew( + memory_config={ + "provider": "mem0", + "config": { + "user_id": "test_user", + "api_key": "ABCDEFGH", + "version": "v2", # Explicitly set to v2 + }, + } + ) + + mem0_storage = Mem0Storage(type="short_term", crew=crew) + return mem0_storage, mock_mem0_memory_client + + +@pytest.fixture +def mem0_storage_with_run_id(mock_mem0_memory_client): + """Fixture to create a Mem0Storage instance with run_id configuration""" + + with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client): + crew = MockCrew( + memory_config={ + "provider": "mem0", + "config": { + "user_id": "test_user", + "api_key": "ABCDEFGH", + "version": "v2", + "run_id": "test-session-123", # Set run_id + }, + } + ) + + mem0_storage = Mem0Storage(type="short_term", crew=crew) + return mem0_storage, mock_mem0_memory_client + + +def test_mem0_storage_v2_initialization(mem0_storage_with_v2_api): + """Test that Mem0Storage initializes correctly with v2 API configuration""" + mem0_storage, _ = mem0_storage_with_v2_api + + assert mem0_storage.version == "v2" + + assert mem0_storage.run_id is None + + +def test_mem0_storage_with_run_id_initialization(mem0_storage_with_run_id): + """Test that Mem0Storage initializes correctly with run_id configuration""" + mem0_storage, _ = mem0_storage_with_run_id + + assert mem0_storage.version == "v2" + + assert mem0_storage.run_id == "test-session-123" + + +def test_save_method_with_v2_api(mem0_storage_with_v2_api): + """Test save method with v2 API""" + mem0_storage, mock_memory_client = mem0_storage_with_v2_api + mock_memory_client.add = MagicMock() + + test_value = "This is a test memory" + test_metadata = {"key": "value"} + + mem0_storage.save(test_value, test_metadata) + + mock_memory_client.add.assert_called_once() + call_args = mock_memory_client.add.call_args[1] + + assert call_args["version"] == "v2" + assert "run_id" not in call_args + assert call_args["agent_id"] == "Test_Agent" + assert call_args["metadata"] == {"type": "short_term", "key": "value"} + + +def test_save_method_with_run_id(mem0_storage_with_run_id): + """Test save method with run_id""" + mem0_storage, mock_memory_client = mem0_storage_with_run_id + mock_memory_client.add = MagicMock() + + test_value = "This is a test memory" + test_metadata = {"key": "value"} + + mem0_storage.save(test_value, test_metadata) + + mock_memory_client.add.assert_called_once() + call_args = mock_memory_client.add.call_args[1] + + assert call_args["version"] == "v2" + assert call_args["run_id"] == "test-session-123" + assert call_args["agent_id"] == "Test_Agent" + assert call_args["metadata"] == {"type": "short_term", "key": "value"} + + +def test_search_method_with_v2_api(mem0_storage_with_v2_api): + """Test search method with v2 API""" + mem0_storage, mock_memory_client = mem0_storage_with_v2_api + mock_results = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]} + mock_memory_client.search = MagicMock(return_value=mock_results) + + results = mem0_storage.search("test query", limit=5, score_threshold=0.5) + + mock_memory_client.search.assert_called_once() + call_args = mock_memory_client.search.call_args[1] + + assert call_args["version"] == "v2" + assert "run_id" not in call_args + assert call_args["query"] == "test query" + assert call_args["limit"] == 5 + + assert len(results) == 1 + assert results[0]["content"] == "Result 1" + + +def test_search_method_with_run_id(mem0_storage_with_run_id): + """Test search method with run_id""" + mem0_storage, mock_memory_client = mem0_storage_with_run_id + mock_results = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]} + mock_memory_client.search = MagicMock(return_value=mock_results) + + results = mem0_storage.search("test query", limit=5, score_threshold=0.5) + + mock_memory_client.search.assert_called_once() + call_args = mock_memory_client.search.call_args[1] + + assert call_args["version"] == "v2" + assert call_args["run_id"] == "test-session-123" + assert call_args["query"] == "test query" + assert call_args["limit"] == 5 + + assert len(results) == 1 + assert results[0]["content"] == "Result 1"