diff --git a/CHANGELOG.md b/CHANGELOG.md index 077e103..746d1f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - `TextFileLoader` and `MarkdownLoader` now convert Unicode decoding and unknown encoding errors to `LoaderError`, with configurable text decoding error handling (closes #24). - Validate chunker size and overlap parameters in `FixedSizeChunker` and `SentenceChunker` (closes #25). +- Ensure retrievers consistently apply last-occurrence-wins semantics for duplicate chunk IDs and return the latest content after upserts. ## [0.3.0] - 2026-09-25 diff --git a/ragframework/base.py b/ragframework/base.py index 850ce7e..3aa8d81 100644 --- a/ragframework/base.py +++ b/ragframework/base.py @@ -119,13 +119,18 @@ class Retriever(ABC): @abstractmethod def add(self, chunks: list[Chunk]) -> None: - """Add *chunks* to the index. + """Add or replace *chunks* in the index. + + Adding a chunk whose :attr:`~Chunk.id` already exists replaces the + previously indexed chunk with the same ID. Retrievers therefore expose + upsert semantics: repeated ingestion of the same chunk ID does not + increase the number of indexed chunks. Chunks must have their :attr:`~Chunk.embedding` set before calling this method. Args: - chunks: Chunks to index. + chunks: Chunks to add or replace in the index. """ @abstractmethod diff --git a/ragframework/retriever/chroma.py b/ragframework/retriever/chroma.py index 657eba4..be06f11 100644 --- a/ragframework/retriever/chroma.py +++ b/ragframework/retriever/chroma.py @@ -54,11 +54,14 @@ def __init__( raise RetrieverError(f"Could not initialize ChromaDB: {exc}") from exc def add(self, chunks: list[Chunk]) -> None: - """Add embedded chunks to the ChromaDB collection.""" + """Upsert chunks; the last occurrence of each ID wins.""" if not chunks: return - embeddings: list[Sequence[float]] = [] + # Validate all chunks and deduplicate IDs. + # Reassigning an existing key keeps its original position, + # while the latest chunk becomes its value. + unique_chunks: dict[str, Chunk] = {} for chunk in chunks: if chunk.embedding is None: @@ -66,23 +69,27 @@ def add(self, chunks: list[Chunk]) -> None: f"Chunk '{chunk.id}' has no embedding. " "Embed chunks before adding them to the retriever." ) + unique_chunks[chunk.id] = chunk - embeddings.append(chunk.embedding) + final_chunks = list(unique_chunks.values()) + final_embeddings: list[Sequence[float]] = [ + chunk.embedding for chunk in final_chunks if chunk.embedding is not None + ] try: self._collection.upsert( - ids=[chunk.id for chunk in chunks], - embeddings=embeddings, - documents=[chunk.content for chunk in chunks], - metadatas=[self._metadata(chunk) for chunk in chunks], + ids=[chunk.id for chunk in final_chunks], + embeddings=final_embeddings, + documents=[chunk.content for chunk in final_chunks], + metadatas=[self._metadata(chunk) for chunk in final_chunks], ) except Exception as exc: raise RetrieverError(f"Failed to add chunks to ChromaDB: {exc}") from exc logger.debug( "Added chunks count=%d dimension=%d", - len(chunks), - len(embeddings[0]), + len(final_chunks), + len(final_embeddings[0]), ) def retrieve( diff --git a/ragframework/retriever/faiss.py b/ragframework/retriever/faiss.py index 93f7de5..5dc2555 100644 --- a/ragframework/retriever/faiss.py +++ b/ragframework/retriever/faiss.py @@ -19,8 +19,10 @@ from ragframework.utils.vectors import validate_vector logger = logging.getLogger(__name__) + _INSTALL_HINT = ( - "FAISS support requires 'ragframework[faiss]'. Install it with: pip install ragframework[faiss]" + "FAISS support requires 'ragframework[faiss]'. " + "Install it with: pip install ragframework[faiss]" ) @@ -28,7 +30,7 @@ class FAISSRetriever(Retriever): """Retrieve chunks with an HNSW index and cosine similarity. Chunk and query vectors are L2-normalized before being passed to an - ``IndexHNSWFlat`` configured for inner product. Inner product between + ``IndexHNSWFlat`` configured for inner product. Inner product between normalized vectors is cosine similarity. Args: @@ -41,9 +43,15 @@ class FAISSRetriever(Retriever): RetrieverError: If an HNSW parameter is not positive. """ - def __init__(self, m: int = 32, ef_construction: int = 40, ef_search: int = 16) -> None: + def __init__( + self, + m: int = 32, + ef_construction: int = 40, + ef_search: int = 16, + ) -> None: if not _FAISS_AVAILABLE: raise ImportError(_INSTALL_HINT) + if any( not isinstance(value, int) or isinstance(value, bool) or value <= 0 for value in (m, ef_construction, ef_search) @@ -53,87 +61,205 @@ def __init__(self, m: int = 32, ef_construction: int = 40, ef_search: int = 16) self._m = m self._ef_construction = ef_construction self._ef_search = ef_search + self._chunks: list[Chunk] = [] self._dimension: int | None = None self._index: Any | None = None + # Maps logical chunk IDs to their current physical FAISS labels. + # Old labels remain in the FAISS index as tombstones after replacement. + self._id_to_index: dict[str, int] = {} + + # Logical chunk ID -> position in self._chunks. + self._id_to_position: dict[str, int] = {} + + def _build_index( + self, + vectors: np.ndarray[Any, Any], + dimension: int, + ) -> Any: + """Build an HNSW index from normalized vectors.""" + index = faiss.IndexHNSWFlat( + dimension, + self._m, + faiss.METRIC_INNER_PRODUCT, + ) + + index.hnsw.efConstruction = self._ef_construction + index.hnsw.efSearch = self._ef_search + + index.add(vectors) + + return index + def add(self, chunks: list[Chunk]) -> None: - """Validate, normalize, and add embedded chunks to the HNSW index. + """Validate and upsert chunks; the last duplicate ID in a batch wins. - An empty batch is a no-op. Every vector in a batch is validated before - either the FAISS index or the chunk mapping is changed. + Existing IDs retain their logical position. New IDs are appended. + Replaced physical FAISS vectors become tombstones. """ if not chunks: return expected_dimension = self._dimension - vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + validated: dict[str, tuple[Chunk, np.ndarray[Any, np.dtype[np.float32]]]] = {} + + # Validate the entire batch before mutating retriever state. + # Dict reassignment keeps the first-seen order but the last value. for chunk in chunks: if chunk.embedding is None: raise RetrieverError( f"Chunk '{chunk.id}' has no embedding. " "Embed chunks before adding them to the retriever." ) - vector = validate_vector(chunk.embedding, f"Chunk '{chunk.id}' embedding") + + vector = validate_vector( + chunk.embedding, + f"Chunk '{chunk.id}' embedding", + ) + if expected_dimension is None: expected_dimension = int(vector.shape[0]) elif vector.shape[0] != expected_dimension: raise RetrieverError( - f"Chunk '{chunk.id}' embedding has dimension {vector.shape[0]}; " - f"expected {expected_dimension}." + f"Chunk '{chunk.id}' embedding has dimension " + f"{vector.shape[0]}; expected {expected_dimension}." ) - vectors.append(vector) + + validated[chunk.id] = (chunk, vector) assert expected_dimension is not None - matrix = np.ascontiguousarray(np.stack(vectors), dtype=np.float32) - new_index = self._index - try: - if new_index is None: - new_index = faiss.IndexHNSWFlat( - expected_dimension, self._m, faiss.METRIC_INNER_PRODUCT - ) - new_index.hnsw.efConstruction = self._ef_construction - new_index.hnsw.efSearch = self._ef_search - new_index.add(matrix) - except Exception as exc: - raise RetrieverError("FAISS failed to add chunk embeddings.") from exc - self._index = new_index - self._dimension = expected_dimension - self._chunks.extend(chunks) + unique_chunks = [item[0] for item in validated.values()] + vectors = [item[1] for item in validated.values()] + matrix = np.ascontiguousarray( + np.stack(vectors), + dtype=np.float32, + ) + + # Build a fresh index for the first insertion. + if self._index is None: + try: + new_index = self._build_index(matrix, expected_dimension) + except Exception as exc: + raise RetrieverError("FAISS failed to add chunk embeddings.") from exc + + self._index = new_index + self._dimension = expected_dimension + + for physical_label, chunk in enumerate(unique_chunks): + logical_position = len(self._chunks) + self._chunks.append(chunk) + self._id_to_index[chunk.id] = physical_label + self._id_to_position[chunk.id] = logical_position + + else: + # Add only the final version of each ID to FAISS. + start_label = int(self._index.ntotal) + try: + self._index.add(matrix) + except Exception as exc: + raise RetrieverError("FAISS failed to add chunk embeddings.") from exc + + # Update the logical mapping only after FAISS succeeds. + for offset, chunk in enumerate(unique_chunks): + physical_label = start_label + offset + existing_position = self._id_to_position.get(chunk.id) + + if existing_position is None: + logical_position = len(self._chunks) + self._chunks.append(chunk) + self._id_to_position[chunk.id] = logical_position + else: + logical_position = existing_position + self._chunks[logical_position] = chunk + + self._id_to_index[chunk.id] = physical_label + + self._dimension = expected_dimension + logger.debug( "Added chunks count=%d dimension=%d", - len(chunks), + len(unique_chunks), matrix.shape[1], ) - def retrieve(self, query_embedding: list[float], top_k: int = 5) -> list[Chunk]: - """Return up to ``top_k`` chunks ordered by approximate cosine similarity. - - Non-positive ``top_k`` values return an empty list. An empty index also - returns an empty list. Invalid FAISS labels (for example ``-1`` when - fewer than ``top_k`` neighbours exist) are ignored. - """ + def retrieve( + self, + query_embedding: list[float], + top_k: int = 5, + ) -> list[Chunk]: + """Retrieve active chunks using progressively expanded FAISS searches.""" if not isinstance(top_k, int) or isinstance(top_k, bool): raise RetrieverError("top_k must be an integer.") - if top_k <= 0 or self._index is None or self._dimension is None: + + if top_k <= 0 or self._index is None or self._dimension is None or not self._chunks: return [] query = validate_vector(query_embedding, "Query embedding") + if query.shape[0] != self._dimension: raise RetrieverError( - f"Query embedding has dimension {query.shape[0]}; expected {self._dimension}." + f"Query embedding has dimension {query.shape[0]}; " f"expected {self._dimension}." ) - count = min(top_k, len(self._chunks)) + total = int(self._index.ntotal) + + # Current physical label -> current logical position. + active_labels = { + physical_label: self._id_to_position[chunk_id] + for chunk_id, physical_label in self._id_to_index.items() + if chunk_id in self._id_to_position + } + + # Start with a bounded over-fetch, not every physical vector. + search_count = min(total, max(top_k * 2, top_k + 8)) + candidates: dict[int, tuple[float, int, Chunk]] = {} + try: - _scores, labels = self._index.search(query.reshape(1, -1), count) + while True: + scores, labels = self._index.search( + query.reshape(1, -1), + search_count, + ) + + for score, label in zip( + scores[0], + labels[0], + strict=True, + ): + physical_label = int(label) + if physical_label < 0: + continue + + logical_position = active_labels.get(physical_label) + if logical_position is None: + continue + + chunk = self._chunks[logical_position] + candidates[physical_label] = ( + float(score), + logical_position, + chunk, + ) + + if len(candidates) >= min(top_k, len(self._chunks)): + break + + if search_count >= total: + break + + search_count = min(total, search_count * 2) + except Exception as exc: raise RetrieverError("FAISS failed to search the index.") from exc - return [ - self._chunks[int(label)] for label in labels[0] if 0 <= int(label) < len(self._chunks) - ] + ordered = sorted( + candidates.values(), + key=lambda item: (-item[0], item[1]), + ) + + return [chunk for _score, _position, chunk in ordered[:top_k]] def __len__(self) -> int: """Return the number of indexed chunks.""" diff --git a/ragframework/retriever/in_memory.py b/ragframework/retriever/in_memory.py index 104b38e..ad5fbdd 100644 --- a/ragframework/retriever/in_memory.py +++ b/ragframework/retriever/in_memory.py @@ -15,80 +15,139 @@ class InMemoryRetriever(Retriever): - """A simple retriever that stores chunks in RAM and ranks by cosine similarity. - - Suitable for small corpora and quick experiments. For large-scale use, - swap this out with a FAISS or ChromaDB retriever (see - ``.github/GOOD_FIRST_ISSUES.md``). - """ + """Store chunks in RAM and rank them by cosine similarity.""" def __init__(self) -> None: self._chunks: list[Chunk] = [] - self._matrix: np.ndarray[Any, Any] | None = None # shape (N, dim) + self._matrix: np.ndarray[Any, Any] | None = None + self._matrix_blocks: list[np.ndarray[Any, Any]] = [] + self._block_ids: list[list[str]] = [] + self._id_to_index: dict[str, int] = {} + self._id_to_location: dict[str, tuple[int, int]] = {} self._dimension: int | None = None def add(self, chunks: list[Chunk]) -> None: - """Validate and normalize a batch before changing the stored chunks. - - An empty batch is a no-op. Invalid vectors or inconsistent dimensions - raise ``RetrieverError`` without adding any part of the batch. - """ + """Validate and upsert chunks; the last duplicate ID wins.""" if not chunks: return expected_dimension = self._dimension - vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + validated: dict[str, tuple[Chunk, np.ndarray[Any, np.dtype[np.float32]]]] = {} + + # Validate the entire batch before changing stored state. for chunk in chunks: if chunk.embedding is None: raise RetrieverError( f"Chunk '{chunk.id}' has no embedding. " "Embed chunks before adding them to the retriever." ) + vector = validate_vector(chunk.embedding, f"Chunk '{chunk.id}' embedding") + if expected_dimension is None: expected_dimension = int(vector.shape[0]) elif vector.shape[0] != expected_dimension: raise RetrieverError( - f"Chunk '{chunk.id}' embedding has dimension {vector.shape[0]}; " - f"expected {expected_dimension}." + f"Chunk '{chunk.id}' embedding has dimension " + f"{vector.shape[0]}; expected {expected_dimension}." ) - vectors.append(vector) - matrix = np.stack(vectors) - if self._matrix is not None: - matrix = np.concatenate((self._matrix, matrix)) + # Dict assignment makes the last occurrence win. + validated[chunk.id] = (chunk, vector) - self._matrix = matrix + assert expected_dimension is not None + + # Prepare the new state separately so failed preparation cannot + # partially mutate the existing retriever. + final_chunks = list(self._chunks) + final_positions = dict(self._id_to_index) + final_locations = dict(self._id_to_location) + final_blocks = list(self._matrix_blocks) + final_block_ids = [list(ids) for ids in self._block_ids] + + new_chunks = list(validated.values()) + new_ids = list(validated.keys()) + new_matrix = np.ascontiguousarray( + np.stack([vector for _, vector in new_chunks]), + dtype=np.float32, + ) + + block_index = len(final_blocks) + final_blocks.append(new_matrix) + final_block_ids.append(new_ids) + + for row_index, (chunk, _) in enumerate(new_chunks): + position = final_positions.get(chunk.id) + + if position is None: + position = len(final_chunks) + final_positions[chunk.id] = position + final_chunks.append(chunk) + else: + final_chunks[position] = chunk + + # The latest location supersedes any older physical row. + final_locations[chunk.id] = (block_index, row_index) + + self._chunks = final_chunks + self._id_to_index = final_positions + self._id_to_location = final_locations + self._matrix_blocks = final_blocks + self._block_ids = final_block_ids + self._matrix = None self._dimension = expected_dimension - self._chunks.extend(chunks) + logger.debug( "Added chunks count=%d dimension=%d", - len(chunks), - matrix.shape[1], + len(validated), + expected_dimension, ) - def retrieve(self, query_embedding: list[float], top_k: int = 5) -> list[Chunk]: - """Rank chunks by cosine similarity. + def _consolidate(self) -> None: + """Build one matrix from the latest vector for each logical chunk.""" + if self._matrix is not None: + return - Non-positive ``top_k`` values return an empty list. - An empty index also returns an empty list. Invalid query vectors - or dimensions raise ``RetrieverError``. Non-integer / bool - ``top_k`` raises ``RetrieverError``. - """ + if not self._chunks: + return + + vectors: list[np.ndarray[Any, Any]] = [] + + for chunk in self._chunks: + block_index, row_index = self._id_to_location[chunk.id] + vectors.append(self._matrix_blocks[block_index][row_index]) + + matrix = np.ascontiguousarray(np.stack(vectors), dtype=np.float32) + + self._matrix = matrix + self._matrix_blocks = [matrix] + self._block_ids = [[chunk.id for chunk in self._chunks]] + self._id_to_location = {chunk.id: (0, index) for index, chunk in enumerate(self._chunks)} + + def retrieve(self, query_embedding: list[float], top_k: int = 5) -> list[Chunk]: + """Return the most similar chunks first.""" if not isinstance(top_k, int) or isinstance(top_k, bool): raise RetrieverError("top_k must be an integer.") - if top_k <= 0 or not self._chunks or self._matrix is None: + + if top_k <= 0 or not self._chunks: return [] + q = validate_vector(query_embedding, "Query embedding") + if q.shape[0] != self._dimension: raise RetrieverError( - f"Query embedding has dimension {q.shape[0]}; expected {self._dimension}." + f"Query embedding has dimension {q.shape[0]}; " f"expected {self._dimension}." ) + + self._consolidate() + assert self._matrix is not None + scores: np.ndarray[Any, Any] = self._matrix @ q k = min(top_k, len(self._chunks)) top_indices = np.argpartition(scores, -k)[-k:] top_indices = top_indices[np.argsort(scores[top_indices])[::-1]] - return [self._chunks[i] for i in top_indices] + + return [self._chunks[int(index)] for index in top_indices] def __len__(self) -> int: return len(self._chunks) diff --git a/tests/test_retriever/test_chroma.py b/tests/test_retriever/test_chroma.py index a46bbca..77e07a9 100644 --- a/tests/test_retriever/test_chroma.py +++ b/tests/test_retriever/test_chroma.py @@ -321,3 +321,71 @@ def test_user_metadata_with_old_empty_marker_is_preserved(): assert len(results) == 1 assert results[0].metadata == metadata + + +def test_repeated_upsert_returns_latest_content_and_metadata(): + retriever = ChromaRetriever( + collection_name="test_repeated_upsert_latest", + ) + + old = Chunk( + id="same-id", + content="Old content", + metadata={"version": "old"}, + embedding=[1.0, 0.0], + ) + + latest = Chunk( + id="same-id", + content="Latest content", + metadata={"version": "latest"}, + embedding=[0.0, 1.0], + ) + + retriever.add([old]) + retriever.add([latest]) + + results = retriever.retrieve([0.0, 1.0], top_k=1) + + assert len(retriever) == 1 + assert len(results) == 1 + assert results[0].id == "same-id" + assert results[0].content == "Latest content" + assert results[0].metadata == {"version": "latest"} + + +def test_duplicate_ids_in_batch_last_occurrence_wins(): + retriever = ChromaRetriever( + collection_name="test_duplicate_ids_last_wins", + ) + + old = Chunk( + id="duplicate", + content="Old content", + metadata={"version": "old"}, + embedding=[1.0, 0.0], + ) + + other = Chunk( + id="other", + content="Other content", + metadata={}, + embedding=[1.0, 0.0], + ) + + latest = Chunk( + id="duplicate", + content="Latest content", + metadata={"version": "latest"}, + embedding=[0.0, 1.0], + ) + + retriever.add([old, other, latest]) + + results = retriever.retrieve([0.0, 1.0], top_k=2) + by_id = {chunk.id: chunk for chunk in results} + + assert len(retriever) == 2 + assert len(by_id) == 2 + assert by_id["duplicate"].content == "Latest content" + assert by_id["duplicate"].metadata == {"version": "latest"} diff --git a/tests/test_retriever/test_faiss.py b/tests/test_retriever/test_faiss.py index c180502..69c3bf9 100644 --- a/tests/test_retriever/test_faiss.py +++ b/tests/test_retriever/test_faiss.py @@ -41,15 +41,85 @@ def test_real_faiss_search_uses_cosine_similarity_and_preserves_chunks() -> None assert results[0].metadata == {"source": "north"} -def test_batches_append_without_replacing_existing_chunks() -> None: +def test_batches_upsert_without_duplicate_ids() -> None: retriever = FAISSRetriever(ef_search=64) + first = make_chunk("first", [1.0, 0.0, 0.0]) second = make_chunk("second", [0.0, 1.0, 0.0]) + replacement = make_chunk("first", [0.0, 1.0, 0.0]) + retriever.add([first]) retriever.add([second]) + retriever.add([replacement]) assert len(retriever) == 2 - assert retriever.retrieve([0.0, 1.0, 0.0], top_k=2)[0] is second + assert retriever._chunks[0] is replacement + assert retriever._chunks[1] is second + + results = retriever.retrieve([0.0, 1.0, 0.0], top_k=2) + assert len(results) == 2 + assert {chunk.id for chunk in results} == {"first", "second"} + assert retriever._chunks[0] is replacement + + +def test_add_batch_replaces_existing_and_adds_new_chunks() -> None: + retriever = FAISSRetriever(ef_search=64) + + first = make_chunk("first", [1.0, 0.0, 0.0]) + second = make_chunk("second", [0.0, 1.0, 0.0]) + replacement = make_chunk("first", [0.0, 1.0, 0.0]) + third = make_chunk("third", [1.0, 1.0, 0.0]) + + retriever.add([first, second]) + retriever.add([replacement, third]) + + results = retriever.retrieve([0.0, 1.0, 0.0], top_k=3) + + assert len(results) == 3 + assert {chunk.id for chunk in results} == {"first", "second", "third"} + assert retriever._chunks[0] is replacement + assert retriever._chunks[1] is second + assert retriever._chunks[2] is third + + +def test_add_same_id_replaces_existing_chunk() -> None: + retriever = FAISSRetriever(ef_search=64) + + old = make_chunk("same", [1.0, 0.0]) + new = make_chunk("same", [0.0, 1.0]) + + retriever.add([old]) + retriever.add([new]) + + assert len(retriever) == 1 + assert retriever.retrieve([0.0, 1.0], top_k=1) == [new] + assert retriever.retrieve([0.0, 1.0], top_k=1)[0].content == new.content + assert retriever.retrieve([0.0, 1.0], top_k=1)[0].embedding == new.embedding + + +def test_failed_upsert_preserves_existing_index( + monkeypatch: pytest.MonkeyPatch, +) -> None: + retriever = FAISSRetriever() + existing = make_chunk("existing", [1.0, 0.0]) + retriever.add([existing]) + + def fail(_vectors: np.ndarray) -> None: + raise RuntimeError("native failure") + + monkeypatch.setattr(retriever._index, "add", fail) + + replacement = make_chunk("existing", [0.0, 1.0]) + + with pytest.raises(RetrieverError, match="failed to add") as exc_info: + retriever.add([replacement]) + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert len(retriever) == 1 + assert retriever._chunks[0] is existing + + results = retriever.retrieve([1.0, 0.0], top_k=1) + assert results == [existing] def test_add_logs_count_and_dimension(caplog) -> None: @@ -197,18 +267,30 @@ def test_faiss_add_failure_is_wrapped_without_extending_chunk_mapping( monkeypatch: pytest.MonkeyPatch, ) -> None: retriever = FAISSRetriever() - retriever.add([make_chunk("existing", [1.0, 0.0])]) + existing = make_chunk("existing", [1.0, 0.0]) + retriever.add([existing]) - def fail(_vectors: np.ndarray) -> None: + original_add = faiss.IndexHNSWFlat.add + + def fail(self: object, _vectors: np.ndarray) -> None: raise RuntimeError("native failure") - monkeypatch.setattr(retriever._index, "add", fail) + monkeypatch.setattr(faiss.IndexHNSWFlat, "add", fail) + with pytest.raises(RetrieverError, match="failed to add") as exc_info: retriever.add([make_chunk("new", [0.0, 1.0])]) + assert len(retriever) == 1 - assert retriever._chunks[0].id == "existing" + assert retriever._chunks[0] is existing assert isinstance(exc_info.value.__cause__, RuntimeError) + # The original index is still usable. + assert retriever.retrieve([1.0, 0.0], top_k=1) == [existing] + + # Keep the local binding intentionally referenced so linters don't + # complain if the implementation changes. + assert original_add is not None + def test_missing_faiss_fails_only_when_backend_is_constructed() -> None: script = """ @@ -230,3 +312,106 @@ def blocked(name, *args, **kwargs): """ result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) assert result.returncode == 0, result.stderr + + +def test_duplicate_ids_in_initial_batch_last_occurrence_wins() -> None: + retriever = FAISSRetriever(ef_search=64) + old = make_chunk("same", [1.0, 0.0]) + latest = make_chunk("same", [0.0, 1.0]) + other = make_chunk("other", [1.0, 0.0]) + + retriever.add([old, other, latest]) + + assert len(retriever) == 2 + assert [chunk.id for chunk in retriever._chunks] == ["same", "other"] + assert retriever._chunks[0] is latest + assert retriever.retrieve([0.0, 1.0], top_k=1) == [latest] + + +def test_repeated_replacement_returns_latest_content_and_embedding() -> None: + retriever = FAISSRetriever(ef_search=64) + old = make_chunk("same", [1.0, 0.0]) + middle = Chunk( + id="same", + content="Middle version", + metadata={}, + embedding=[0.0, 1.0], + ) + latest = Chunk( + id="same", + content="Latest version", + metadata={}, + embedding=[-1.0, 0.0], + ) + + retriever.add([old]) + retriever.add([middle]) + retriever.add([latest]) + + results = retriever.retrieve([-1.0, 0.0], top_k=1) + + assert len(retriever) == 1 + assert results == [latest] + assert results[0].content == "Latest version" + assert results[0].embedding == [-1.0, 0.0] + + +def test_search_starts_with_bounded_candidate_count( + monkeypatch: pytest.MonkeyPatch, +) -> None: + retriever = FAISSRetriever() + retriever.add([make_chunk(f"chunk-{i}", [1.0, float(i + 1)]) for i in range(30)]) + + original_search = retriever._index.search + requested_counts: list[int] = [] + + def track_search(query: np.ndarray, count: int): + requested_counts.append(count) + return original_search(query, count) + + monkeypatch.setattr(retriever._index, "search", track_search) + + retriever.retrieve([1.0, 0.1], top_k=1) + + assert requested_counts + assert requested_counts[0] < retriever._index.ntotal + + +def test_search_expands_when_initial_candidates_are_stale( + monkeypatch: pytest.MonkeyPatch, +) -> None: + retriever = FAISSRetriever() + old = make_chunk("same", [1.0, 0.0]) + latest = make_chunk("same", [0.0, 1.0]) + + retriever.add([old]) + retriever.add([latest]) + retriever.add([make_chunk(f"other-{i}", [1.0, float(i + 1)]) for i in range(10)]) + + original_search = retriever._index.search + total = int(retriever._index.ntotal) + requested_counts: list[int] = [] + + def simulate_stale_first_search(query: np.ndarray, count: int): + requested_counts.append(count) + + if count < total: + labels = np.full((1, count), -1, dtype=np.int64) + labels[0, 0] = 0 # Obsolete physical label for "same". + scores = np.zeros((1, count), dtype=np.float32) + return scores, labels + + return original_search(query, count) + + monkeypatch.setattr( + retriever._index, + "search", + simulate_stale_first_search, + ) + + results = retriever.retrieve([0.0, 1.0], top_k=1) + + assert len(requested_counts) >= 2 + assert requested_counts[0] < total + assert requested_counts[-1] == total + assert results == [latest] diff --git a/tests/test_retriever/test_in_memory.py b/tests/test_retriever/test_in_memory.py index 88fb664..b1ef56f 100644 --- a/tests/test_retriever/test_in_memory.py +++ b/tests/test_retriever/test_in_memory.py @@ -172,3 +172,73 @@ def test_bool_top_k_raises_retriever_error(self): with pytest.raises(RetrieverError, match="top_k must be an integer"): r.retrieve([1.0, 0.0], top_k=True) + + def test_add_same_id_replaces_existing_chunk(self): + r = InMemoryRetriever() + + old = make_chunk("same", [1.0, 0.0]) + new = make_chunk("same", [0.0, 1.0]) + + r.add([old]) + r.add([new]) + + assert len(r) == 1 + assert r.retrieve([0.0, 1.0], top_k=1) == [new] + assert r.retrieve([0.0, 1.0], top_k=1)[0].content == new.content + assert r.retrieve([0.0, 1.0], top_k=1)[0].embedding == new.embedding + + def test_add_batch_replaces_existing_and_adds_new_chunks(self): + r = InMemoryRetriever() + + first = make_chunk("first", [1.0, 0.0]) + second = make_chunk("second", [0.0, 1.0]) + replacement = make_chunk("first", [0.0, 1.0]) + third = make_chunk("third", [1.0, 1.0]) + + r.add([first, second]) + r.add([replacement, third]) + + assert len(r) == 3 + assert r._chunks[0] is replacement + assert r._chunks[1] is second + assert r._chunks[2] is third + + def test_duplicate_ids_in_initial_batch_last_occurrence_wins(self): + retriever = InMemoryRetriever() + old = make_chunk("same", [1.0, 0.0]) + other = make_chunk("other", [1.0, 0.0]) + latest = make_chunk("same", [0.0, 1.0]) + + retriever.add([old, other, latest]) + + assert len(retriever) == 2 + assert [chunk.id for chunk in retriever._chunks] == ["same", "other"] + assert retriever._chunks[0] is latest + assert retriever.retrieve([0.0, 1.0], top_k=1) == [latest] + + def test_upsert_after_matrix_consolidation_uses_latest_vector(self): + retriever = InMemoryRetriever() + + old = make_chunk("same", [1.0, 0.0]) + other = make_chunk("other", [0.0, 1.0]) + + retriever.add([old, other]) + retriever.retrieve([1.0, 1.0], top_k=2) + + # The first retrieval consolidates the matrix. + assert retriever._matrix is not None + assert len(retriever._matrix_blocks) == 1 + + latest = make_chunk("same", [0.0, 1.0]) + retriever.add([latest]) + + # The replacement must invalidate the consolidated matrix. + assert retriever._matrix is None + assert len(retriever) == 2 + + results = retriever.retrieve([0.0, 1.0], top_k=2) + + by_id = {chunk.id: chunk for chunk in results} + assert by_id["same"] is latest + assert by_id["same"].content == latest.content + assert by_id["same"].embedding == latest.embedding