Fix Memory OSS compatibility and implement code review suggestions

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-05-07 21:53:07 +00:00
parent 85ff023b76
commit c53cd2775a
2 changed files with 258 additions and 33 deletions

View File

@@ -65,24 +65,54 @@ def mem0_storage_with_run_id(mock_mem0_memory_client):
return mem0_storage, mock_mem0_memory_client
@pytest.fixture
def mem0_storage_with_v1_api(mock_mem0_memory_client):
"""Fixture to create a Mem0Storage instance with v1.1 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": "v1.1", # Explicitly set to v1.1
},
}
)
mem0_storage = Mem0Storage(type="short_term", crew=crew)
return mem0_storage, mock_mem0_memory_client
@pytest.mark.v2_api
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
@pytest.mark.v2_api
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"
@pytest.mark.v1_api
def test_mem0_storage_v1_initialization(mem0_storage_with_v1_api):
"""Test that Mem0Storage initializes correctly with v1.1 API configuration"""
mem0_storage, _ = mem0_storage_with_v1_api
assert mem0_storage.version == "v1.1"
assert mem0_storage.run_id is None
@pytest.mark.v2_api
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
@@ -102,6 +132,7 @@ def test_save_method_with_v2_api(mem0_storage_with_v2_api):
assert call_args["metadata"] == {"type": "short_term", "key": "value"}
@pytest.mark.v2_api
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
@@ -121,6 +152,7 @@ def test_save_method_with_run_id(mem0_storage_with_run_id):
assert call_args["metadata"] == {"type": "short_term", "key": "value"}
@pytest.mark.v2_api
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
@@ -141,6 +173,7 @@ def test_search_method_with_v2_api(mem0_storage_with_v2_api):
assert results[0]["content"] == "Result 1"
@pytest.mark.v2_api
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
@@ -159,3 +192,99 @@ def test_search_method_with_run_id(mem0_storage_with_run_id):
assert len(results) == 1
assert results[0]["content"] == "Result 1"
@pytest.mark.v2_api
def test_search_method_with_different_result_formats(mem0_storage_with_v2_api):
"""Test search method with different result formats"""
mem0_storage, mock_memory_client = mem0_storage_with_v2_api
mock_results_dict = {"results": [{"score": 0.9, "content": "Result 1"}, {"score": 0.4, "content": "Result 2"}]}
mock_memory_client.search = MagicMock(return_value=mock_results_dict)
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
assert len(results) == 1
assert results[0]["content"] == "Result 1"
mock_results_list = [{"score": 0.9, "content": "Result 3"}, {"score": 0.4, "content": "Result 4"}]
mock_memory_client.search = MagicMock(return_value=mock_results_list)
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
assert len(results) == 1
assert results[0]["content"] == "Result 3"
mock_memory_client.search = MagicMock(return_value="unexpected format")
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
assert len(results) == 0
@pytest.mark.parametrize("run_id", [None, "", "test-123", "a" * 256])
@pytest.mark.v2_api
def test_run_id_edge_cases(mock_mem0_memory_client, run_id):
"""Test edge cases for run_id parameter"""
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": run_id,
},
}
)
if run_id == "":
mem0_storage = Mem0Storage(type="short_term", crew=crew)
assert mem0_storage.run_id == ""
mock_mem0_memory_client.add = MagicMock()
mem0_storage.save("test", {})
assert "run_id" not in mock_mem0_memory_client.add.call_args[1]
else:
mem0_storage = Mem0Storage(type="short_term", crew=crew)
assert mem0_storage.run_id == run_id
if run_id is not None:
mock_mem0_memory_client.add = MagicMock()
mem0_storage.save("test", {})
assert mock_mem0_memory_client.add.call_args[1].get("run_id") == run_id
def test_invalid_version_handling(mock_mem0_memory_client):
"""Test handling of invalid version"""
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": "invalid",
},
}
)
with pytest.raises(ValueError, match="Unsupported version"):
Mem0Storage(type="short_term", crew=crew)
def test_invalid_run_id_type(mock_mem0_memory_client):
"""Test handling of invalid run_id type"""
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": 123, # Not a string
},
}
)
with pytest.raises(ValueError, match="run_id must be a string"):
Mem0Storage(type="short_term", crew=crew)