Merge b3cdc52b44 into c1de31ab2c

reme/components/file_store/faiss_local_file_store.py

@@ -3,6 +3,7 @@
 import asyncio
 import hashlib
 import json
+from collections.abc import Callable
 from contextlib import suppress
 from uuid import uuid4
 
@@ -664,6 +665,7 @@ class FaissLocalFileStore(LocalFileStore):
             return []
 
         q = self._prepare(query_embedding)
+        matches = self._prepare_search_filter(search_filter)
         ntotal = index.ntotal
 
         # Small-index shortcut: when ntotal is below the brute-force threshold,
@@ -683,14 +685,14 @@ class FaissLocalFileStore(LocalFileStore):
         if ntotal < max(1, (max(limit, ef_search)) ** 0.5) * self.hnsw_m:
             k = ntotal if search_filter else min(ntotal, limit + len(self._tombstones))
             scores, rows = index.storage.search(q, k)
-            return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter)
+            return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter, matches=matches)
 
         if not search_filter:
             # No filter: simple over-fetch to cover tombstones.
             k = min(ntotal, limit + len(self._tombstones))
             self._set_ef_search(index, limit)
             scores, rows = index.search(q, k)
-            return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter)
+            return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter, matches=matches)
 
         # With filter: progressively increase k until we collect enough results
         # or exhaust the entire index.
@@ -698,7 +700,7 @@ class FaissLocalFileStore(LocalFileStore):
         while True:
             self._set_ef_search(index, limit)
             scores, rows = index.search(q, k)
-            results = self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter)
+            results = self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter, matches=matches)
             if len(results) >= limit or k >= ntotal:
                 return results
             k = min(ntotal, k * 2)
@@ -709,15 +711,19 @@ class FaissLocalFileStore(LocalFileStore):
         scores: list[float],
         limit: int,
         search_filter: dict | None = None,
+        *,
+        matches: Callable[[FileChunk], bool] | None = None,
     ) -> list[FileChunk]:
         """Map raw FAISS rows back to chunks, skipping tombstones and stale ids."""
+        if matches is None:
+            matches = self._prepare_search_filter(search_filter)
         results: list[FileChunk] = []
         for raw_row, score in zip(rows, scores):
             row = int(raw_row)
             if row < 0 or row in self._tombstones or row >= len(self._id_map):
                 continue
             chunk = self.file_chunks.get(self._id_map[row])
-            if chunk is None or not self._matches_search_filter(chunk, search_filter):
+            if chunk is None or not matches(chunk):
                 continue
             results.append(chunk.model_copy(update={"scores": {"vector": float(score), "score": float(score)}}))
             if len(results) >= limit:

reme/components/file_store/local_file_store.py

@@ -6,7 +6,7 @@ import datetime
 import heapq
 import json
 import time
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from contextlib import suppress
 
 import numpy as np
@@ -941,6 +941,7 @@ class LocalFileStore(BaseFileStore):
         if query_embedding is None:
             return []
 
+        matches = self._prepare_search_filter(search_filter)
         top: list[tuple[float, int, FileChunk]] = []
         candidates: list[FileChunk] = []
         embeddings: list[np.ndarray] = []
@@ -964,10 +965,7 @@ class LocalFileStore(BaseFileStore):
             embeddings.clear()
 
         for candidate in self.file_chunks.values():
-            if not self._embedding_dim_matches(candidate.embedding) or not self._matches_search_filter(
-                candidate,
-                search_filter,
-            ):
+            if not self._embedding_dim_matches(candidate.embedding) or not matches(candidate):
                 continue
             candidates.append(candidate)
             embeddings.append(candidate.embedding)
@@ -989,10 +987,9 @@ class LocalFileStore(BaseFileStore):
         if not query:
             return []
 
+        matches = self._prepare_search_filter(search_filter)
         if search_filter:
-            eligible_ids = {
-                chunk.id for chunk in self.file_chunks.values() if self._matches_search_filter(chunk, search_filter)
-            }
+            eligible_ids = {chunk.id for chunk in self.file_chunks.values() if matches(chunk)}
             if not eligible_ids:
                 return []
             try:
