fix: key error when llm omits entities key tool call (#4313)

mem0/graphs/neptune/base.py

@@ -97,7 +97,7 @@ class NeptuneBase(ABC):
             for tool_call in search_results["tool_calls"]:
                 if tool_call["name"] != "extract_entities":
                     continue
-                for item in tool_call["arguments"]["entities"]:
+                for item in tool_call.get("arguments", {}).get("entities", []):
                     entity_type_map[item["entity"]] = item["entity_type"]
         except Exception as e:
             logger.exception(
@@ -144,7 +144,7 @@ class NeptuneBase(ABC):
 
         entities = []
         if extracted_entities["tool_calls"]:
-            entities = extracted_entities["tool_calls"][0]["arguments"]["entities"]
+            entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
 
         entities = self._remove_spaces_from_entities(entities)
         logger.debug(f"Extracted entities: {entities}")

mem0/memory/graph_memory.py

@@ -215,7 +215,7 @@ class MemoryGraph:
             for tool_call in search_results["tool_calls"]:
                 if tool_call["name"] != "extract_entities":
                     continue
-                for item in tool_call["arguments"]["entities"]:
+                for item in tool_call.get("arguments", {}).get("entities", []):
                     entity_type_map[item["entity"]] = item["entity_type"]
         except Exception as e:
             logger.exception(

mem0/memory/kuzu_memory.py

@@ -241,7 +241,7 @@ class MemoryGraph:
             for tool_call in search_results["tool_calls"]:
                 if tool_call["name"] != "extract_entities":
                     continue
-                for item in tool_call["arguments"]["entities"]:
+                for item in tool_call.get("arguments", {}).get("entities", []):
                     entity_type_map[item["entity"]] = item["entity_type"]
         except Exception as e:
             logger.exception(

mem0/memory/memgraph_memory.py

@@ -218,7 +218,7 @@ class MemoryGraph:
             for tool_call in search_results["tool_calls"]:
                 if tool_call["name"] != "extract_entities":
                     continue
-                for item in tool_call["arguments"]["entities"]:
+                for item in tool_call.get("arguments", {}).get("entities", []):
                     if "entity" in item and "entity_type" in item:
                         entity_type_map[item["entity"]] = item["entity_type"]
         except Exception as e:

tests/memory/test_kuzu.py

@@ -1,6 +1,8 @@
+from unittest.mock import MagicMock, Mock, patch
+
 import numpy as np
 import pytest
-from unittest.mock import Mock, patch
+
 from mem0.memory.kuzu_memory import MemoryGraph
 
 
@@ -190,6 +192,48 @@ class TestKuzu:
         assert get_node_count(kuzu_memory) == 0
         assert get_edge_count(kuzu_memory) == 0
 
+def _make_kuzu_instance():
+    with patch.object(MemoryGraph, "__init__", return_value=None):
+        instance = MemoryGraph.__new__(MemoryGraph)
+        instance.llm_provider = "openai"
+        instance.llm = MagicMock()
+        instance.embedding_model = MagicMock()
+        instance.config = MagicMock()
+        instance.config.graph_store.custom_prompt = None
+        return instance
+
+
+class TestRetrieveNodesFromData:
+    """Tests for _retrieve_nodes_from_data in KuzuMemoryGraph."""
+
+    def test_missing_entities_key_returns_empty(self):
+        """LLM returns extract_entities tool call without 'entities' key — should not crash.
+        Reproduces the exact scenario from issue #4238."""
+        instance = _make_kuzu_instance()
+        instance.llm.generate_response.return_value = {
+            "tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
+        }
+        result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
+        assert result == {}
+
+    def test_normal_entities_extracted(self):
+        instance = _make_kuzu_instance()
+        instance.llm.generate_response.return_value = {
+            "tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
+                {"entity": "Alice", "entity_type": "person"},
+                {"entity": "hiking", "entity_type": "activity"},
+            ]}}]
+        }
+        result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
+        assert result == {"alice": "person", "hiking": "activity"}
+
+    def test_none_tool_calls_returns_empty(self):
+        instance = _make_kuzu_instance()
+        instance.llm.generate_response.return_value = {"tool_calls": None}
+        result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
+        assert result == {}
+
+
 def get_node_count(kuzu_memory):
     results = kuzu_memory.kuzu_execute(
         """

tests/memory/test_memgraph_memory.py

@@ -10,6 +10,7 @@ patch.dict("sys.modules", {
 }).start()
 
 from mem0.memory.memgraph_memory import MemoryGraph as MemgraphMemoryGraph  # noqa: E402
+
 MemoryGraph = MemgraphMemoryGraph
 
 
@@ -54,6 +55,16 @@ class TestRetrieveNodesFromData:
         assert "relu" in result
         assert "task" not in result
 
+    def test_missing_entities_key_returns_empty(self):
+        """LLM returns extract_entities tool call without 'entities' key — should not crash.
+        Reproduces the exact scenario from issue #4238."""
+        instance = _make_instance()
+        instance.llm.generate_response.return_value = {
+            "tool_calls": [{"name": "extract_entities", "arguments": {"text": "Hello."}}]
+        }
+        result = instance._retrieve_nodes_from_data("Hello.", {"user_id": "u1"})
+        assert result == {}
+
     def test_none_tool_calls_returns_empty(self):
         instance = _make_instance()
         instance.llm.generate_response.return_value = {"tool_calls": None}