From cb82d16842822ceb5398fa9166521b785047e5aa Mon Sep 17 00:00:00 2001 From: Uday Kumar Date: Sun, 27 Sep 2026 19:37:14 +0530 Subject: [PATCH 1/3] fix(retriever): avoid re-normalizing full in-memory matrix --- CHANGELOG.md | 5 +++++ ragframework/retriever/in_memory.py | 11 +++++++---- tests/test_retriever/test_in_memory.py | 23 +++++++++++++++++++++++ 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e6be596..1d625cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- Optimize `InMemoryRetriever.add()` to normalize only newly added vectors and avoid rebuilding the full matrix on every call (closes #27). + ## [0.3.0] - 2026-09-25 ### Added @@ -26,6 +29,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `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). - Enforce LF line endings with `.gitattributes` across platforms while keeping PNG files binary (closes #48). - `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 6393f5f..a986561 100644 --- a/ragframework/retriever/in_memory.py +++ b/ragframework/retriever/in_memory.py @@ -22,6 +22,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: @@ -52,10 +53,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) @@ -69,13 +69,16 @@ 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: + self._matrix = np.concatenate(self._matrix_blocks) + 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_retriever/test_in_memory.py b/tests/test_retriever/test_in_memory.py index 16b558b..e134a2f 100644 --- a/tests/test_retriever/test_in_memory.py +++ b/tests/test_retriever/test_in_memory.py @@ -158,3 +158,26 @@ 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 From 199e73c1f4a982c80a7656a732c9679346a9b529 Mon Sep 17 00:00:00 2001 From: Uday Kumar Date: Mon, 28 Sep 2026 12:00:53 +0530 Subject: [PATCH 2/3] fix: consolidate in-memory matrix blocks --- ragframework/retriever/in_memory.py | 9 ++++- tests/test_retriever/test_in_memory.py | 49 ++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/ragframework/retriever/in_memory.py b/ragframework/retriever/in_memory.py index a986561..24e29a1 100644 --- a/ragframework/retriever/in_memory.py +++ b/ragframework/retriever/in_memory.py @@ -77,7 +77,14 @@ def retrieve(self, query_embedding: list[float], top_k: int = 5) -> list[Chunk]: f"Query embedding has dimension {q.shape[0]}; expected {self._dimension}." ) if self._matrix is None: - self._matrix = np.concatenate(self._matrix_blocks) + 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)) diff --git a/tests/test_retriever/test_in_memory.py b/tests/test_retriever/test_in_memory.py index e134a2f..858cc41 100644 --- a/tests/test_retriever/test_in_memory.py +++ b/tests/test_retriever/test_in_memory.py @@ -181,3 +181,52 @@ def test_batched_add_matches_single_batch(self): 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 From ddd0998659df1f11e44cae795a8874283da90137 Mon Sep 17 00:00:00 2001 From: Uday Kumar Date: Mon, 28 Sep 2026 12:08:20 +0530 Subject: [PATCH 3/3] style: fix test import ordering --- tests/test_document/test_chunkers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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():