diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 4860d0ad..06daa870 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -185,10 +185,10 @@ class BM25Index(BaseKeywordIndex): self._idf_cache[token_id] = idf return idf - def _prepare_doc(self, doc_id: str, content: str) -> tuple[np.ndarray, int, Counter] | None: - """Tokenize and count terms; retire any prior version of `doc_id`.""" + def _prepare_doc(self, doc_id: str, tokens: list[str]) -> tuple[np.ndarray, int, Counter] | None: + """Count prepared tokens; retire any prior version of `doc_id`.""" self._remove_doc(doc_id) - token_ids = self._tokens_to_ids(self._tokenize(content)) + token_ids = self._tokens_to_ids(tokens) if not token_ids: return None counts = Counter(token_ids) @@ -252,14 +252,17 @@ class BM25Index(BaseKeywordIndex): if not docs_dict: return + # Tokenization can fail; finish it before retiring documents or allocating slots. + tokenized_docs = {doc_id: self._tokenize(content) for doc_id, content in docs_dict.items()} + new_doc_ids: list[str] = [] new_doc_lens: list[int] = [] new_doc_token_ids: list[np.ndarray] = [] pending: dict[int, list[tuple[int, int]]] = {} next_idx = len(self._doc_ids) - for doc_id, content in docs_dict.items(): - prepared = self._prepare_doc(doc_id, content) + for doc_id, tokens in tokenized_docs.items(): + prepared = self._prepare_doc(doc_id, tokens) if prepared is None: continue unique_tids, n_tokens, token_counts = prepared diff --git a/tests/unit/test_bm25_live_length.py b/tests/unit/test_bm25_live_length.py index 5ba674d3..219beff1 100644 --- a/tests/unit/test_bm25_live_length.py +++ b/tests/unit/test_bm25_live_length.py @@ -87,23 +87,42 @@ async def test_snapshot_roundtrip_rebuilds_length_without_new_fields(tmp_path, m @pytest.mark.asyncio -async def test_failed_batch_counts_only_published_arrays(tmp_path, monkeypatch): - monkeypatch.chdir(tmp_path) +@pytest.mark.parametrize( + "batch", + [ + {"existing": "unavailable"}, + {"new": "fresh memory", "failing": "unavailable"}, + {"new": "fresh memory", "existing": "unavailable"}, + {"existing": "replacement memory", "failing": "unavailable"}, + ], +) +async def test_failed_tokenization_preserves_documents_and_allows_retry(batch): + """A tokenizer error must not retire old documents or leave incomplete new slots.""" index = make_index() - await index.add_docs({"old": "alpha beta", "keep": "gamma"}) - tokenize = index._tokenize + tokenizer = index.tokenizer + await index.add_docs({"existing": "original memory", "untouched": "stable memory"}) + before = await index.retrieve("original stable fresh replacement", limit=10) + tokenize = tokenizer.tokenize - def fail(text): - if text == "fail": - raise ValueError("tokenizer failed") - return tokenize(text) + def fail_on_unavailable(texts, **kwargs): + if "unavailable" in texts: + raise RuntimeError("tokenizer unavailable") + return tokenize(texts, **kwargs) + + with patch.object(tokenizer, "tokenize", side_effect=fail_on_unavailable): + with pytest.raises(RuntimeError, match="tokenizer unavailable"): + await index.add_docs(batch) - with patch.object(index, "_tokenize", side_effect=fail): - with pytest.raises(ValueError, match="tokenizer failed"): - await index.add_docs({"pending": "alpha alpha", "old": "fail"}) - # Existing partial-batch behavior retires old but does not append pending. assert_length(index) - assert index.total_len == 1 + assert index.total_len == 4 + assert set(index.document_ids) == {"existing", "untouched"} + assert await index.retrieve("original stable fresh replacement", limit=10) == before + + retry = {doc_id: "recovered memory" if text == "unavailable" else text for doc_id, text in batch.items()} + await index.add_docs(retry) + assert set(index.document_ids) == {"existing", "untouched"} | set(batch) + assert set(await index.retrieve("fresh replacement recovered", limit=10)) == set(batch) + assert set(await index.retrieve("stable")) == {"untouched"} @pytest.mark.asyncio