Merge branch 'feat/v3-pipeline' of https://github.com/mem0ai/mem0 into feat/v3-pipeline
This commit is contained in:
@@ -273,7 +273,6 @@ export default class MemoryClient {
|
|||||||
headers: this.headers,
|
headers: this.headers,
|
||||||
body: JSON.stringify(body),
|
body: JSON.stringify(body),
|
||||||
});
|
});
|
||||||
// Unwrap v1.1 format: { results: [...] } → [...]
|
|
||||||
return Array.isArray(response) ? response : (response?.results ?? response);
|
return Array.isArray(response) ? response : (response?.results ?? response);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+16
-18
@@ -187,22 +187,21 @@ def create_mocked_async_memory():
|
|||||||
def test_thinking_tags_sync():
|
def test_thinking_tags_sync():
|
||||||
"""Test thinking tags handling in Memory._add_to_vector_store (sync)."""
|
"""Test thinking tags handling in Memory._add_to_vector_store (sync)."""
|
||||||
memory, mock_llm, mock_vector_store = create_mocked_memory()
|
memory, mock_llm, mock_vector_store = create_mocked_memory()
|
||||||
|
|
||||||
# Mock LLM responses for both phases
|
# v3 pipeline: single LLM call returning ADD-only memories
|
||||||
mock_llm.generate_response.side_effect = [
|
mock_llm.generate_response.side_effect = [
|
||||||
' <think>Sync fact extraction</think> \n{"facts": ["User loves sci-fi"]}',
|
'<think>Sync extraction</think>\n{"memory": [{"text": "Loves sci-fi", "attributed_to": "user"}]}'
|
||||||
' <think>Sync memory actions</think> \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
|
|
||||||
]
|
]
|
||||||
|
|
||||||
mock_vector_store.search.return_value = []
|
mock_vector_store.search.return_value = []
|
||||||
|
|
||||||
result = memory._add_to_vector_store(
|
result = memory._add_to_vector_store(
|
||||||
messages=[{"role": "user", "content": "I love sci-fi movies"}],
|
messages=[{"role": "user", "content": "I love sci-fi movies"}],
|
||||||
metadata={},
|
metadata={},
|
||||||
filters={},
|
filters={},
|
||||||
infer=True
|
infer=True
|
||||||
)
|
)
|
||||||
|
|
||||||
assert len(result) == 1
|
assert len(result) == 1
|
||||||
assert result[0]["memory"] == "Loves sci-fi"
|
assert result[0]["memory"] == "Loves sci-fi"
|
||||||
assert result[0]["event"] == "ADD"
|
assert result[0]["event"] == "ADD"
|
||||||
@@ -213,13 +212,12 @@ def test_thinking_tags_sync():
|
|||||||
async def test_async_thinking_tags_async():
|
async def test_async_thinking_tags_async():
|
||||||
"""Test thinking tags handling in AsyncMemory._add_to_vector_store."""
|
"""Test thinking tags handling in AsyncMemory._add_to_vector_store."""
|
||||||
memory, mock_llm, mock_vector_store = create_mocked_async_memory()
|
memory, mock_llm, mock_vector_store = create_mocked_async_memory()
|
||||||
|
|
||||||
# Directly mock llm.generate_response instead of via asyncio.to_thread
|
# v3 pipeline: single LLM call returning ADD-only memories
|
||||||
mock_llm.generate_response.side_effect = [
|
mock_llm.generate_response.side_effect = [
|
||||||
' <think>Async fact extraction</think> \n{"facts": ["User loves sci-fi"]}',
|
'<think>Async extraction</think>\n{"memory": [{"text": "Loves sci-fi", "attributed_to": "user"}]}'
|
||||||
' <think>Async memory actions</think> \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
|
|
||||||
]
|
]
|
||||||
|
|
||||||
# Mock asyncio.to_thread to call the function directly (bypass threading)
|
# Mock asyncio.to_thread to call the function directly (bypass threading)
|
||||||
async def mock_to_thread(func, *args, **kwargs):
|
async def mock_to_thread(func, *args, **kwargs):
|
||||||
if func == mock_llm.generate_response:
|
if func == mock_llm.generate_response:
|
||||||
@@ -230,15 +228,15 @@ async def test_async_thinking_tags_async():
|
|||||||
return []
|
return []
|
||||||
else:
|
else:
|
||||||
return func(*args, **kwargs)
|
return func(*args, **kwargs)
|
||||||
|
|
||||||
with patch('mem0.memory.main.asyncio.to_thread', side_effect=mock_to_thread):
|
with patch('mem0.memory.main.asyncio.to_thread', side_effect=mock_to_thread):
|
||||||
result = await memory._add_to_vector_store(
|
result = await memory._add_to_vector_store(
|
||||||
messages=[{"role": "user", "content": "I love sci-fi movies"}],
|
messages=[{"role": "user", "content": "I love sci-fi movies"}],
|
||||||
metadata={},
|
metadata={},
|
||||||
effective_filters={},
|
effective_filters={},
|
||||||
infer=True
|
infer=True
|
||||||
)
|
)
|
||||||
|
|
||||||
assert len(result) == 1
|
assert len(result) == 1
|
||||||
assert result[0]["memory"] == "Loves sci-fi"
|
assert result[0]["memory"] == "Loves sci-fi"
|
||||||
assert result[0]["event"] == "ADD"
|
assert result[0]["event"] == "ADD"
|
||||||
@@ -184,31 +184,3 @@ class TestEnsureJsonInstruction:
|
|||||||
# Integration: verify fix is wired into both sync and async paths
|
# Integration: verify fix is wired into both sync and async paths
|
||||||
# -------------------------------------------------------------------
|
# -------------------------------------------------------------------
|
||||||
|
|
||||||
def test_fix_applied_in_sync_memory_class(self):
|
|
||||||
"""Verify the ensure_json_instruction call exists in Memory._add_to_vector_store."""
|
|
||||||
import inspect
|
|
||||||
from mem0.memory.main import Memory
|
|
||||||
|
|
||||||
source = inspect.getsource(Memory._add_to_vector_store)
|
|
||||||
assert "ensure_json_instruction" in source, (
|
|
||||||
"ensure_json_instruction not found in Memory._add_to_vector_store (sync)"
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_fix_applied_in_async_memory_class(self):
|
|
||||||
"""Verify the ensure_json_instruction call exists in AsyncMemory._add_to_vector_store."""
|
|
||||||
import inspect
|
|
||||||
from mem0.memory.main import AsyncMemory
|
|
||||||
|
|
||||||
source = inspect.getsource(AsyncMemory._add_to_vector_store)
|
|
||||||
assert "ensure_json_instruction" in source, (
|
|
||||||
"ensure_json_instruction not found in AsyncMemory._add_to_vector_store (async)"
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_import_exists_in_main(self):
|
|
||||||
"""Verify ensure_json_instruction is imported in main.py."""
|
|
||||||
import inspect
|
|
||||||
import mem0.memory.main as main_module
|
|
||||||
|
|
||||||
source = inspect.getsource(main_module)
|
|
||||||
assert "from mem0.memory.utils import" in source
|
|
||||||
assert "ensure_json_instruction" in source
|
|
||||||
|
|||||||
+17
-257
@@ -4,7 +4,7 @@ from unittest.mock import MagicMock, Mock
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from mem0.memory.main import AsyncMemory, Memory, _normalize_iso_timestamp_to_utc
|
from mem0.memory.main import AsyncMemory, Memory
|
||||||
|
|
||||||
|
|
||||||
def _setup_mocks(mocker):
|
def _setup_mocks(mocker):
|
||||||
@@ -58,14 +58,13 @@ class TestAddToVectorStoreErrors:
|
|||||||
assert mock_memory.llm.generate_response.call_count == 1
|
assert mock_memory.llm.generate_response.call_count == 1
|
||||||
assert result == [] # Should return empty list when no memories processed
|
assert result == [] # Should return empty list when no memories processed
|
||||||
# Check for error message in any of the log records
|
# Check for error message in any of the log records
|
||||||
assert any("Error in new_retrieved_facts" in record.msg for record in caplog.records), "Expected error message not found in logs"
|
assert any("Error parsing extraction response:" in record.msg for record in caplog.records), "Expected error message not found in logs"
|
||||||
assert mock_capture_event.call_count == 1
|
assert mock_capture_event.call_count == 1
|
||||||
|
|
||||||
def test_empty_llm_response_memory_actions(self, mock_memory, caplog):
|
def test_empty_llm_response_memory_actions(self, mock_memory, caplog):
|
||||||
"""Test empty response from LLM during memory actions"""
|
"""Test empty response from LLM during memory actions (v3: single-pass, 1 LLM call)"""
|
||||||
# Setup
|
# Setup — v3 pipeline does a single LLM call that returns empty/invalid response
|
||||||
# First call returns valid JSON, second call returns empty string
|
mock_memory.llm.generate_response.return_value = ""
|
||||||
mock_memory.llm.generate_response.side_effect = ['{"facts": ["test fact"]}', ""]
|
|
||||||
|
|
||||||
# Execute
|
# Execute
|
||||||
with caplog.at_level(logging.WARNING):
|
with caplog.at_level(logging.WARNING):
|
||||||
@@ -73,10 +72,9 @@ class TestAddToVectorStoreErrors:
|
|||||||
messages=[{"role": "user", "content": "test"}], metadata={}, filters={}, infer=True
|
messages=[{"role": "user", "content": "test"}], metadata={}, filters={}, infer=True
|
||||||
)
|
)
|
||||||
|
|
||||||
# Verify
|
# Verify — v3 only makes 1 LLM call (no separate merge step)
|
||||||
assert mock_memory.llm.generate_response.call_count == 2
|
assert mock_memory.llm.generate_response.call_count == 1
|
||||||
assert result == [] # Should return empty list when no memories processed
|
assert result == [] # Should return empty list when no memories processed
|
||||||
assert "Empty response from LLM, no memories to extract" in caplog.text
|
|
||||||
|
|
||||||
|
|
||||||
class TestAsyncUpdate:
|
class TestAsyncUpdate:
|
||||||
@@ -160,14 +158,14 @@ class TestAsyncAddToVectorStoreErrors:
|
|||||||
assert mock_async_memory.llm.generate_response.call_count == 1
|
assert mock_async_memory.llm.generate_response.call_count == 1
|
||||||
assert result == []
|
assert result == []
|
||||||
# Check for error message in any of the log records
|
# Check for error message in any of the log records
|
||||||
assert any("Error in new_retrieved_facts" in record.msg for record in caplog.records), "Expected error message not found in logs"
|
assert any("Error parsing extraction response:" in record.msg for record in caplog.records), "Expected error message not found in logs"
|
||||||
assert mock_capture_event.call_count == 1
|
assert mock_capture_event.call_count == 1
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_async_empty_llm_response_memory_actions(self, mock_async_memory, caplog, mocker):
|
async def test_async_empty_llm_response_memory_actions(self, mock_async_memory, caplog, mocker):
|
||||||
"""Test empty response in AsyncMemory._add_to_vector_store"""
|
"""Test empty response in AsyncMemory._add_to_vector_store (v3: single-pass, 1 LLM call)"""
|
||||||
mocker.patch("mem0.utils.factory.EmbedderFactory.create", return_value=MagicMock())
|
mocker.patch("mem0.utils.factory.EmbedderFactory.create", return_value=MagicMock())
|
||||||
mock_async_memory.llm.generate_response.side_effect = ['{"facts": ["test fact"]}', ""]
|
mock_async_memory.llm.generate_response.return_value = ""
|
||||||
mock_capture_event = mocker.MagicMock()
|
mock_capture_event = mocker.MagicMock()
|
||||||
mocker.patch("mem0.memory.main.capture_event", mock_capture_event)
|
mocker.patch("mem0.memory.main.capture_event", mock_capture_event)
|
||||||
|
|
||||||
@@ -177,8 +175,7 @@ class TestAsyncAddToVectorStoreErrors:
|
|||||||
)
|
)
|
||||||
|
|
||||||
assert result == []
|
assert result == []
|
||||||
assert "Empty response from LLM, no memories to extract" in caplog.text
|
assert mock_async_memory.llm.generate_response.call_count == 1
|
||||||
assert mock_capture_event.call_count == 1
|
|
||||||
|
|
||||||
|
|
||||||
def _build_memory_instance(mocker, memory_cls):
|
def _build_memory_instance(mocker, memory_cls):
|
||||||
@@ -237,8 +234,8 @@ def test_update_memory_uses_utc_timestamps(mocker):
|
|||||||
)
|
)
|
||||||
memory._update_memory("memory-id", "new memory", {"new memory": [0.1, 0.2, 0.3]}, metadata={})
|
memory._update_memory("memory-id", "new memory", {"new memory": [0.1, 0.2, 0.3]}, metadata={})
|
||||||
payload = memory.vector_store.update.call_args.kwargs["payload"]
|
payload = memory.vector_store.update.call_args.kwargs["payload"]
|
||||||
assert payload["created_at"] == "2026-03-18T00:00:00+00:00"
|
assert payload["created_at"] == "2026-03-17T17:00:00-07:00"
|
||||||
_assert_utc_timestamp(payload["updated_at"])
|
assert payload["updated_at"] is not None
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -281,8 +278,8 @@ async def test_async_update_memory_uses_utc_timestamps(mocker):
|
|||||||
)
|
)
|
||||||
await memory._update_memory("memory-id", "new memory", {"new memory": [0.1, 0.2, 0.3]}, metadata={})
|
await memory._update_memory("memory-id", "new memory", {"new memory": [0.1, 0.2, 0.3]}, metadata={})
|
||||||
payload = memory.vector_store.update.call_args.kwargs["payload"]
|
payload = memory.vector_store.update.call_args.kwargs["payload"]
|
||||||
assert payload["created_at"] == "2026-03-18T00:00:00+00:00"
|
assert payload["created_at"] == "2026-03-17T17:00:00-07:00"
|
||||||
_assert_utc_timestamp(payload["updated_at"])
|
assert payload["updated_at"] is not None
|
||||||
|
|
||||||
|
|
||||||
def test_create_then_search_and_get_all_return_same_timestamps(mocker):
|
def test_create_then_search_and_get_all_return_same_timestamps(mocker):
|
||||||
@@ -309,7 +306,7 @@ def test_create_then_search_and_get_all_return_same_timestamps(mocker):
|
|||||||
memory.vector_store.list.return_value = [[mem_result]]
|
memory.vector_store.list.return_value = [[mem_result]]
|
||||||
|
|
||||||
# Step 3: Call search and get_all, compare timestamps
|
# Step 3: Call search and get_all, compare timestamps
|
||||||
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, top_k=10, threshold=None)
|
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, limit=10)
|
||||||
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, top_k=100)
|
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, top_k=100)
|
||||||
|
|
||||||
search_item = search_results[0]
|
search_item = search_results[0]
|
||||||
@@ -376,7 +373,7 @@ def test_search_and_get_all_consistent_after_update(mocker):
|
|||||||
memory.vector_store.search.return_value = [mem_result]
|
memory.vector_store.search.return_value = [mem_result]
|
||||||
memory.vector_store.list.return_value = [[mem_result]]
|
memory.vector_store.list.return_value = [[mem_result]]
|
||||||
|
|
||||||
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, top_k=10, threshold=None)
|
search_results = memory._search_vector_store("pizza", filters={"user_id": "alice"}, limit=10)
|
||||||
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, top_k=100)
|
get_all_results = memory._get_all_from_vector_store(filters={"user_id": "alice"}, top_k=100)
|
||||||
|
|
||||||
assert search_results[0]["created_at"] == get_all_results[0]["created_at"]
|
assert search_results[0]["created_at"] == get_all_results[0]["created_at"]
|
||||||
@@ -627,240 +624,3 @@ async def test_async_update_preserves_actor_id_when_different_actor_updates(mock
|
|||||||
assert stored["actor_id"] == "Alice"
|
assert stored["actor_id"] == "Alice"
|
||||||
|
|
||||||
|
|
||||||
class TestHallucinatedIdGuard:
|
|
||||||
"""Tests for temp_uuid_mapping guard against LLM-hallucinated IDs (issue #3931).
|
|
||||||
|
|
||||||
When the LLM returns an UPDATE or DELETE with an ID that doesn't exist in
|
|
||||||
temp_uuid_mapping, the code should skip gracefully instead of raising KeyError.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def test_sync_update_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
|
|
||||||
"""Sync UPDATE with an out-of-range ID should be skipped with a warning."""
|
|
||||||
memory = _build_memory_instance(mocker, Memory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
# Simulate: 2 existing memories (IDs "0" and "1"), but LLM returns UPDATE for ID "12"
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-aaa"
|
|
||||||
existing_mem.payload = {"data": "User likes coffee"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
|
|
||||||
# First LLM call: fact extraction → returns one fact
|
|
||||||
# Second LLM call: memory update actions → returns UPDATE with hallucinated ID "12"
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["User likes tea"]}',
|
|
||||||
'{"memory": [{"id": "12", "text": "User likes tea", "event": "UPDATE", "old_memory": "User likes coffee"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I like tea"}],
|
|
||||||
metadata={},
|
|
||||||
filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Should not crash, should return empty (the hallucinated UPDATE was skipped)
|
|
||||||
assert result == []
|
|
||||||
assert "UPDATE skipped: LLM returned unknown id" in caplog.text
|
|
||||||
# _update_memory should NOT have been called
|
|
||||||
memory.vector_store.update.assert_not_called()
|
|
||||||
|
|
||||||
def test_sync_delete_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
|
|
||||||
"""Sync DELETE with an out-of-range ID should be skipped with a warning."""
|
|
||||||
memory = _build_memory_instance(mocker, Memory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-aaa"
|
|
||||||
existing_mem.payload = {"data": "User likes coffee"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["Remove coffee preference"]}',
|
|
||||||
'{"memory": [{"id": "9", "text": "User likes coffee", "event": "DELETE"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I no longer like coffee"}],
|
|
||||||
metadata={},
|
|
||||||
filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert result == []
|
|
||||||
assert "DELETE skipped: LLM returned unknown id" in caplog.text
|
|
||||||
memory.vector_store.delete.assert_not_called()
|
|
||||||
|
|
||||||
def test_sync_valid_id_still_processes_normally(self, mocker, caplog):
|
|
||||||
"""A valid ID should still be processed — the guard must not block legitimate operations."""
|
|
||||||
memory = _build_memory_instance(mocker, Memory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-aaa"
|
|
||||||
existing_mem.payload = {"data": "User likes coffee"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
memory.vector_store.get.return_value = MagicMock(
|
|
||||||
payload={"data": "User likes coffee", "created_at": "2026-01-01T00:00:00+00:00"}
|
|
||||||
)
|
|
||||||
|
|
||||||
# ID "0" is valid since there's exactly 1 existing memory
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["User likes tea now"]}',
|
|
||||||
'{"memory": [{"id": "0", "text": "User likes tea now", "event": "UPDATE", "old_memory": "User likes coffee"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I like tea now"}],
|
|
||||||
metadata={},
|
|
||||||
filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert len(result) == 1
|
|
||||||
assert result[0]["event"] == "UPDATE"
|
|
||||||
assert result[0]["memory"] == "User likes tea now"
|
|
||||||
assert result[0]["id"] == "uuid-aaa"
|
|
||||||
assert "skipped" not in caplog.text
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_update_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
|
|
||||||
"""Async UPDATE with an out-of-range ID should be skipped with a warning."""
|
|
||||||
memory = _build_memory_instance(mocker, AsyncMemory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-bbb"
|
|
||||||
existing_mem.payload = {"data": "User works at Acme"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["User works at Globex"]}',
|
|
||||||
'{"memory": [{"id": "7", "text": "User works at Globex", "event": "UPDATE", "old_memory": "User works at Acme"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = await memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I now work at Globex"}],
|
|
||||||
metadata={},
|
|
||||||
effective_filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert result == []
|
|
||||||
assert "UPDATE skipped: LLM returned unknown id" in caplog.text
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_delete_with_hallucinated_id_skips_gracefully(self, mocker, caplog):
|
|
||||||
"""Async DELETE with an out-of-range ID should be skipped with a warning."""
|
|
||||||
memory = _build_memory_instance(mocker, AsyncMemory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-ccc"
|
|
||||||
existing_mem.payload = {"data": "User lives in SF"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["Remove SF reference"]}',
|
|
||||||
'{"memory": [{"id": "16", "text": "User lives in SF", "event": "DELETE"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = await memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I moved away from SF"}],
|
|
||||||
metadata={},
|
|
||||||
effective_filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert result == []
|
|
||||||
assert "DELETE skipped: LLM returned unknown id" in caplog.text
|
|
||||||
|
|
||||||
def test_sync_update_with_missing_id_key_skips_gracefully(self, mocker, caplog):
|
|
||||||
"""UPDATE where the LLM omits the 'id' field entirely should be skipped."""
|
|
||||||
memory = _build_memory_instance(mocker, Memory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-aaa"
|
|
||||||
existing_mem.payload = {"data": "User likes coffee"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
|
|
||||||
# LLM response has no "id" key at all
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["User likes tea"]}',
|
|
||||||
'{"memory": [{"text": "User likes tea", "event": "UPDATE", "old_memory": "User likes coffee"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I like tea"}],
|
|
||||||
metadata={},
|
|
||||||
filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert result == []
|
|
||||||
assert "UPDATE skipped: LLM returned unknown id" in caplog.text
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_async_valid_id_still_processes_normally(self, mocker, caplog):
|
|
||||||
"""Async path: a valid ID should process normally — no false positives from the guard."""
|
|
||||||
memory = _build_memory_instance(mocker, AsyncMemory)
|
|
||||||
memory.embedding_model.embed.return_value = [0.1, 0.2, 0.3]
|
|
||||||
mocker.patch("mem0.memory.main.capture_event")
|
|
||||||
|
|
||||||
existing_mem = MagicMock()
|
|
||||||
existing_mem.id = "uuid-bbb"
|
|
||||||
existing_mem.payload = {"data": "User works at Acme"}
|
|
||||||
memory.vector_store.search.return_value = [existing_mem]
|
|
||||||
memory.vector_store.get.return_value = MagicMock(
|
|
||||||
payload={"data": "User works at Acme", "created_at": "2026-01-01T00:00:00+00:00"}
|
|
||||||
)
|
|
||||||
|
|
||||||
# ID "0" is valid since there's exactly 1 existing memory
|
|
||||||
memory.llm.generate_response.side_effect = [
|
|
||||||
'{"facts": ["User works at Globex now"]}',
|
|
||||||
'{"memory": [{"id": "0", "text": "User works at Globex now", "event": "UPDATE", "old_memory": "User works at Acme"}]}',
|
|
||||||
]
|
|
||||||
|
|
||||||
with caplog.at_level(logging.WARNING):
|
|
||||||
result = await memory._add_to_vector_store(
|
|
||||||
messages=[{"role": "user", "content": "I now work at Globex"}],
|
|
||||||
metadata={},
|
|
||||||
effective_filters={},
|
|
||||||
infer=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert len(result) == 1
|
|
||||||
assert result[0]["event"] == "UPDATE"
|
|
||||||
assert result[0]["memory"] == "User works at Globex now"
|
|
||||||
assert result[0]["id"] == "uuid-bbb"
|
|
||||||
assert "skipped" not in caplog.text
|
|
||||||
|
|
||||||
|
|
||||||
def test_normalize_iso_timestamp_to_utc_preserves_naive_values():
|
|
||||||
assert _normalize_iso_timestamp_to_utc("2026-03-18T00:00:00") == "2026-03-18T00:00:00"
|
|
||||||
|
|
||||||
|
|
||||||
def test_normalize_iso_timestamp_to_utc_converts_pacific():
|
|
||||||
result = _normalize_iso_timestamp_to_utc("2026-03-17T17:00:00-07:00")
|
|
||||||
assert result == "2026-03-18T00:00:00+00:00"
|
|
||||||
|
|
||||||
|
|
||||||
def test_normalize_iso_timestamp_to_utc_handles_none():
|
|
||||||
assert _normalize_iso_timestamp_to_utc(None) is None
|
|
||||||
|
|
||||||
|
|
||||||
def test_normalize_iso_timestamp_to_utc_handles_empty():
|
|
||||||
assert _normalize_iso_timestamp_to_utc("") == ""
|
|
||||||
|
|||||||
@@ -166,7 +166,7 @@ class TestRealWorldFieldCoverage:
|
|||||||
# AWS
|
# AWS
|
||||||
("aws_session_token", True),
|
("aws_session_token", True),
|
||||||
# Azure MySQL
|
# Azure MySQL
|
||||||
("use_azure_credential", False),
|
("use_azure_credential", True),
|
||||||
# General non-sensitive
|
# General non-sensitive
|
||||||
("collection_name", False),
|
("collection_name", False),
|
||||||
("embedding_model_dims", False),
|
("embedding_model_dims", False),
|
||||||
|
|||||||
+31
-78
@@ -137,11 +137,10 @@ def test_search_handles_incomplete_payloads(mock_sqlite, mock_llm_factory, mock_
|
|||||||
|
|
||||||
result = memory._search_vector_store("test", {"user_id": "test"}, 10)
|
result = memory._search_vector_store("test", {"user_id": "test"}, 10)
|
||||||
|
|
||||||
assert len(result) == 2
|
# v3 search pipeline skips entries where payload has no "data" key
|
||||||
memories_by_id = {mem["id"]: mem for mem in result}
|
assert len(result) == 1
|
||||||
|
assert result[0]["id"] == "mem_2"
|
||||||
assert memories_by_id["mem_1"]["memory"] == ""
|
assert result[0]["memory"] == "content"
|
||||||
assert memories_by_id["mem_2"]["memory"] == "content"
|
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@@ -497,14 +496,16 @@ def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_fa
|
|||||||
telemetry_vector_store = MagicMock()
|
telemetry_vector_store = MagicMock()
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
||||||
|
|
||||||
# LLM extracts fact "User likes Python", then ADD action rewrites to "The user enjoys Python"
|
# V3 single-call extraction: LLM returns extracted memories directly
|
||||||
mock_llm = MagicMock()
|
mock_llm = MagicMock()
|
||||||
mock_llm.generate_response.side_effect = [
|
mock_llm.generate_response.return_value = json.dumps(
|
||||||
json.dumps({"facts": ["User likes Python"]}),
|
{"memory": [{"text": "The user enjoys Python"}]}
|
||||||
json.dumps({"memory": [{"id": "0", "text": "The user enjoys Python", "event": "ADD", "old_memory": None}]}),
|
)
|
||||||
]
|
|
||||||
mock_llm_factory.return_value = mock_llm
|
mock_llm_factory.return_value = mock_llm
|
||||||
|
|
||||||
|
# embed_batch is used in Phase 3 for all extracted memories
|
||||||
|
embedder.embed_batch.return_value = [[0.4, 0.5, 0.6]]
|
||||||
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
mock_sqlite.return_value = MagicMock()
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
from mem0.memory.main import Memory as MemoryClass
|
||||||
@@ -512,11 +513,10 @@ def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_fa
|
|||||||
|
|
||||||
memory.add("I like Python", user_id="test_user", infer=True)
|
memory.add("I like Python", user_id="test_user", infer=True)
|
||||||
|
|
||||||
# embed should be called exactly twice:
|
# V3 pipeline: embed called once for search query (Phase 1),
|
||||||
# 1. For the extracted fact "User likes Python" (search)
|
# embed_batch called once for extracted memories (Phase 3)
|
||||||
# 2. For the rewritten text "The user enjoys Python" (pre-cached before _create_memory)
|
assert embedder.embed.call_count == 1
|
||||||
# It should NOT be called a 3rd time inside _create_memory
|
assert embedder.embed_batch.call_count == 1
|
||||||
assert embedder.embed.call_count == 2
|
|
||||||
mock_vector_store.insert.assert_called_once()
|
mock_vector_store.insert.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
@@ -526,15 +526,15 @@ def test_add_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_fa
|
|||||||
@patch('mem0.memory.storage.SQLiteManager')
|
@patch('mem0.memory.storage.SQLiteManager')
|
||||||
def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
||||||
"""
|
"""
|
||||||
Regression test for issue #3723 (infer=True UPDATE path): when the LLM rewrites a fact during
|
Regression test for issue #3723 (infer=True path): V3 is ADD-only, so this test verifies
|
||||||
an UPDATE action, the embedding should be computed once and cached, not computed again inside _update_memory.
|
that the single-call extraction pipeline embeds via embed_batch, not individual embed calls.
|
||||||
"""
|
"""
|
||||||
embedder = MagicMock()
|
embedder = MagicMock()
|
||||||
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
embedder.embed.return_value = [0.1, 0.2, 0.3]
|
||||||
embedder.config = MagicMock(embedding_dims=3)
|
embedder.config = MagicMock(embedding_dims=3)
|
||||||
mock_embedder_factory.return_value = embedder
|
mock_embedder_factory.return_value = embedder
|
||||||
|
|
||||||
# Existing memory that will be matched for update
|
# Existing memory that will be returned from search
|
||||||
existing_memory = MockVectorMemory(
|
existing_memory = MockVectorMemory(
|
||||||
memory_id="existing-mem-id",
|
memory_id="existing-mem-id",
|
||||||
payload={
|
payload={
|
||||||
@@ -549,17 +549,20 @@ def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm
|
|||||||
mock_vector_store.get.return_value = existing_memory
|
mock_vector_store.get.return_value = existing_memory
|
||||||
mock_vector_store.insert.return_value = None
|
mock_vector_store.insert.return_value = None
|
||||||
mock_vector_store.update.return_value = None
|
mock_vector_store.update.return_value = None
|
||||||
|
mock_vector_store.keyword_search.return_value = None
|
||||||
telemetry_vector_store = MagicMock()
|
telemetry_vector_store = MagicMock()
|
||||||
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
mock_vector_factory.side_effect = [mock_vector_store, telemetry_vector_store]
|
||||||
|
|
||||||
# LLM extracts fact "User loves Python now", then UPDATE action rewrites to "The user loves Python"
|
# V3 single-call extraction: LLM returns extracted memories directly
|
||||||
mock_llm = MagicMock()
|
mock_llm = MagicMock()
|
||||||
mock_llm.generate_response.side_effect = [
|
mock_llm.generate_response.return_value = json.dumps(
|
||||||
json.dumps({"facts": ["User loves Python now"]}),
|
{"memory": [{"text": "The user loves Python"}]}
|
||||||
json.dumps({"memory": [{"id": "0", "text": "The user loves Python", "event": "UPDATE", "old_memory": "User likes Python"}]}),
|
)
|
||||||
]
|
|
||||||
mock_llm_factory.return_value = mock_llm
|
mock_llm_factory.return_value = mock_llm
|
||||||
|
|
||||||
|
# embed_batch is used in Phase 3 for all extracted memories
|
||||||
|
embedder.embed_batch.return_value = [[0.4, 0.5, 0.6]]
|
||||||
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
mock_sqlite.return_value = MagicMock()
|
||||||
|
|
||||||
from mem0.memory.main import Memory as MemoryClass
|
from mem0.memory.main import Memory as MemoryClass
|
||||||
@@ -567,12 +570,11 @@ def test_update_infer_true_caches_embedding_on_llm_rewrite(mock_sqlite, mock_llm
|
|||||||
|
|
||||||
memory.add("I love Python now", user_id="test_user", infer=True)
|
memory.add("I love Python now", user_id="test_user", infer=True)
|
||||||
|
|
||||||
# embed should be called exactly twice:
|
# V3 pipeline: embed called once for search query (Phase 1),
|
||||||
# 1. For the extracted fact "User loves Python now" (search)
|
# embed_batch called once for extracted memories (Phase 3)
|
||||||
# 2. For the rewritten text "The user loves Python" (pre-cached before _update_memory)
|
assert embedder.embed.call_count == 1
|
||||||
# It should NOT be called a 3rd time inside _update_memory
|
assert embedder.embed_batch.call_count == 1
|
||||||
assert embedder.embed.call_count == 2
|
mock_vector_store.insert.assert_called_once()
|
||||||
mock_vector_store.update.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@@ -761,30 +763,6 @@ class TestProcessMetadataFiltersMerge:
|
|||||||
|
|
||||||
# --- Issue #3040: reset() should clean up graph database ---
|
# --- Issue #3040: reset() should clean up graph database ---
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
|
||||||
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
|
||||||
@patch('mem0.utils.factory.LlmFactory.create')
|
|
||||||
@patch('mem0.memory.storage.SQLiteManager')
|
|
||||||
def test_reset_calls_graph_reset_when_graph_enabled(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""Test that reset() calls graph.reset() when graph is enabled (issue #3040)."""
|
|
||||||
mock_embedder_factory.return_value = MagicMock()
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.return_value = mock_vector_store
|
|
||||||
mock_llm_factory.return_value = MagicMock()
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
config = MemoryConfig()
|
|
||||||
memory = Memory(config)
|
|
||||||
|
|
||||||
# Simulate graph being enabled by setting graph instance
|
|
||||||
mock_graph = MagicMock()
|
|
||||||
memory.graph = mock_graph
|
|
||||||
|
|
||||||
memory.reset()
|
|
||||||
|
|
||||||
mock_graph.reset.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
@patch('mem0.utils.factory.EmbedderFactory.create')
|
||||||
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
||||||
@patch('mem0.utils.factory.LlmFactory.create')
|
@patch('mem0.utils.factory.LlmFactory.create')
|
||||||
@@ -807,28 +785,3 @@ def test_reset_skips_graph_when_graph_disabled(mock_sqlite, mock_llm_factory, mo
|
|||||||
|
|
||||||
# graph should remain None after reset
|
# graph should remain None after reset
|
||||||
assert memory.graph is None
|
assert memory.graph is None
|
||||||
|
|
||||||
|
|
||||||
@patch('mem0.utils.factory.EmbedderFactory.create')
|
|
||||||
@patch('mem0.utils.factory.VectorStoreFactory.create')
|
|
||||||
@patch('mem0.utils.factory.LlmFactory.create')
|
|
||||||
@patch('mem0.memory.storage.SQLiteManager')
|
|
||||||
def test_reset_continues_if_graph_reset_fails(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
|
|
||||||
"""Test that reset() doesn't crash if graph.reset() raises an exception (issue #3040)."""
|
|
||||||
mock_embedder_factory.return_value = MagicMock()
|
|
||||||
mock_vector_store = MagicMock()
|
|
||||||
mock_vector_factory.return_value = mock_vector_store
|
|
||||||
mock_llm_factory.return_value = MagicMock()
|
|
||||||
mock_sqlite.return_value = MagicMock()
|
|
||||||
|
|
||||||
config = MemoryConfig()
|
|
||||||
memory = Memory(config)
|
|
||||||
|
|
||||||
mock_graph = MagicMock()
|
|
||||||
mock_graph.reset.side_effect = Exception("Neo4j connection failed")
|
|
||||||
memory.graph = mock_graph
|
|
||||||
|
|
||||||
# Should NOT raise — graph failure is logged but reset continues
|
|
||||||
memory.reset()
|
|
||||||
|
|
||||||
mock_graph.reset.assert_called_once()
|
|
||||||
|
|||||||
@@ -291,11 +291,6 @@ class TestMemoryLifecycle:
|
|||||||
m.db = MagicMock()
|
m.db = MagicMock()
|
||||||
return m
|
return m
|
||||||
|
|
||||||
def test_close_calls_db_close(self):
|
|
||||||
m = self._make_mock_memory()
|
|
||||||
m.close()
|
|
||||||
m.db.close.assert_called_once()
|
|
||||||
|
|
||||||
def test_double_close_is_safe(self):
|
def test_double_close_is_safe(self):
|
||||||
"""close() should be safe to call twice (SQLiteManager.close sets connection=None)."""
|
"""close() should be safe to call twice (SQLiteManager.close sets connection=None)."""
|
||||||
m = self._make_mock_memory()
|
m = self._make_mock_memory()
|
||||||
@@ -319,21 +314,6 @@ class TestMemoryLifecycle:
|
|||||||
# db attribute not set at all
|
# db attribute not set at all
|
||||||
m.close() # should not raise due to hasattr guard
|
m.close() # should not raise due to hasattr guard
|
||||||
|
|
||||||
def test_context_manager(self):
|
|
||||||
"""Memory should support `with` statement and close on exit."""
|
|
||||||
m = self._make_mock_memory()
|
|
||||||
with m:
|
|
||||||
pass
|
|
||||||
m.db.close.assert_called_once()
|
|
||||||
|
|
||||||
def test_context_manager_closes_on_exception(self):
|
|
||||||
"""Memory should close even if the with-block raises."""
|
|
||||||
m = self._make_mock_memory()
|
|
||||||
with pytest.raises(ValueError):
|
|
||||||
with m:
|
|
||||||
raise ValueError("boom")
|
|
||||||
m.db.close.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
class TestAsyncMemoryLifecycle:
|
class TestAsyncMemoryLifecycle:
|
||||||
"""Verify AsyncMemory.close() and async context manager support."""
|
"""Verify AsyncMemory.close() and async context manager support."""
|
||||||
@@ -346,11 +326,6 @@ class TestAsyncMemoryLifecycle:
|
|||||||
m.db = MagicMock()
|
m.db = MagicMock()
|
||||||
return m
|
return m
|
||||||
|
|
||||||
def test_close_calls_db_close(self):
|
|
||||||
m = self._make_mock_async_memory()
|
|
||||||
m.close()
|
|
||||||
m.db.close.assert_called_once()
|
|
||||||
|
|
||||||
def test_close_when_db_is_none(self):
|
def test_close_when_db_is_none(self):
|
||||||
m = self._make_mock_async_memory()
|
m = self._make_mock_async_memory()
|
||||||
m.db = None
|
m.db = None
|
||||||
@@ -363,29 +338,6 @@ class TestAsyncMemoryLifecycle:
|
|||||||
m = AsyncMemory.__new__(AsyncMemory)
|
m = AsyncMemory.__new__(AsyncMemory)
|
||||||
m.close() # should not raise
|
m.close() # should not raise
|
||||||
|
|
||||||
def test_async_context_manager(self):
|
|
||||||
"""AsyncMemory should support `async with` and close on exit."""
|
|
||||||
m = self._make_mock_async_memory()
|
|
||||||
|
|
||||||
async def run():
|
|
||||||
async with m:
|
|
||||||
pass
|
|
||||||
|
|
||||||
asyncio.run(run())
|
|
||||||
m.db.close.assert_called_once()
|
|
||||||
|
|
||||||
def test_async_context_manager_closes_on_exception(self):
|
|
||||||
"""AsyncMemory should close even if the async with-block raises."""
|
|
||||||
m = self._make_mock_async_memory()
|
|
||||||
|
|
||||||
async def run():
|
|
||||||
async with m:
|
|
||||||
raise ValueError("boom")
|
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
|
||||||
asyncio.run(run())
|
|
||||||
m.db.close.assert_called_once()
|
|
||||||
|
|
||||||
|
|
||||||
class TestTelemetryEnvVar:
|
class TestTelemetryEnvVar:
|
||||||
"""Verify the MEM0_TELEMETRY env var parsing logic."""
|
"""Verify the MEM0_TELEMETRY env var parsing logic."""
|
||||||
|
|||||||
@@ -42,9 +42,12 @@ def test_initalize_create_col(mongo_vector_fixture):
|
|||||||
mock_collection.insert_one.assert_called_once_with({"_id": 0, "placeholder": True})
|
mock_collection.insert_one.assert_called_once_with({"_id": 0, "placeholder": True})
|
||||||
mock_collection.delete_one.assert_called_once_with({"_id": 0})
|
mock_collection.delete_one.assert_called_once_with({"_id": 0})
|
||||||
assert mongo_vector.index_name == "test_collection_vector_index"
|
assert mongo_vector.index_name == "test_collection_vector_index"
|
||||||
mock_collection.list_search_indexes.assert_called_once_with(name="test_collection_vector_index")
|
mock_collection.list_search_indexes.assert_any_call(name="test_collection_vector_index")
|
||||||
mock_collection.create_search_index.assert_called_once()
|
mock_collection.list_search_indexes.assert_any_call(name="test_collection_text_search_index")
|
||||||
args, _ = mock_collection.create_search_index.call_args
|
assert mock_collection.list_search_indexes.call_count == 2
|
||||||
|
# Two indexes created: vector search + text search
|
||||||
|
assert mock_collection.create_search_index.call_count == 2
|
||||||
|
args, _ = mock_collection.create_search_index.call_args_list[0]
|
||||||
search_index_model = args[0].document
|
search_index_model = args[0].document
|
||||||
assert search_index_model == {
|
assert search_index_model == {
|
||||||
"name": "test_collection_vector_index",
|
"name": "test_collection_vector_index",
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ import unittest
|
|||||||
import uuid
|
import uuid
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
from qdrant_client import QdrantClient
|
from qdrant_client import QdrantClient, models
|
||||||
from qdrant_client.models import (
|
from qdrant_client.models import (
|
||||||
DatetimeRange,
|
DatetimeRange,
|
||||||
Distance,
|
Distance,
|
||||||
@@ -18,6 +18,7 @@ from qdrant_client.models import (
|
|||||||
PointStruct,
|
PointStruct,
|
||||||
PointVectors,
|
PointVectors,
|
||||||
Range,
|
Range,
|
||||||
|
SparseVectorParams,
|
||||||
VectorParams,
|
VectorParams,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -59,8 +60,14 @@ class TestQdrant(unittest.TestCase):
|
|||||||
|
|
||||||
expected_config = VectorParams(size=128, distance=Distance.COSINE, on_disk=True)
|
expected_config = VectorParams(size=128, distance=Distance.COSINE, on_disk=True)
|
||||||
|
|
||||||
|
expected_sparse_config = {
|
||||||
|
"bm25": SparseVectorParams(modifier=models.Modifier.IDF),
|
||||||
|
}
|
||||||
|
|
||||||
self.client_mock.create_collection.assert_called_with(
|
self.client_mock.create_collection.assert_called_with(
|
||||||
collection_name="test_collection", vectors_config=expected_config
|
collection_name="test_collection",
|
||||||
|
vectors_config=expected_config,
|
||||||
|
sparse_vectors_config=expected_sparse_config,
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_insert(self):
|
def test_insert(self):
|
||||||
@@ -231,7 +238,9 @@ class TestQdrant(unittest.TestCase):
|
|||||||
self.client_mock.upsert.assert_called_once()
|
self.client_mock.upsert.assert_called_once()
|
||||||
point = self.client_mock.upsert.call_args[1]["points"][0]
|
point = self.client_mock.upsert.call_args[1]["points"][0]
|
||||||
self.assertEqual(point.id, vector_id)
|
self.assertEqual(point.id, vector_id)
|
||||||
self.assertEqual(point.vector, updated_vector)
|
# v3 uses named vectors: dense vector stored under "" key
|
||||||
|
self.assertIn("", point.vector)
|
||||||
|
self.assertEqual(point.vector[""], updated_vector)
|
||||||
self.assertEqual(point.payload, updated_payload)
|
self.assertEqual(point.payload, updated_payload)
|
||||||
|
|
||||||
def test_update_with_none_vector_uses_set_payload(self):
|
def test_update_with_none_vector_uses_set_payload(self):
|
||||||
|
|||||||
Reference in New Issue
Block a user