@@ -1011,7 +1008,7 @@ class LocalFileStore(BaseFileStore):
         results = []
         for doc_id, score in doc_id_score_dict.items():
             chunk = self.file_chunks.get(doc_id)
-            if chunk and self._matches_search_filter(chunk, search_filter):
+            if chunk and matches(chunk):
                 results.append(chunk.model_copy(update={"scores": {"keyword": score, "score": score}}))
                 if len(results) >= limit:
                     break
@@ -1020,6 +1017,78 @@ class LocalFileStore(BaseFileStore):
 
     # -- extensions -----------------------------------------------------------
 
+    @classmethod
+    def _prepare_search_filter(cls, search_filter: dict | None) -> Callable[[FileChunk], bool]:
+        """Prepare query-local membership tests; retain the dict-based compatibility path."""
+
+        def legacy(chunk):
+            return cls._matches_search_filter(chunk, search_filter)
+
+        # Subclasses may customize matching without implementing this new helper.
+        if (
+            cls._matches_search_filter.__func__ is not LocalFileStore._matches_search_filter.__func__
+            or cls._value_matches.__func__ is not LocalFileStore._value_matches.__func__
+        ):
+            return legacy
+        if not search_filter:
+            return lambda chunk: True
+        try:
+            exact_paths = set()
+            has_exact_paths = False
+            for key in ("path", "paths"):
+                if key in search_filter:
+                    has_exact_paths = True
+                    exact_paths.update(cls._as_filter_values(search_filter[key]))
+            prefixes = tuple(
+                str(value)
+                for key in ("path_prefix", "path_prefixes", "prefix", "prefixes")
+                if key in search_filter
+                for value in cls._as_filter_values(search_filter[key])
+            )
+            start_date = search_filter.get("start_date")
+            end_date = search_filter.get("end_date")
+            strict_date = bool(search_filter.get("strict_date_filter", False))
+            metadata = dict(search_filter.get("metadata") or {})
+            reserved = {
+                "path",
+                "paths",
+                "path_prefix",
+                "path_prefixes",
+                "prefix",
+                "prefixes",
+                "metadata",
+                "start_date",
+                "end_date",
+                "strict_date_filter",
+            }
+            metadata.update((key, value) for key, value in search_filter.items() if key not in reserved)
+            conditions = []
+            for key, value in metadata.items():
+                collection = isinstance(value, (list, tuple, set, frozenset))
+                conditions.append((key, frozenset(value) if collection else value, collection))
+        except (TypeError, ValueError):
+            # Malformed late-stage conditions must still short-circuit as before.
+            return legacy
+
+        def matches(chunk: FileChunk) -> bool:
+            if has_exact_paths and chunk.path not in exact_paths:
+                return False
+            if prefixes and not chunk.path.startswith(prefixes):
+                return False
+            if start_date or end_date:
+                path_date = cls._extract_date_from_path(chunk.path)
+                if not path_date:
+                    if strict_date:
+                        return False
+                elif (start_date and path_date < start_date) or (end_date and path_date > end_date):
+                    return False
+            return all(
+                chunk.metadata.get(key) in expected if collection else chunk.metadata.get(key) == expected
+                for key, expected, collection in conditions
+            )
+
+        return matches
+
     @staticmethod
     def _as_filter_values(value) -> set:
         if isinstance(value, (list, tuple, set, frozenset)):

reme/components/file_store/zvec_local_file_store.py

@@ -4,6 +4,7 @@ import hashlib
 import json
 import shutil
 import time
+from collections.abc import Callable
 from uuid import uuid4
 
 import aiofiles
@@ -455,18 +456,19 @@ class ZvecLocalFileStore(LocalFileStore):
             return []
 
         vector = np.asarray(query_embedding, dtype=np.float32).tolist()
+        matches = self._prepare_search_filter(search_filter)
         ntotal = len(self._indexed_ids)
 
         if not search_filter:
             hits = self._query_collection(collection, vector, min(limit, ntotal))
