mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
fix: handle missing parameters in Mem0Storage.search() to prevent KeyError
- Replace unsafe del statements with safe .pop() calls - Add comprehensive tests for missing parameter scenarios - Fixes issue #3266 where KeyError occurred with local_mem0_config The issue occurred when using Memory (OSS) with local_mem0_config where parameters like 'metadata', 'version', 'output_format' might not be present in the params dictionary. Using .pop() with None default safely removes parameters without raising KeyError. Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
@@ -159,9 +159,10 @@ class Mem0Storage(Storage):
|
||||
params['threshold'] = score_threshold
|
||||
|
||||
if isinstance(self.memory, Memory):
|
||||
del params["metadata"], params["version"], params['output_format']
|
||||
if params.get("run_id"):
|
||||
del params["run_id"]
|
||||
params.pop("metadata", None)
|
||||
params.pop("version", None)
|
||||
params.pop("output_format", None)
|
||||
params.pop("run_id", None)
|
||||
|
||||
results = self.memory.search(**params)
|
||||
return [r for r in results["results"]]
|
||||
|
||||
@@ -363,3 +363,55 @@ def test_search_method_with_agent_id_and_user_id():
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0]["content"] == "Result 1"
|
||||
|
||||
|
||||
def test_search_method_with_memory_oss_missing_params():
|
||||
"""Test search method handles missing parameters gracefully when using Memory (OSS)"""
|
||||
config = {
|
||||
"user_id": "test_user",
|
||||
"local_mem0_config": {
|
||||
"vector_store": {"provider": "mock_vector_store"},
|
||||
"llm": {"provider": "mock_llm"},
|
||||
}
|
||||
}
|
||||
|
||||
mock_memory = MagicMock(spec=Memory)
|
||||
mock_results = {"results": [{"score": 0.9, "content": "Result 1"}]}
|
||||
mock_memory.search = MagicMock(return_value=mock_results)
|
||||
|
||||
with patch("mem0.memory.main.Memory.from_config", return_value=mock_memory):
|
||||
mem0_storage = Mem0Storage(type="external", config=config)
|
||||
|
||||
results = mem0_storage.search("test query", limit=5, score_threshold=0.5)
|
||||
|
||||
mock_memory.search.assert_called_once_with(
|
||||
query="test query",
|
||||
limit=5,
|
||||
user_id="test_user",
|
||||
filters={"AND": [{"user_id": "test_user"}]},
|
||||
threshold=0.5
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["content"] == "Result 1"
|
||||
|
||||
|
||||
def test_search_method_with_memory_oss_no_optional_params():
|
||||
"""Test search method works when no optional parameters are present"""
|
||||
mock_memory = MagicMock(spec=Memory)
|
||||
mock_results = {"results": []}
|
||||
mock_memory.search = MagicMock(return_value=mock_results)
|
||||
|
||||
with patch("mem0.memory.main.Memory", return_value=mock_memory):
|
||||
mem0_storage = Mem0Storage(type="external", config={})
|
||||
|
||||
results = mem0_storage.search("test query")
|
||||
|
||||
mock_memory.search.assert_called_once_with(
|
||||
query="test query",
|
||||
limit=3,
|
||||
filters={},
|
||||
threshold=0.35
|
||||
)
|
||||
|
||||
assert len(results) == 0
|
||||
|
||||
Reference in New Issue
Block a user