diff --git a/mem0/graphs/neptune/base.py b/mem0/graphs/neptune/base.py index 7cb643d2f..458c51755 100644 --- a/mem0/graphs/neptune/base.py +++ b/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}") diff --git a/mem0/memory/graph_memory.py b/mem0/memory/graph_memory.py index 0b911dff5..ec1cff960 100644 --- a/mem0/memory/graph_memory.py +++ b/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( diff --git a/mem0/memory/kuzu_memory.py b/mem0/memory/kuzu_memory.py index a56778664..a025aec36 100644 --- a/mem0/memory/kuzu_memory.py +++ b/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( diff --git a/mem0/memory/memgraph_memory.py b/mem0/memory/memgraph_memory.py index 14a2e5119..84b6d0ca8 100644 --- a/mem0/memory/memgraph_memory.py +++ b/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: diff --git a/tests/memory/test_kuzu.py b/tests/memory/test_kuzu.py index 5cb3f2bce..14ed1d175 100644 --- a/tests/memory/test_kuzu.py +++ b/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( """ diff --git a/tests/memory/test_memgraph_memory.py b/tests/memory/test_memgraph_memory.py index 4bf51a400..58d6acce4 100644 --- a/tests/memory/test_memgraph_memory.py +++ b/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}