diff --git a/mem0/memory/main.py b/mem0/memory/main.py index 91ed15810..17b13bcba 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -1081,13 +1081,15 @@ class Memory(MemoryBase): return original_memories - def update(self, memory_id, data): + def update(self, memory_id, data, metadata: Optional[Dict[str, Any]] = None): """ Update a memory by ID. Args: memory_id (str): ID of the memory to update. data (str): New content to update the memory with. + metadata (dict, optional): Additional metadata to update. Existing metadata fields + not specified here will be preserved. Defaults to None. Returns: dict: Success message indicating the memory was updated. @@ -1095,12 +1097,14 @@ class Memory(MemoryBase): Example: >>> m.update(memory_id="mem_123", data="Likes to play tennis on weekends") {'message': 'Memory updated successfully!'} + >>> m.update(memory_id="mem_123", data="Likes tennis", metadata={"category": "sports"}) + {'message': 'Memory updated successfully!'} """ capture_event("mem0.update", self, {"memory_id": memory_id, "sync_type": "sync"}) existing_embeddings = {data: self.embedding_model.embed(data, "update")} - self._update_memory(memory_id, data, existing_embeddings) + self._update_memory(memory_id, data, existing_embeddings, metadata) return {"message": "Memory updated successfully!"} def delete(self, memory_id): @@ -2201,13 +2205,15 @@ class AsyncMemory(MemoryBase): return original_memories - async def update(self, memory_id, data): + async def update(self, memory_id, data, metadata: Optional[Dict[str, Any]] = None): """ Update a memory by ID asynchronously. Args: memory_id (str): ID of the memory to update. data (str): New content to update the memory with. + metadata (dict, optional): Additional metadata to update. Existing metadata fields + not specified here will be preserved. Defaults to None. Returns: dict: Success message indicating the memory was updated. @@ -2215,13 +2221,15 @@ class AsyncMemory(MemoryBase): Example: >>> await m.update(memory_id="mem_123", data="Likes to play tennis on weekends") {'message': 'Memory updated successfully!'} + >>> await m.update(memory_id="mem_123", data="Likes tennis", metadata={"category": "sports"}) + {'message': 'Memory updated successfully!'} """ capture_event("mem0.update", self, {"memory_id": memory_id, "sync_type": "async"}) embeddings = await asyncio.to_thread(self.embedding_model.embed, data, "update") existing_embeddings = {data: embeddings} - await self._update_memory(memory_id, data, existing_embeddings) + await self._update_memory(memory_id, data, existing_embeddings, metadata) return {"message": "Memory updated successfully!"} async def delete(self, memory_id): diff --git a/tests/memory/test_main.py b/tests/memory/test_main.py index a7555292a..fba12feef 100644 --- a/tests/memory/test_main.py +++ b/tests/memory/test_main.py @@ -1,6 +1,6 @@ import logging from datetime import datetime, timezone -from unittest.mock import MagicMock +from unittest.mock import MagicMock, Mock import pytest @@ -79,6 +79,57 @@ class TestAddToVectorStoreErrors: assert "Empty response from LLM, no memories to extract" in caplog.text +class TestAsyncUpdate: + @pytest.fixture + def mock_async_memory(self, mocker): + """Fixture for AsyncMemory with mocker-based mocks""" + _setup_mocks(mocker) + memory = AsyncMemory() + return memory + + @pytest.mark.asyncio + async def test_async_update_without_metadata(self, mock_async_memory, mocker): + """Test async update passes None metadata by default""" + mock_async_memory.embedding_model = Mock() + mock_async_memory.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3]) + mock_async_memory._update_memory = mocker.AsyncMock() + + result = await mock_async_memory.update("test_id", "Updated memory") + + mock_async_memory._update_memory.assert_called_once_with( + "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}, None + ) + assert result["message"] == "Memory updated successfully!" + + @pytest.mark.asyncio + async def test_async_update_with_metadata(self, mock_async_memory, mocker): + """Test async update correctly forwards metadata""" + mock_async_memory.embedding_model = Mock() + mock_async_memory.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3]) + mock_async_memory._update_memory = mocker.AsyncMock() + metadata = {"category": "sports", "priority": "high"} + + result = await mock_async_memory.update("test_id", "Updated memory", metadata=metadata) + + mock_async_memory._update_memory.assert_called_once_with( + "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}, metadata + ) + assert result["message"] == "Memory updated successfully!" + + @pytest.mark.asyncio + async def test_async_update_with_empty_metadata(self, mock_async_memory, mocker): + """Test async update with empty metadata dict""" + mock_async_memory.embedding_model = Mock() + mock_async_memory.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3]) + mock_async_memory._update_memory = mocker.AsyncMock() + + await mock_async_memory.update("test_id", "Updated memory", metadata={}) + + mock_async_memory._update_memory.assert_called_once_with( + "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}, {} + ) + + @pytest.mark.asyncio class TestAsyncAddToVectorStoreErrors: @pytest.fixture diff --git a/tests/test_main.py b/tests/test_main.py index 96bcb77fc..37cb11015 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -175,12 +175,41 @@ def test_update(memory_instance): result = memory_instance.update("test_id", "Updated memory") memory_instance._update_memory.assert_called_once_with( - "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]} + "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}, None ) assert result["message"] == "Memory updated successfully!" +def test_update_with_metadata(memory_instance): + memory_instance.embedding_model = Mock() + memory_instance.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3]) + + memory_instance._update_memory = Mock() + metadata = {"category": "sports", "priority": "high"} + + result = memory_instance.update("test_id", "Updated memory", metadata=metadata) + + memory_instance._update_memory.assert_called_once_with( + "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}, metadata + ) + + assert result["message"] == "Memory updated successfully!" + + +def test_update_with_empty_metadata(memory_instance): + memory_instance.embedding_model = Mock() + memory_instance.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3]) + + memory_instance._update_memory = Mock() + + memory_instance.update("test_id", "Updated memory", metadata={}) + + memory_instance._update_memory.assert_called_once_with( + "test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}, {} + ) + + def test_delete(memory_instance): memory_instance._delete_memory = Mock()