From caf6661af988904f3d0cd72f024fd25572a78695 Mon Sep 17 00:00:00 2001 From: Uday Kumar Date: Mon, 28 Sep 2026 11:35:34 +0530 Subject: [PATCH 1/2] fix: add retriever upsert semantics --- ragframework/base.py | 9 +- ragframework/retriever/faiss.py | 354 ++++++++++++++++++++++--- ragframework/retriever/in_memory.py | 56 +++- tests/test_retriever/test_faiss.py | 92 ++++++- tests/test_retriever/test_in_memory.py | 30 +++ 5 files changed, 495 insertions(+), 46 deletions(-) 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/faiss.py b/ragframework/retriever/faiss.py index 93f7de5..32e984c 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,98 +43,386 @@ 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 + not isinstance(value, int) + or isinstance(value, bool) + or value <= 0 for value in (m, ef_construction, ef_search) ): - raise RetrieverError("m, ef_construction, and ef_search must be positive integers.") + raise RetrieverError( + "m, ef_construction, and ef_search must be positive integers." + ) 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] = {} + + 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 embedded chunks into the HNSW index. + + An empty batch is a no-op. - 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. + Every vector in the incoming batch is validated before changing + the Python chunk mapping. + + Chunks with an existing ID replace the previously indexed chunk. + New IDs are appended. + + FAISS stores physical vectors permanently. When a chunk is replaced, + its old physical vector becomes a tombstone and the ID mapping points + to the newly appended vector. """ if not chunks: return expected_dimension = self._dimension + vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + + # --------------------------------------------------------------- + # 1. Validate the COMPLETE incoming batch first. + # --------------------------------------------------------------- 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) 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: + + matrix = np.ascontiguousarray( + np.stack(vectors), + dtype=np.float32, + ) + + # --------------------------------------------------------------- + # 2. Create the index for the first insertion. + # --------------------------------------------------------------- + if self._index is None: + try: new_index = faiss.IndexHNSWFlat( - expected_dimension, self._m, faiss.METRIC_INNER_PRODUCT + 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) + + new_index.add(matrix) + + except Exception as exc: + raise RetrieverError( + "FAISS failed to add chunk embeddings." + ) from exc + + # Only update Python state after FAISS succeeds. + self._index = new_index + self._dimension = expected_dimension + + start_label = 0 + + for offset, chunk in enumerate(chunks): + self._id_to_index[chunk.id] = start_label + offset + self._chunks.append(chunk) + + logger.debug( + "Added chunks count=%d dimension=%d", + len(chunks), + matrix.shape[1], + ) + + return + + # --------------------------------------------------------------- + # 3. Existing index. + # + # Append all incoming vectors to FAISS. + # + # This is intentionally done on the EXISTING index so a native + # FAISS failure leaves the Python mapping unchanged. + # --------------------------------------------------------------- + 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 + raise RetrieverError( + "FAISS failed to add chunk embeddings." + ) from exc + + # --------------------------------------------------------------- + # 4. FAISS succeeded. + # + # Now update the logical mapping. + # --------------------------------------------------------------- + for offset, chunk in enumerate(chunks): + physical_label = start_label + offset + + existing_logical_index = None + + for logical_index, existing_chunk in enumerate(self._chunks): + if existing_chunk.id == chunk.id: + existing_logical_index = logical_index + break + + if existing_logical_index is not None: + # Replace the logical chunk while preserving its position. + self._chunks[existing_logical_index] = chunk + else: + # New logical chunk. + self._chunks.append(chunk) + + # Point the ID to the newly appended physical vector. + self._id_to_index[chunk.id] = physical_label - self._index = new_index self._dimension = expected_dimension - self._chunks.extend(chunks) + logger.debug( "Added chunks count=%d dimension=%d", len(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. + return + + # --------------------------------------------------------------- + # 4. Upsert path. + # + # HNSW does not provide a reliable in-place "replace vector at + # arbitrary logical position" operation for maintaining our + # Python chunk mapping. + # + # Therefore we construct the desired final chunk list and rebuild + # the index. Before rebuilding, we exercise the current FAISS + # index's add operation so FAISS-native failures are surfaced + # without changing the Python chunk mapping. + # --------------------------------------------------------------- + + final_chunks = list(self._chunks) - 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. + existing_indices = { + chunk.id: index + for index, chunk in enumerate(self._chunks) + } + + # Replace existing chunks in their original positions. + for chunk in chunks: + existing_index = existing_indices.get(chunk.id) + + if existing_index is not None: + final_chunks[existing_index] = chunk + + # Append genuinely new chunks. + for chunk in chunks: + if chunk.id not in existing_indices: + final_chunks.append(chunk) + + # --------------------------------------------------------------- + # Build the final normalized matrix. + # --------------------------------------------------------------- + final_vectors: list[ + np.ndarray[Any, np.dtype[np.float32]] + ] = [] + + incoming_by_id: dict[str, np.ndarray[Any, np.dtype[np.float32]]] = { + chunk.id: vector + for chunk, vector in zip(chunks, vectors, strict=True) + } + + for chunk in final_chunks: + if chunk.id in incoming_by_id: + final_vectors.append(incoming_by_id[chunk.id]) + else: + if chunk.embedding is None: + raise RetrieverError( + f"Chunk '{chunk.id}' has no embedding. " + "Embed chunks before adding them to the retriever." + ) + + final_vectors.append( + validate_vector( + chunk.embedding, + f"Chunk '{chunk.id}' embedding", + ) + ) + + final_matrix = np.ascontiguousarray( + np.stack(final_vectors), + dtype=np.float32, + ) + + # --------------------------------------------------------------- + # 5. Build replacement index before changing retriever state. + # --------------------------------------------------------------- + try: + new_index = self._build_index( + final_matrix, + expected_dimension, + ) + except Exception as exc: + raise RetrieverError( + "FAISS failed to add chunk embeddings." + ) from exc + + # --------------------------------------------------------------- + # 6. Commit only after FAISS successfully built the new index. + # --------------------------------------------------------------- + self._index = new_index + self._dimension = expected_dimension + self._chunks = final_chunks + + logger.debug( + "Added chunks count=%d dimension=%d", + len(chunks), + final_matrix.shape[1], + ) + + def retrieve( + self, + query_embedding: list[float], + top_k: int = 5, + ) -> list[Chunk]: + """Return active chunks ordered by approximate cosine similarity. + + Superseded physical FAISS vectors are ignored. + + When multiple active chunks have identical similarity scores, + their logical insertion order is used as a deterministic + tie-breaker. """ 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") + 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)) + # Search all physical vectors because some of them may be + # tombstones belonging to superseded chunks. + count = int(self._index.ntotal) + try: - _scores, labels = self._index.search(query.reshape(1, -1), count) + scores, labels = self._index.search( + query.reshape(1, -1), + count, + ) except Exception as exc: - raise RetrieverError("FAISS failed to search the index.") from exc + raise RetrieverError( + "FAISS failed to search the index." + ) from exc + + # Reverse mapping: + # physical FAISS label -> logical chunk position. + active_labels: dict[int, int] = { + physical_label: logical_index + for logical_index, chunk in enumerate(self._chunks) + for physical_label in [self._id_to_index.get(chunk.id)] + if physical_label is not None + } + + candidates: list[tuple[float, int, Chunk]] = [] + + for score, label in zip( + scores[0], + labels[0], + strict=True, + ): + physical_label = int(label) + + if physical_label < 0: + continue + + logical_index = active_labels.get(physical_label) + + # Ignore tombstoned / unknown physical labels. + if logical_index is None: + continue + + candidates.append( + ( + float(score), + logical_index, + self._chunks[logical_index], + ) + ) + + # Highest cosine similarity first. + # Logical index breaks exact score ties deterministically. + candidates.sort( + key=lambda item: (-item[0], item[1]) + ) return [ - self._chunks[int(label)] for label in labels[0] if 0 <= int(label) < len(self._chunks) + chunk + for _score, _logical_index, chunk in candidates[:top_k] ] def __len__(self) -> int: diff --git a/ragframework/retriever/in_memory.py b/ragframework/retriever/in_memory.py index 104b38e..b1a296c 100644 --- a/ragframework/retriever/in_memory.py +++ b/ragframework/retriever/in_memory.py @@ -25,26 +25,31 @@ class InMemoryRetriever(Retriever): def __init__(self) -> None: self._chunks: list[Chunk] = [] self._matrix: np.ndarray[Any, Any] | None = None # shape (N, dim) + self._id_to_index: dict[str, int] = {} self._dimension: int | None = None def add(self, chunks: list[Chunk]) -> None: - """Validate and normalize a batch before changing the stored chunks. + """Validate and upsert 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. + raise ``RetrieverError`` without adding or replacing any part of the + batch. """ if not chunks: return expected_dimension = self._dimension vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + 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: @@ -52,15 +57,54 @@ def add(self, chunks: list[Chunk]) -> None: f"Chunk '{chunk.id}' embedding has dimension {vector.shape[0]}; " f"expected {expected_dimension}." ) + vectors.append(vector) matrix = np.stack(vectors) - if self._matrix is not None: - matrix = np.concatenate((self._matrix, matrix)) - self._matrix = matrix + # Build the initial index. + if self._matrix is None: + self._matrix = matrix + self._dimension = expected_dimension + + for chunk in chunks: + self._id_to_index[chunk.id] = len(self._chunks) + self._chunks.append(chunk) + + logger.debug( + "Added chunks count=%d dimension=%d", + len(chunks), + matrix.shape[1], + ) + return + + # Apply replacements and collect genuinely new chunks. + new_chunks: list[Chunk] = [] + new_vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + + for chunk, vector in zip(chunks, vectors, strict=True): + existing_index = self._id_to_index.get(chunk.id) + + if existing_index is not None: + self._chunks[existing_index] = chunk + self._matrix[existing_index] = vector + else: + new_chunks.append(chunk) + new_vectors.append(vector) + + if new_chunks: + start_index = len(self._chunks) + self._chunks.extend(new_chunks) + + self._matrix = np.concatenate( + (self._matrix, np.stack(new_vectors)) + ) + + for offset, chunk in enumerate(new_chunks): + self._id_to_index[chunk.id] = start_index + offset + self._dimension = expected_dimension - self._chunks.extend(chunks) + logger.debug( "Added chunks count=%d dimension=%d", len(chunks), diff --git a/tests/test_retriever/test_faiss.py b/tests/test_retriever/test_faiss.py index c180502..fcad298 100644 --- a/tests/test_retriever/test_faiss.py +++ b/tests/test_retriever/test_faiss.py @@ -41,16 +41,84 @@ 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: retriever = FAISSRetriever() @@ -197,18 +265,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 = """ diff --git a/tests/test_retriever/test_in_memory.py b/tests/test_retriever/test_in_memory.py index 88fb664..3f1de41 100644 --- a/tests/test_retriever/test_in_memory.py +++ b/tests/test_retriever/test_in_memory.py @@ -172,3 +172,33 @@ 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 From 93c30893efe5c3ec710f0e625c03b1b93168dc91 Mon Sep 17 00:00:00 2001 From: Uday Kumar Date: Fri, 23 Oct 2026 13:00:59 +0530 Subject: [PATCH 2/2] fix: address retriever upsert review feedback --- CHANGELOG.md | 1 + ragframework/retriever/chroma.py | 25 +- ragframework/retriever/faiss.py | 348 +++++++------------------ ragframework/retriever/in_memory.py | 143 +++++----- tests/test_retriever/test_chroma.py | 68 +++++ tests/test_retriever/test_faiss.py | 107 +++++++- tests/test_retriever/test_in_memory.py | 40 +++ 7 files changed, 402 insertions(+), 330 deletions(-) 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/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 32e984c..5dc2555 100644 --- a/ragframework/retriever/faiss.py +++ b/ragframework/retriever/faiss.py @@ -53,14 +53,10 @@ def __init__( raise ImportError(_INSTALL_HINT) if any( - not isinstance(value, int) - or isinstance(value, bool) - or value <= 0 + not isinstance(value, int) or isinstance(value, bool) or value <= 0 for value in (m, ef_construction, ef_search) ): - raise RetrieverError( - "m, ef_construction, and ef_search must be positive integers." - ) + raise RetrieverError("m, ef_construction, and ef_search must be positive integers.") self._m = m self._ef_construction = ef_construction @@ -74,6 +70,9 @@ def __init__( # 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], @@ -94,30 +93,19 @@ def _build_index( return index def add(self, chunks: list[Chunk]) -> None: - """Validate and upsert embedded chunks into the HNSW index. - - An empty batch is a no-op. - - Every vector in the incoming batch is validated before changing - the Python chunk mapping. + """Validate and upsert chunks; the last duplicate ID in a batch wins. - Chunks with an existing ID replace the previously indexed chunk. - New IDs are appended. - - FAISS stores physical vectors permanently. When a chunk is replaced, - its old physical vector becomes a tombstone and the ID mapping points - to the newly appended vector. + Existing IDs retain their logical position. New IDs are appended. + Replaced physical FAISS vectors become tombstones. """ if not chunks: return expected_dimension = self._dimension + validated: dict[str, tuple[Chunk, np.ndarray[Any, np.dtype[np.float32]]]] = {} - vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] - - # --------------------------------------------------------------- - # 1. Validate the COMPLETE incoming batch first. - # --------------------------------------------------------------- + # 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( @@ -138,292 +126,140 @@ def add(self, chunks: list[Chunk]) -> None: f"{vector.shape[0]}; expected {expected_dimension}." ) - vectors.append(vector) + validated[chunk.id] = (chunk, vector) assert expected_dimension is not None + 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, ) - # --------------------------------------------------------------- - # 2. Create the index for the first insertion. - # --------------------------------------------------------------- + # Build a fresh index for the first insertion. if self._index is None: try: - 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) - + new_index = self._build_index(matrix, expected_dimension) except Exception as exc: - raise RetrieverError( - "FAISS failed to add chunk embeddings." - ) from exc + raise RetrieverError("FAISS failed to add chunk embeddings.") from exc - # Only update Python state after FAISS succeeds. self._index = new_index self._dimension = expected_dimension - start_label = 0 - - for offset, chunk in enumerate(chunks): - self._id_to_index[chunk.id] = start_label + offset + 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 - logger.debug( - "Added chunks count=%d dimension=%d", - len(chunks), - matrix.shape[1], - ) - - return - - # --------------------------------------------------------------- - # 3. Existing index. - # - # Append all incoming vectors to FAISS. - # - # This is intentionally done on the EXISTING index so a native - # FAISS failure leaves the Python mapping unchanged. - # --------------------------------------------------------------- - start_label = int(self._index.ntotal) + 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 - try: - self._index.add(matrix) - except Exception as exc: - raise RetrieverError( - "FAISS failed to add chunk embeddings." - ) from exc - - # --------------------------------------------------------------- - # 4. FAISS succeeded. - # - # Now update the logical mapping. - # --------------------------------------------------------------- - for offset, chunk in enumerate(chunks): - physical_label = start_label + offset - - existing_logical_index = None - - for logical_index, existing_chunk in enumerate(self._chunks): - if existing_chunk.id == chunk.id: - existing_logical_index = logical_index - break + # 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_logical_index is not None: - # Replace the logical chunk while preserving its position. - self._chunks[existing_logical_index] = chunk - else: - # New logical chunk. - self._chunks.append(chunk) + 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 - # Point the ID to the newly appended physical vector. - self._id_to_index[chunk.id] = physical_label + self._id_to_index[chunk.id] = physical_label - self._dimension = expected_dimension + self._dimension = expected_dimension logger.debug( "Added chunks count=%d dimension=%d", - len(chunks), + len(unique_chunks), matrix.shape[1], ) - return - - # --------------------------------------------------------------- - # 4. Upsert path. - # - # HNSW does not provide a reliable in-place "replace vector at - # arbitrary logical position" operation for maintaining our - # Python chunk mapping. - # - # Therefore we construct the desired final chunk list and rebuild - # the index. Before rebuilding, we exercise the current FAISS - # index's add operation so FAISS-native failures are surfaced - # without changing the Python chunk mapping. - # --------------------------------------------------------------- - - final_chunks = list(self._chunks) - - existing_indices = { - chunk.id: index - for index, chunk in enumerate(self._chunks) - } - - # Replace existing chunks in their original positions. - for chunk in chunks: - existing_index = existing_indices.get(chunk.id) - - if existing_index is not None: - final_chunks[existing_index] = chunk - - # Append genuinely new chunks. - for chunk in chunks: - if chunk.id not in existing_indices: - final_chunks.append(chunk) - - # --------------------------------------------------------------- - # Build the final normalized matrix. - # --------------------------------------------------------------- - final_vectors: list[ - np.ndarray[Any, np.dtype[np.float32]] - ] = [] - - incoming_by_id: dict[str, np.ndarray[Any, np.dtype[np.float32]]] = { - chunk.id: vector - for chunk, vector in zip(chunks, vectors, strict=True) - } - - for chunk in final_chunks: - if chunk.id in incoming_by_id: - final_vectors.append(incoming_by_id[chunk.id]) - else: - if chunk.embedding is None: - raise RetrieverError( - f"Chunk '{chunk.id}' has no embedding. " - "Embed chunks before adding them to the retriever." - ) - - final_vectors.append( - validate_vector( - chunk.embedding, - f"Chunk '{chunk.id}' embedding", - ) - ) - - final_matrix = np.ascontiguousarray( - np.stack(final_vectors), - dtype=np.float32, - ) - - # --------------------------------------------------------------- - # 5. Build replacement index before changing retriever state. - # --------------------------------------------------------------- - try: - new_index = self._build_index( - final_matrix, - expected_dimension, - ) - except Exception as exc: - raise RetrieverError( - "FAISS failed to add chunk embeddings." - ) from exc - - # --------------------------------------------------------------- - # 6. Commit only after FAISS successfully built the new index. - # --------------------------------------------------------------- - self._index = new_index - self._dimension = expected_dimension - self._chunks = final_chunks - - logger.debug( - "Added chunks count=%d dimension=%d", - len(chunks), - final_matrix.shape[1], - ) - def retrieve( self, query_embedding: list[float], top_k: int = 5, ) -> list[Chunk]: - """Return active chunks ordered by approximate cosine similarity. - - Superseded physical FAISS vectors are ignored. - - When multiple active chunks have identical similarity scores, - their logical insertion order is used as a deterministic - tie-breaker. - """ + """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 - or not self._chunks - ): + 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", - ) + query = validate_vector(query_embedding, "Query embedding") if query.shape[0] != self._dimension: raise RetrieverError( - f"Query embedding has dimension {query.shape[0]}; " - f"expected {self._dimension}." + f"Query embedding has dimension {query.shape[0]}; " f"expected {self._dimension}." ) - # Search all physical vectors because some of them may be - # tombstones belonging to superseded chunks. - count = int(self._index.ntotal) + total = int(self._index.ntotal) - try: - scores, labels = self._index.search( - query.reshape(1, -1), - count, - ) - except Exception as exc: - raise RetrieverError( - "FAISS failed to search the index." - ) from exc - - # Reverse mapping: - # physical FAISS label -> logical chunk position. - active_labels: dict[int, int] = { - physical_label: logical_index - for logical_index, chunk in enumerate(self._chunks) - for physical_label in [self._id_to_index.get(chunk.id)] - if physical_label is not None + # 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 } - candidates: list[tuple[float, int, Chunk]] = [] + # 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]] = {} - for score, label in zip( - scores[0], - labels[0], - strict=True, - ): - physical_label = int(label) + try: + 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 physical_label < 0: - continue + if len(candidates) >= min(top_k, len(self._chunks)): + break - logical_index = active_labels.get(physical_label) + if search_count >= total: + break - # Ignore tombstoned / unknown physical labels. - if logical_index is None: - continue + search_count = min(total, search_count * 2) - candidates.append( - ( - float(score), - logical_index, - self._chunks[logical_index], - ) - ) + except Exception as exc: + raise RetrieverError("FAISS failed to search the index.") from exc - # Highest cosine similarity first. - # Logical index breaks exact score ties deterministically. - candidates.sort( - key=lambda item: (-item[0], item[1]) + ordered = sorted( + candidates.values(), + key=lambda item: (-item[0], item[1]), ) - return [ - chunk - for _score, _logical_index, chunk in candidates[:top_k] - ] + 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 b1a296c..ad5fbdd 100644 --- a/ragframework/retriever/in_memory.py +++ b/ragframework/retriever/in_memory.py @@ -15,32 +15,26 @@ 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 upsert a batch before changing the stored chunks. - - An empty batch is a no-op. Invalid vectors or inconsistent dimensions - raise ``RetrieverError`` without adding or replacing 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( @@ -54,85 +48,106 @@ def add(self, chunks: list[Chunk]) -> 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) + # Dict assignment makes the last occurrence win. + validated[chunk.id] = (chunk, vector) - # Build the initial index. - if self._matrix is None: - self._matrix = matrix - self._dimension = expected_dimension + assert expected_dimension is not None - for chunk in chunks: - self._id_to_index[chunk.id] = len(self._chunks) - self._chunks.append(chunk) + # 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] - logger.debug( - "Added chunks count=%d dimension=%d", - len(chunks), - matrix.shape[1], - ) - return + 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, + ) - # Apply replacements and collect genuinely new chunks. - new_chunks: list[Chunk] = [] - new_vectors: list[np.ndarray[Any, np.dtype[np.float32]]] = [] + block_index = len(final_blocks) + final_blocks.append(new_matrix) + final_block_ids.append(new_ids) - for chunk, vector in zip(chunks, vectors, strict=True): - existing_index = self._id_to_index.get(chunk.id) + for row_index, (chunk, _) in enumerate(new_chunks): + position = final_positions.get(chunk.id) - if existing_index is not None: - self._chunks[existing_index] = chunk - self._matrix[existing_index] = vector + if position is None: + position = len(final_chunks) + final_positions[chunk.id] = position + final_chunks.append(chunk) else: - new_chunks.append(chunk) - new_vectors.append(vector) + final_chunks[position] = chunk - if new_chunks: - start_index = len(self._chunks) - self._chunks.extend(new_chunks) - - self._matrix = np.concatenate( - (self._matrix, np.stack(new_vectors)) - ) - - for offset, chunk in enumerate(new_chunks): - self._id_to_index[chunk.id] = start_index + offset + # 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 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 + + 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]) - 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``. - """ + 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 fcad298..69c3bf9 100644 --- a/tests/test_retriever/test_faiss.py +++ b/tests/test_retriever/test_faiss.py @@ -61,6 +61,7 @@ def test_batches_upsert_without_duplicate_ids() -> None: 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) @@ -80,6 +81,7 @@ def test_add_batch_replaces_existing_and_adds_new_chunks() -> None: 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) @@ -95,7 +97,6 @@ def test_add_same_id_replaces_existing_chunk() -> None: 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: @@ -120,6 +121,7 @@ def fail(_vectors: np.ndarray) -> None: results = retriever.retrieve([1.0, 0.0], top_k=1) assert results == [existing] + def test_add_logs_count_and_dimension(caplog) -> None: retriever = FAISSRetriever() chunks = [ @@ -310,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 3f1de41..b1ef56f 100644 --- a/tests/test_retriever/test_in_memory.py +++ b/tests/test_retriever/test_in_memory.py @@ -202,3 +202,43 @@ def test_add_batch_replaces_existing_and_adds_new_chunks(self): 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