mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
Fix #2591: Fix KeyError in contextual memory with Mem0 provider
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -56,7 +56,7 @@ class ContextualMemory:
|
|||||||
stm_results = self.stm.search(query)
|
stm_results = self.stm.search(query)
|
||||||
formatted_results = "\n".join(
|
formatted_results = "\n".join(
|
||||||
[
|
[
|
||||||
f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
|
f"- {result['context']}"
|
||||||
for result in stm_results
|
for result in stm_results
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
@@ -89,7 +89,7 @@ class ContextualMemory:
|
|||||||
em_results = self.em.search(query)
|
em_results = self.em.search(query)
|
||||||
formatted_results = "\n".join(
|
formatted_results = "\n".join(
|
||||||
[
|
[
|
||||||
f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
|
f"- {result['context']}"
|
||||||
for result in em_results
|
for result in em_results
|
||||||
] # type: ignore # Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
|
] # type: ignore # Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
|
||||||
)
|
)
|
||||||
|
|||||||
45
tests/memory/contextual/test_contextual_memory.py
Normal file
45
tests/memory/contextual/test_contextual_memory.py
Normal file
@@ -0,0 +1,45 @@
|
|||||||
|
import pytest
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
from crewai.memory.contextual.contextual_memory import ContextualMemory
|
||||||
|
|
||||||
|
|
||||||
|
def test_contextual_memory_with_mem0_provider():
|
||||||
|
"""Test that contextual memory properly handles Mem0 results."""
|
||||||
|
stm_mock = MagicMock()
|
||||||
|
ltm_mock = MagicMock()
|
||||||
|
em_mock = MagicMock()
|
||||||
|
um_mock = MagicMock()
|
||||||
|
exm_mock = MagicMock() # External memory mock
|
||||||
|
|
||||||
|
mock_result = {
|
||||||
|
'id': 'test-id',
|
||||||
|
'metadata': 'some metadata',
|
||||||
|
'context': 'test context data',
|
||||||
|
'score': 0.95
|
||||||
|
}
|
||||||
|
|
||||||
|
stm_mock.search.return_value = [mock_result]
|
||||||
|
em_mock.search.return_value = [mock_result]
|
||||||
|
um_mock.search.return_value = [{'memory': 'user memory'}] # User memory has different structure
|
||||||
|
exm_mock.search.return_value = [{'memory': 'external memory'}] # External memory structure
|
||||||
|
|
||||||
|
memory_config = {"provider": "mem0"}
|
||||||
|
context_memory = ContextualMemory(
|
||||||
|
memory_config=memory_config,
|
||||||
|
stm=stm_mock,
|
||||||
|
ltm=ltm_mock,
|
||||||
|
em=em_mock,
|
||||||
|
um=um_mock,
|
||||||
|
exm=exm_mock
|
||||||
|
)
|
||||||
|
|
||||||
|
result = context_memory._fetch_stm_context("test query")
|
||||||
|
|
||||||
|
stm_mock.search.assert_called_once_with("test query")
|
||||||
|
|
||||||
|
assert "test context data" in result
|
||||||
|
|
||||||
|
entity_result = context_memory._fetch_entity_context("test query")
|
||||||
|
em_mock.search.assert_called_once_with("test query")
|
||||||
|
assert "test context data" in entity_result
|
||||||
Reference in New Issue
Block a user