mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
Fix Memory OSS compatibility and implement code review suggestions
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user