fix: support to add memories to Mem0 with agent_id (#3217)

* fix: support to add memories to Mem0 with agent_id

* feat: removing memory_type checkings from Mem0Storage

* feat: ensure agent_id is always present while saving memory into Mem0

* fix: use OR operator when querying Mem0 memories with both user_id and agent_id

tests/storage/test_mem0_storage.py

@@ -191,17 +191,39 @@ def test_save_method_with_memory_oss(mem0_storage_with_mocked_config):
     """Test save method for different memory types"""
     mem0_storage, _, _ = mem0_storage_with_mocked_config
     mem0_storage.memory.add = MagicMock()
-    
+
     # Test short_term memory type (already set in fixture)
     test_value = "This is a test memory"
     test_metadata = {"key": "value"}
-    
+
     mem0_storage.save(test_value, test_metadata)
-    
+
     mem0_storage.memory.add.assert_called_once_with(
-        [{'role': 'assistant' , 'content': test_value}],
+        [{"role": "assistant" , "content": test_value}],
         infer=True,
         metadata={"type": "short_term", "key": "value"},
+        run_id="my_run_id",
+        user_id="test_user",
+        agent_id='Test_Agent'
+    )
+
+def test_save_method_with_multiple_agents(mem0_storage_with_mocked_config):
+    mem0_storage, _, _ = mem0_storage_with_mocked_config
+    mem0_storage.crew.agents = [MagicMock(role="Test Agent"), MagicMock(role="Test Agent 2"), MagicMock(role="Test Agent 3")]
+    mem0_storage.memory.add = MagicMock()
+
+    test_value = "This is a test memory"
+    test_metadata = {"key": "value"}
+
+    mem0_storage.save(test_value, test_metadata)
+
+    mem0_storage.memory.add.assert_called_once_with(
+        [{"role": "assistant" , "content": test_value}],
+        infer=True,
+        metadata={"type": "short_term", "key": "value"},
+        run_id="my_run_id",
+        user_id="test_user",
+        agent_id='Test_Agent_Test_Agent_2_Test_Agent_3'
     )
 
 
@@ -209,13 +231,13 @@ def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_co
     """Test save method for different memory types"""
     mem0_storage = mem0_storage_with_memory_client_using_config_from_crew
     mem0_storage.memory.add = MagicMock()
-    
+
     # Test short_term memory type (already set in fixture)
     test_value = "This is a test memory"
     test_metadata = {"key": "value"}
-    
+
     mem0_storage.save(test_value, test_metadata)
-    
+
     mem0_storage.memory.add.assert_called_once_with(
         [{'role': 'assistant' , 'content': test_value}],
         infer=True,
@@ -224,7 +246,9 @@ def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_co
         run_id="my_run_id",
         includes="include1",
         excludes="exclude1",
-        output_format='v1.1'
+        output_format='v1.1',
+        user_id='test_user',
+         agent_id='Test_Agent'
     )
 
 
@@ -237,10 +261,10 @@ def test_search_method_with_memory_oss(mem0_storage_with_mocked_config):
     results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
 
     mem0_storage.memory.search.assert_called_once_with(
-        query="test query", 
-        limit=5, 
+        query="test query",
+        limit=5,
         user_id="test_user",
-        filters={'AND': [{'run_id': 'my_run_id'}]}, 
+        filters={'AND': [{'run_id': 'my_run_id'}]},
         threshold=0.5
     )
 
@@ -257,8 +281,8 @@ def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_
     results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
 
     mem0_storage.memory.search.assert_called_once_with(
-        query="test query", 
-        limit=5, 
+        query="test query",
+        limit=5,
         metadata={"type": "short_term"},
         user_id="test_user",
         version='v2',
@@ -286,4 +310,56 @@ def test_mem0_storage_default_infer_value(mock_mem0_memory_client):
         )
 
         mem0_storage = Mem0Storage(type="short_term", crew=crew)
-        assert mem0_storage.infer is True
+        assert mem0_storage.infer is True
+
+def test_save_memory_using_agent_entity(mock_mem0_memory_client):
+    config = {
+        "agent_id": "agent-123",
+    }
+
+    mock_memory = MagicMock(spec=Memory)
+    with patch.object(Memory, "__new__", return_value=mock_memory):
+        mem0_storage = Mem0Storage(type="external", config=config)
+        mem0_storage.save("test memory", {"key": "value"})
+        mem0_storage.memory.add.assert_called_once_with(
+            [{'role': 'assistant' , 'content': 'test memory'}],
+            infer=True,
+            metadata={"type": "external", "key": "value"},
+            agent_id="agent-123",
+        )
+
+def test_search_method_with_agent_entity():
+    mem0_storage = Mem0Storage(type="external", config={"agent_id": "agent-123"})
+    mock_results = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]}
+    mem0_storage.memory.search = MagicMock(return_value=mock_results)
+
+    results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
+
+    mem0_storage.memory.search.assert_called_once_with(
+        query="test query",
+        limit=5,
+        filters={"AND": [{"agent_id": "agent-123"}]},
+        threshold=0.5,
+    )
+
+    assert len(results) == 2
+    assert results[0]["content"] == "Result 1"
+
+
+def test_search_method_with_agent_id_and_user_id():
+    mem0_storage = Mem0Storage(type="external", config={"agent_id": "agent-123", "user_id": "user-123"})
+    mock_results = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]}
+    mem0_storage.memory.search = MagicMock(return_value=mock_results)
+
+    results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
+
+    mem0_storage.memory.search.assert_called_once_with(
+        query="test query",
+        limit=5,
+        user_id='user-123',
+        filters={"OR": [{"user_id": "user-123"}, {"agent_id": "agent-123"}]},
+        threshold=0.5,
+    )
+
+    assert len(results) == 2
+    assert results[0]["content"] == "Result 1"