From 85cd4ebc5f94401b177de22b2d26d821cf9e4d2b Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Sat, 12 Apr 2025 01:46:56 +0000
Subject: [PATCH] Fix #2591: Fix KeyError in contextual memory with Mem0
 provider

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 .../memory/contextual/contextual_memory.py    |  4 +-
 .../contextual/test_contextual_memory.py      | 45 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 tests/memory/contextual/test_contextual_memory.py

diff --git a/src/crewai/memory/contextual/contextual_memory.py b/src/crewai/memory/contextual/contextual_memory.py
index 23e79986f..42d77cb40 100644
--- a/src/crewai/memory/contextual/contextual_memory.py
+++ b/src/crewai/memory/contextual/contextual_memory.py
@@ -56,7 +56,7 @@ class ContextualMemory:
         stm_results = self.stm.search(query)
         formatted_results = "\n".join(
             [
-                f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
+                f"- {result['context']}"
                 for result in stm_results
             ]
         )
@@ -89,7 +89,7 @@ class ContextualMemory:
         em_results = self.em.search(query)
         formatted_results = "\n".join(
             [
-                f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
+                f"- {result['context']}"
                 for result in em_results
             ]  # type: ignore #  Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
         )
diff --git a/tests/memory/contextual/test_contextual_memory.py b/tests/memory/contextual/test_contextual_memory.py
new file mode 100644
index 000000000..bded48760
--- /dev/null
+++ b/tests/memory/contextual/test_contextual_memory.py
@@ -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