-            return self._collect_hits(hits, limit, search_filter)
+            return self._collect_hits(hits, limit, search_filter, matches=matches)
 
         # With a post-filter: progressively increase k until we collect enough
         # results or exhaust the reachable index.
         k = min(ntotal, 3 * limit)
         while True:
             hits = self._query_collection(collection, vector, k)
-            results = self._collect_hits(hits, limit, search_filter)
+            results = self._collect_hits(hits, limit, search_filter, matches=matches)
             if len(results) >= limit or k >= ntotal:
                 return results
             k = min(ntotal, k * 2)
@@ -484,12 +486,16 @@ class ZvecLocalFileStore(LocalFileStore):
         hits: list[tuple[str, float]],
         limit: int,
         search_filter: dict | None = None,
+        *,
+        matches: Callable[[FileChunk], bool] | None = None,
     ) -> list[FileChunk]:
         """Map zvec hits back to chunks, skipping stale ids and filtered-out chunks."""
+        if matches is None:
+            matches = self._prepare_search_filter(search_filter)
         results: list[FileChunk] = []
         for chunk_id, score in hits:
             chunk = self.file_chunks.get(chunk_id)
-            if chunk is None or not self._matches_search_filter(chunk, search_filter):
+            if chunk is None or not matches(chunk):
                 continue
             results.append(chunk.model_copy(update={"scores": {"vector": score, "score": score}}))
             if len(results) >= limit:

scripts/benchmark_search_filter.py

