diff --git a/CHANGELOG.md b/CHANGELOG.md index 077e103..d84758b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,16 +8,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added - - `TokenChunker` for token-level chunking using tiktoken with a new `[tokens]` optional extra (closes #44) - - Structured logging for ingest, query, and retriever stages, including timing and empty-ingest warnings (closes #41). - `RecursiveChunker.from_config()` for creating a chunker from pipeline configuration (closes #25). ### Fixed +- Optimize `InMemoryRetriever.add()` to normalize only newly added vectors and avoid rebuilding the full matrix on every call (closes #27). - `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). + ## [0.3.0] - 2026-09-25 ### Added @@ -37,6 +37,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Enforce LF line endings with `.gitattributes` across platforms while keeping PNG files binary (closes #48). - `InMemoryRetriever.retrieve()` now returns an empty list for non-positive `top_k` values and raises `RetrieverError` for non-integer or boolean `top_k` values (closes #23). - `InMemoryRetriever` now raises `RetrieverError` for invalid vectors and dimension mismatches, validates complete batches before updating stored data, and treats empty batches as a no-op. Vector validation and normalization are shared with `FAISSRetriever` (closes #26). +### Fixed +- Optimize `InMemoryRetriever.add()` to normalize only newly added vectors and avoid rebuilding the full matrix on every call (closes #27). ## [0.2.0] - 2026-09-19 diff --git a/ragframework/retriever/in_memory.py b/ragframework/retriever/in_memory.py index 104b38e..0d37340 100644 --- a/ragframework/retriever/in_memory.py +++ b/ragframework/retriever/in_memory.py @@ -25,6 +25,7 @@ class InMemoryRetriever(Retriever): def __init__(self) -> None: self._chunks: list[Chunk] = [] self._matrix: np.ndarray[Any, Any] | None = None # shape (N, dim) + self._matrix_blocks: list[np.ndarray[Any, Any]] = [] self._dimension: int | None = None def add(self, chunks: list[Chunk]) -> None: @@ -55,10 +56,9 @@ def add(self, chunks: list[Chunk]) -> None: vectors.append(vector) matrix = np.stack(vectors) - if self._matrix is not None: - matrix = np.concatenate((self._matrix, matrix)) + self._matrix_blocks.append(matrix) - self._matrix = matrix + self._matrix = None self._dimension = expected_dimension self._chunks.extend(chunks) logger.debug( @@ -77,13 +77,23 @@ def retrieve(self, query_embedding: list[float], top_k: int = 5) -> list[Chunk]: """ 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}." ) + if self._matrix is None: + if len(self._matrix_blocks) == 1: + # Reuse the existing array; do not create another copy. + self._matrix = self._matrix_blocks[0] + else: + # Consolidate all blocks into one cached matrix and replace the + # old block references so the original arrays can be released. + self._matrix = np.concatenate(self._matrix_blocks) + self._matrix_blocks = [self._matrix] + scores: np.ndarray[Any, Any] = self._matrix @ q k = min(top_k, len(self._chunks)) top_indices = np.argpartition(scores, -k)[-k:] diff --git a/tests/test_document/test_chunkers.py b/tests/test_document/test_chunkers.py index bc0ad00..1835688 100644 --- a/tests/test_document/test_chunkers.py +++ b/tests/test_document/test_chunkers.py @@ -3,6 +3,7 @@ import builtins import sys import types +from typing import Any import pytest @@ -12,7 +13,6 @@ RecursiveChunker, SentenceChunker, ) -from typing import Any def test_recursive_chunker_from_config(): diff --git a/tests/test_retriever/test_in_memory.py b/tests/test_retriever/test_in_memory.py index 88fb664..86204e5 100644 --- a/tests/test_retriever/test_in_memory.py +++ b/tests/test_retriever/test_in_memory.py @@ -172,3 +172,75 @@ 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_batched_add_matches_single_batch(self): + chunks = [ + make_chunk("a", [1.0, 0.0]), + make_chunk("b", [0.0, 1.0]), + make_chunk("c", [1.0, 1.0]), + make_chunk("d", [-1.0, 0.0]), + ] + + single_batch = InMemoryRetriever() + single_batch.add(chunks) + + batched = InMemoryRetriever() + batched.add(chunks[:2]) + batched.add(chunks[2:]) + + query = [1.0, 1.0] + + single_results = single_batch.retrieve(query, top_k=4) + batched_results = batched.retrieve(query, top_k=4) + + assert len(batched_results) == len(single_results) + assert batched_results[0].id == single_results[0].id + +def test_retrieve_consolidates_matrix_blocks_without_duplicate_storage(): + retriever = InMemoryRetriever() + + chunks = [ + make_chunk("a", [1.0, 0.0]), + make_chunk("b", [0.0, 1.0]), + make_chunk("c", [1.0, 1.0]), + ] + + retriever.add(chunks[:1]) + retriever.add(chunks[1:]) + + assert len(retriever._matrix_blocks) == 2 + assert retriever._matrix is None + + retriever.retrieve([1.0, 1.0], top_k=3) + + assert retriever._matrix is not None + assert len(retriever._matrix_blocks) == 1 + assert retriever._matrix_blocks[0] is retriever._matrix + +def test_add_after_consolidation_preserves_all_results(): + retriever = InMemoryRetriever() + + first = [ + make_chunk("a", [1.0, 0.0]), + make_chunk("b", [0.0, 1.0]), + ] + second = [ + make_chunk("c", [1.0, 1.0]), + ] + + retriever.add(first) + retriever.retrieve([1.0, 1.0], top_k=2) + + assert len(retriever._matrix_blocks) == 1 + assert retriever._matrix_blocks[0] is retriever._matrix + + retriever.add(second) + + assert retriever._matrix is None + assert len(retriever._matrix_blocks) == 2 + + results = retriever.retrieve([1.0, 1.0], top_k=3) + + assert [chunk.id for chunk in results] == ["c", "b", "a"] + assert len(retriever._matrix_blocks) == 1 + assert retriever._matrix_blocks[0] is retriever._matrix