Update Mem0Storage to support v2 API with run_id parameter

Co-Authored-By: Joe Moura <joao@crewai.com>

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):
@@ -26,6 +29,11 @@ class Mem0Storage(Storage):
         # 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()
         if type == "user" and not 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
+        # Add version and run_id for MemoryClient
-        # automatically when the crew is created.
+        if isinstance(self.memory, MemoryClient):
-        if isinstance(self.memory, Memory):
+            params["version"] = self.version
-            del params["metadata"], params["output_format"]
+            
+            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]

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

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"