@@ -0,0 +1,113 @@
+"""Compare prepared filtering with the unchanged dict matcher on real searches.
+
+Run: python -m scripts.benchmark_search_filter --repeats 7
+Setup is excluded; deterministic chunks and fake query embeddings need no files/network.
+"""
+
+# pylint: disable=protected-access
+
+import argparse
+import asyncio
+import json
+import platform
+import statistics
+import time
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import numpy as np
+
+from reme.components.file_store import LocalFileStore
+from reme.components.keyword_index import BM25Index
+from reme.components.tokenizer import RegexTokenizer
+from reme.schema import FileChunk
+
+
+async def measure(operation, config, repeats):
+    """Warm up and report per-call latency in milliseconds."""
+    for _ in range(2):
+        await operation(config)
+    samples = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        await operation(config)
+        samples.append((time.perf_counter() - start) * 1000)
+    return {"median_ms": statistics.median(samples), "p95_ms": float(np.percentile(samples, 95))}
+
+
+async def benchmark(size, repeats):
+    """Measure predicate traversal, keyword_search and vector_search separately."""
+    store = LocalFileStore(embedding_store="")
+    index = BM25Index()
+    index.tokenizer = RegexTokenizer(filter_stopwords=False)
+    store.keyword_index = index
+    store.embedding_store = SimpleNamespace(dimensions=2, is_healthy=True)
+    store._get_query_embedding = AsyncMock(return_value=np.array([1.0, 0.0]))
+    for i in range(size):
+        store.file_chunks[str(i)] = FileChunk(
+            id=str(i),
+            path=f"daily/{i % 2000}.md",
+            text="alpha beta",
+            embedding=np.array([1.0, 0.0]),
+        )
+    await index.add_docs({key: chunk.text for key, chunk in store.file_chunks.items()})
+
+    async def predicate(config):
+        matches = store._prepare_search_filter(config)
+        return [chunk.id for chunk in store.file_chunks.values() if matches(chunk)]
+
+    async def keyword(config):
+        return [(c.id, c.scores) for c in await store.keyword_search("alpha", 10, config)]
+
+    async def vector(config):
+        return [(c.id, c.scores) for c in await store.vector_search("alpha", 10, config)]
+
+    def legacy(search_filter):
+        return lambda chunk: store._matches_search_filter(chunk, search_filter)
+
+    results = []
+    for path_count in (1, 100, 1000):
+        config = {"paths": [f"daily/{i}.md" for i in range(path_count)]}
+        for name, operation in (("predicate", predicate), ("keyword_search", keyword), ("vector_search", vector)):
+            expected = await operation(config)
+            prepared = await measure(operation, config, repeats)
+            with patch.object(store, "_prepare_search_filter", side_effect=legacy):
+                assert await operation(config) == expected
+                baseline = await measure(operation, config, repeats)
+            results.append(
+                {
+                    "chunks": size,
+                    "paths": path_count,
+                    "operation": name,
+                    "baseline": baseline,
+                    "prepared": prepared,
+                },
+            )
+    return results
+
+
+async def main():
+    """Print reproducible environment and results as JSON."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--sizes", nargs="+", type=int, default=[20000, 100000])
+    parser.add_argument("--repeats", type=int, default=7)
+    args = parser.parse_args()
+    rows = []
+    for size in args.sizes:
+        rows.extend(await benchmark(size, args.repeats))
+    print(
+        json.dumps(
+            {
+                "python": platform.python_version(),
+                "platform": platform.platform(),
+                "numpy": np.__version__,
+                "repeats": args.repeats,
+                "results": rows,
+            },
+            indent=2,
+        ),
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/unit/test_prepared_search_filter.py

@@ -0,0 +1,176 @@
+"""Prepared filters preserve the dict matcher and all search entrypoints."""
+
+# pylint: disable=protected-access,missing-function-docstring
+
+import copy
+import random
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock, patch
+
+import numpy as np
+import pytest
+
+from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore
+from reme.components.keyword_index import BM25Index
+from reme.components.tokenizer import RegexTokenizer
+from reme.schema import FileChunk
+
+
+def test_differential_filter_combinations():
+    rng = random.Random(714)
+    paths = ["daily/2026-01-01/n.md", "daily/2026-02-01.md", "digest/topic.md", "daily/invalid.md"]
+    chunks = [
+        FileChunk(id=str(i), path=rng.choice(paths), metadata={"group": rng.choice([None, "a", "b"]), "n": i % 3})
+        for i in range(80)
+    ]
+    path_filters = [{}, {"paths": []}, {"path": paths[0]}, {"path": paths[0], "paths": paths[1:]}]
+    prefix_filters = [{}, {"prefixes": []}, {"prefix": "daily/"}, {"path_prefix": "digest", "prefixes": ["daily/"]}]
+    date_filters = [
+        {},
+        {"strict_date_filter": True},
+        {"start_date": "2026-01-15"},
+        {
+            "end_date": "2026-01-31",
+            "strict_date_filter": True,
+        },
+    ]
+    metadata_filters = [
+        {},
+        {"group": ["a", "b"]},
+        {"metadata": {"group": "a"}, "group": "b"},
+        {
+            "metadata": {"n": (1, 2), "absent": None},
+        },
+    ]
+    for _ in range(250):
+        config = {
+            **rng.choice(path_filters),
+            **rng.choice(prefix_filters),
+            **rng.choice(date_filters),
+            **rng.choice(metadata_filters),
+        }
+        original = copy.deepcopy(config)
+        prepared = LocalFileStore._prepare_search_filter(config)
+        assert [prepared(chunk) for chunk in chunks] == [
+            LocalFileStore._matches_search_filter(chunk, config) for chunk in chunks
+        ]
+        assert config == original
+    assert LocalFileStore._prepare_search_filter(None)(chunks[0])
+
+
+def test_preparation_is_query_local_and_preserves_subclass_hooks():
+    candidate = FileChunk(path="one")
+    config = {"paths": ["one"]}
+    prepared = LocalFileStore._prepare_search_filter(config)
+    config["paths"][:] = ["two"]
+    assert prepared(candidate)
+    assert not LocalFileStore._prepare_search_filter(config)(candidate)
+
+    class CustomStore(LocalFileStore):
+        """Existing subclass hook still governs the new hot path."""
+
+        @classmethod
+        def _matches_search_filter(cls, chunk, search_filter):
+            return chunk.path == "special"
+
+    assert not CustomStore._prepare_search_filter({})(candidate)
+    assert CustomStore._prepare_search_filter({})(FileChunk(path="special"))
+
+
+@pytest.mark.parametrize(
+    "config",
+    [
+        {"paths": [], "metadata": {"bad": [[1]]}},
+        {"paths": ["one"], "metadata": {"bad": [[1]]}},
+        {"prefix": "else", "metadata": 123},
+        {"paths": [[1]]},
+    ],
+)
+def test_malformed_conditions_keep_short_circuit_behavior(config):
+    candidate = FileChunk(path="one")
+    prepared = LocalFileStore._prepare_search_filter(config)
+    try:
+        expected = LocalFileStore._matches_search_filter(candidate, config)
+    except (TypeError, ValueError) as error:
+        with pytest.raises(type(error)):
+            prepared(candidate)
+    else:
+        assert prepared(candidate) == expected
+
+
+def make_store(cls=LocalFileStore):
+    # Native ANN boundaries are faked below; execute the real backend search methods.
+    store = object.__new__(cls)
+    LocalFileStore.__init__(store, embedding_store="")
+    store.file_chunks = {
+        str(i): FileChunk(id=str(i), path=f"daily/{i}.md", text="alpha beta", embedding=np.array([1.0, 0.0]))
+        for i in range(30)
+    }
+    store.embedding_store = SimpleNamespace(dimensions=2, is_healthy=True)
+    store._get_query_embedding = AsyncMock(return_value=np.array([1.0, 0.0]))
+    return store
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fallback", [False, True])
+async def test_keyword_prepares_once_and_preserves_fallback(fallback):
+    store = make_store()
+    index = BM25Index()
+    index.tokenizer = RegexTokenizer(filter_stopwords=False)
+    await index.add_docs({key: chunk.text for key, chunk in store.file_chunks.items()})
+    store.keyword_index = index
+    if fallback:
+        index.retrieve_filtered = AsyncMock(side_effect=NotImplementedError)
+    config = {"paths": ["daily/28.md", "daily/29.md"]}
+    with patch.object(store, "_prepare_search_filter", wraps=store._prepare_search_filter) as prepare:
+        result = await store.keyword_search("alpha", 2, config)
+        assert prepare.call_count == 1
+    with patch.object(
+        store,
+        "_prepare_search_filter",
+        side_effect=lambda filt: lambda c: store._matches_search_filter(c, filt),
+    ):
+        original = await store.keyword_search("alpha", 2, config)
+    assert [(c.id, c.scores) for c in result] == [(c.id, c.scores) for c in original]
+    assert {c.path for c in result} == set(config["paths"])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("backend", [LocalFileStore, FaissLocalFileStore, ZvecLocalFileStore])
+async def test_vector_prepares_once_including_progressive_recall(backend):
+    store = make_store(backend)
+    searches = []
+    if backend is FaissLocalFileStore:
+
+        def search(_query, count):
+            searches.append(count)
+            return np.ones((1, count)), np.arange(count).reshape(1, -1)
+
+        store._faiss_index = SimpleNamespace(ntotal=30, search=search, hnsw=SimpleNamespace(efSearch=0))
+        store._prepare = Mock(side_effect=lambda value: value.reshape(1, -1))
+        store.hnsw_m = 1
+        store._id_map = list(store.file_chunks)
+        store._tombstones = set()
+    elif backend is ZvecLocalFileStore:
+
+        def query(_collection, _vector, count):
+            searches.append(count)
+            return [(str(i), 1.0) for i in range(count)]
+
+        store._collection = object()
+        store._indexed_ids = set(store.file_chunks)
+        store._query_collection = Mock(side_effect=query)
+    config = {"paths": ["daily/29.md"]}
+    with patch.object(store, "_prepare_search_filter", wraps=store._prepare_search_filter) as prepare:
+        result = await store.vector_search("alpha", 1, config)
+        assert prepare.call_count == 1
+    if backend is not LocalFileStore:
+        assert searches == [3, 6, 12, 24, 30]
+    with patch.object(
+        store,
+        "_prepare_search_filter",
+        side_effect=lambda filt: lambda c: store._matches_search_filter(c, filt),
+    ):
+        original = await store.vector_search("alpha", 1, config)
+    assert [(c.id, c.scores) for c in result] == [(c.id, c.scores) for c in original]
+    assert [c.id for c in result] == ["29"]