diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index b187b4053..03a00b839 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -24,8 +24,10 @@ from basic_memory.repository.script_ngrams import analyze_script_query, build_script_ngrams from basic_memory.repository.semantic_chunking import VectorChunkRecord from basic_memory.repository.search_repository_base import ( + SearchIndexKey, SearchRepositoryBase, VectorChunkState, + candidate_key_restriction_condition, file_path_prefix_condition, metadata_contains_like_condition, metadata_filter_content_type_condition, @@ -980,6 +982,7 @@ async def _build_fts_query_parts( file_path_prefix: Optional[str] = None, temporal: Optional[TemporalFilter] = None, allow_relaxed: bool = False, + candidate_keys: Sequence[SearchIndexKey] | None = None, ) -> tuple[str, str, dict[str, Any], str, str]: """Build Postgres FTS FROM/WHERE params shared by search and count.""" conditions = [] @@ -1148,6 +1151,13 @@ async def _build_fts_query_parts( if subtree_condition is not None: conditions.append(subtree_condition) + # Handle an explicit candidate-row restriction. Built by the shared helper so + # both backends restrict by the identical rule; see + # candidate_key_restriction_condition for why the vector filter pass asks about + # its candidates rather than paging the filter's whole match set (#1431). + if candidate_keys is not None: + conditions.append(candidate_key_restriction_condition(candidate_keys, params)) + # Handle search item type filter (parameterized for defense-in-depth) if search_item_types: type_placeholders = [] @@ -1394,6 +1404,7 @@ async def search( allow_relaxed: bool = False, session: AsyncSession | None = None, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content using PostgreSQL tsvector.""" @@ -1439,6 +1450,7 @@ async def search( file_path_prefix=file_path_prefix, temporal=temporal, allow_relaxed=allow_relaxed, + candidate_keys=candidate_keys, ) # set limit and offset diff --git a/src/basic_memory/repository/search_repository.py b/src/basic_memory/repository/search_repository.py index 4d7685128..dae254b8a 100644 --- a/src/basic_memory/repository/search_repository.py +++ b/src/basic_memory/repository/search_repository.py @@ -18,7 +18,7 @@ from basic_memory.repository.rerank_provider_factory import create_rerank_provider from basic_memory.repository.postgres_search_repository import PostgresSearchRepository from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import ChunkManifestRow +from basic_memory.repository.search_repository_base import ChunkManifestRow, SearchIndexKey from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.semantic_vector_index_factory import ( create_semantic_vector_index, @@ -89,6 +89,7 @@ async def search( allow_relaxed: bool = False, session: AsyncSession | None = None, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across indexed content.""" diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index acc212f40..d926bd7c9 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -89,6 +89,9 @@ # --- Semantic search constants --- VECTOR_FILTER_SCAN_LIMIT = 50000 +# The shared bind-parameter bound for any statement that carries a list of vector +# candidate keys. Both engines cap bind parameters (asyncpg at 32767), so every such +# list — manifest hydration and the filter intersection alike — is split at this size. VECTOR_HYDRATION_BATCH_SIZE = 250 # Over-fetch factor for the rerank candidate chunk pool: chunks collapse to unique # (type, id) rows before reranking, so fetch several times reranker_candidates chunks @@ -258,6 +261,56 @@ def metadata_contains_like_condition( ) +def candidate_key_restriction_condition( + candidate_keys: Sequence[SearchIndexKey], + params: Dict[str, Any], +) -> str: + """Build the SQL restricting a filter query to an explicit set of search rows. + + This is what turns the vector/hybrid filter pass from "give me a page of everything + the filter admits" into "of *these* candidates, which does the filter admit". The + first question has an answer the size of the project and had to be capped, and every + candidate outside the cap was then read as disallowed (#1431). The second question's + answer is bounded by the candidate set itself, so no cap is needed and none of the + candidates can fall off the end. + + Keys are grouped by row type rather than emitted as one ``(type, id)`` pair per + branch: entity, observation, and relation ids come from independent sequences, so the + type is part of the identity, but a handful of type-scoped ``IN`` lists binds one + parameter per key instead of two and leaves the id list in the shape both planners + can drive an index from. PostgreSQL's ``search_index`` primary key is + ``(id, type, project_id)``. + + An empty candidate set is a real state, not a caller error — a vector search whose + every hit was already dropped — and it admits nothing, so it yields a false + predicate rather than the vacuous truth an empty ``OR`` would collapse to. + + Shared verbatim by both backends for the same reason the subtree scope is: a + restriction that admitted different rows per dialect would give semantic search a + different candidate set depending on which database happened to be underneath. + """ + ids_by_type: dict[str, list[int]] = {} + for row_type, row_id in candidate_keys: + ids_by_type.setdefault(row_type, []).append(row_id) + + branches: list[str] = [] + for type_index, (row_type, row_ids) in enumerate(ids_by_type.items()): + type_param = f"candidate_type_{type_index}" + params[type_param] = row_type + id_params: list[str] = [] + for id_index, row_id in enumerate(dict.fromkeys(row_ids)): + id_param = f"candidate_id_{type_index}_{id_index}" + params[id_param] = row_id + id_params.append(f":{id_param}") + branches.append( + f"(search_index.type = :{type_param} AND search_index.id IN ({', '.join(id_params)}))" + ) + + if not branches: + return "1 = 0" + return f"({' OR '.join(branches)})" + + async def purge_stale_search_index_rows( session_maker: async_sessionmaker[AsyncSession], project_id: int, @@ -424,6 +477,7 @@ async def search( offset: int = 0, allow_relaxed: bool = False, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content. @@ -444,6 +498,11 @@ async def search( be true of the world. Sources without such a claim are excluded. limit: Maximum results to return offset: Number of results to skip + candidate_keys: Restrict results to these ``(type, id)`` search rows. ``None`` + searches the whole project; an empty sequence matches nothing. Honored by + the full-text pass, which is where vector and hybrid retrieval evaluate + their structured filters: that pass asks which of a known candidate set a + filter admits instead of paging the filter's whole match set (#1431). Returns: List of SearchIndexRow results with relevance scores @@ -2454,8 +2513,8 @@ def _log_vector_summary() -> None: ) if filter_requested: - filtered_rows = await self.search( - search_text=None, + allowed_keys = await self._filter_candidate_keys( + list(search_index_rows), permalink=permalink, permalink_match=permalink_match, title=title, @@ -2466,13 +2525,7 @@ def _log_vector_summary() -> None: metadata_filters=metadata_filters, file_path_prefix=file_path_prefix, temporal=temporal, - retrieval_mode=SearchRetrievalMode.FTS, - limit=VECTOR_FILTER_SCAN_LIMIT, - offset=0, ) - # Use (type, id) tuples to avoid collisions between different - # search_index row types that share the same auto-increment id. - allowed_keys = {(row.type, row.id) for row in filtered_rows if row.id is not None} if trace is not None: trace.vector = build_vector_stage( previous=trace.vector, @@ -2555,6 +2608,59 @@ def _log_vector_summary() -> None: _log_vector_summary() return output + async def _filter_candidate_keys( + self, + candidate_keys: Sequence[SearchIndexKey], + *, + permalink: Optional[str], + permalink_match: Optional[str], + title: Optional[str], + note_types: Optional[List[str]], + after_date: Optional[datetime], + search_item_types: Optional[List[SearchItemType]], + categories: Optional[List[str]], + metadata_filters: Optional[dict[str, Any]], + file_path_prefix: Optional[str], + temporal: Optional[TemporalFilter], + ) -> set[SearchIndexKey]: + """Return which of ``candidate_keys`` the structured filters admit. + + Vector retrieval scores embeddings and cannot evaluate a structured filter, so + the surviving candidates are decided by an FTS-mode pass carrying every filter. + Asking that pass for a *page of the filter's whole match set* and intersecting + client-side silently lost any candidate that sorted past the page (#1431); asking + it about the candidates themselves cannot, because the answer is bounded by the + question. + + The candidate list is split at the shared bind-parameter bound, so a deep page + whose candidate pool runs to thousands of rows costs a few small indexed lookups + instead of one unbounded scan. + """ + allowed_keys: set[SearchIndexKey] = set() + for batch_start in range(0, len(candidate_keys), VECTOR_HYDRATION_BATCH_SIZE): + batch = candidate_keys[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] + filtered_rows = await self.search( + search_text=None, + permalink=permalink, + permalink_match=permalink_match, + title=title, + note_types=note_types, + after_date=after_date, + search_item_types=search_item_types, + categories=categories, + metadata_filters=metadata_filters, + file_path_prefix=file_path_prefix, + temporal=temporal, + retrieval_mode=SearchRetrievalMode.FTS, + # The restriction, not this limit, is what bounds the result: one row per + # requested key, since (id, type, project_id) identifies a search row. + limit=len(batch), + offset=0, + candidate_keys=batch, + ) + allowed_keys.update((row.type, row.id) for row in filtered_rows if row.id is not None) + return allowed_keys + async def _fetch_search_index_rows_by_ids( self, row_ids: list[int] ) -> dict[SearchIndexKey, SearchIndexRow]: diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index d48a737d7..4160d7932 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -29,7 +29,9 @@ from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_query import relaxed_query_words from basic_memory.repository.search_repository_base import ( + SearchIndexKey, SearchRepositoryBase, + candidate_key_restriction_condition, file_path_prefix_condition, metadata_contains_like_condition, metadata_filter_content_type_condition, @@ -795,6 +797,7 @@ async def _build_fts_query_parts( metadata_filters: Optional[dict[str, Any]] = None, file_path_prefix: Optional[str] = None, temporal: Optional[TemporalFilter] = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, ) -> tuple[str, str, dict[str, Any], str, str]: """Build SQLite FTS FROM/WHERE params shared by search and count.""" conditions = [] @@ -888,6 +891,13 @@ async def _build_fts_query_parts( if subtree_condition is not None: conditions.append(subtree_condition) + # Handle an explicit candidate-row restriction. Built by the shared helper so + # both backends restrict by the identical rule; see + # candidate_key_restriction_condition for why the vector filter pass asks about + # its candidates rather than paging the filter's whole match set (#1431). + if candidate_keys is not None: + conditions.append(candidate_key_restriction_condition(candidate_keys, params)) + # Handle entity type filter (parameterized for defense-in-depth) if search_item_types: type_placeholders = [] @@ -1110,6 +1120,7 @@ async def search( allow_relaxed: bool = False, session: AsyncSession | None = None, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Search across all indexed content using SQLite FTS5. @@ -1160,6 +1171,7 @@ async def search( metadata_filters=metadata_filters, file_path_prefix=file_path_prefix, temporal=temporal, + candidate_keys=candidate_keys, ) # set limit on search query diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index 05aa2e7d7..2a270cb69 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -6,6 +6,7 @@ 3. Produces zero fused score when the source score is zero """ +from collections.abc import Sequence from dataclasses import dataclass from datetime import datetime from typing import override, Any, Optional, cast @@ -15,7 +16,11 @@ from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import FUSION_BONUS, SearchRepositoryBase +from basic_memory.repository.search_repository_base import ( + FUSION_BONUS, + SearchIndexKey, + SearchRepositoryBase, +) from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode from basic_memory.temporal import TemporalFilter @@ -90,6 +95,7 @@ async def search( offset: int = 0, allow_relaxed: bool = False, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] # pragma: no cover diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index e8154a4e8..3f2c369e8 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -15,6 +15,7 @@ from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_repository_base import ( + SearchIndexKey, SearchRepositoryBase, _PreparedEntityVectorSync, ) @@ -96,6 +97,7 @@ async def search( offset: int = 0, allow_relaxed: bool = False, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index 85c9ec721..307b53030 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -1,6 +1,7 @@ """Focused edge-case coverage for shared semantic vector synchronization.""" import hashlib +from collections.abc import Sequence from contextlib import asynccontextmanager from datetime import datetime from types import SimpleNamespace @@ -12,7 +13,10 @@ from basic_memory.repository import semantic_vector_sync from basic_memory.repository import search_repository_base as search_repository_base_module from basic_memory.repository.search_index_row import SearchIndexRow -from basic_memory.repository.search_repository_base import SearchRepositoryBase +from basic_memory.repository.search_repository_base import ( + SearchIndexKey, + SearchRepositoryBase, +) from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.semantic_chunking import VectorChunkRecord from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -65,6 +69,7 @@ async def search( offset: int = 0, allow_relaxed: bool = False, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] diff --git a/tests/repository/test_vector_filter_candidate_restriction.py b/tests/repository/test_vector_filter_candidate_restriction.py new file mode 100644 index 000000000..bb6549f6f --- /dev/null +++ b/tests/repository/test_vector_filter_candidate_restriction.py @@ -0,0 +1,318 @@ +"""A filtered vector search must not lose candidates to a scan window (#1431). + +Vector and hybrid retrieval do not evaluate structured filters themselves. They build a +candidate set from embeddings and then ask an FTS-mode pass which of those candidates the +filter admits. While that pass was a plain capped page over the filter's *whole* match +set, any candidate whose row fell outside the page was read as disallowed -- a wrong +answer that reruns identically, not a derived-state race that a later write repairs. + +The loss is not even arbitrary. A filter-only pass carries ``search_text=None``, so there +is no relevance signal in the ordering: PostgreSQL falls through to its +``search_index.id ASC`` tiebreak and keeps the earliest-indexed rows, and SQLite has no +tiebreak at all. Newer content is what disappears. + +The fix pushes the candidate keys into the filter query, so the pass answers "which of +*these* rows does the filter admit" instead of "here is a page of everything it admits". + +Every database test here runs against whichever backend the session is configured for -- +SQLite by default, PostgreSQL under BASIC_MEMORY_TEST_POSTGRES=1 -- through the shared +``search_repository`` fixture, because a restriction that admitted different rows per +dialect would hand semantic search a different candidate set on each. +""" + +from datetime import datetime, timezone +from typing import Any, cast +from unittest.mock import AsyncMock, patch + +import pytest +from sqlalchemy import text + +from basic_memory import db +from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.search_repository_base import ( + VECTOR_FILTER_SCAN_LIMIT, + VECTOR_HYDRATION_BATCH_SIZE, + candidate_key_restriction_condition, +) +from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode + +# The scope the admitted rows share, plus a sibling scope the filter must reject. +SCOPE = "notes" +REJECTED_SCOPE = "archive" +# One row past the old window: the smallest seed that makes the window bind at all. +FILLER_ROW_COUNT = VECTOR_FILTER_SCAN_LIMIT + 1 +# Indexed after every filler and numbered above them, so it is last under both the +# PostgreSQL id tiebreak and SQLite's insertion order -- i.e. genuinely outside the page. +TARGET_ROW_ID = 10_000_000 +TARGET_CONTENT = "the answer this query is looking for" +# Enough extra candidates to force the candidate list past the shared bind bound. +EXTRA_CANDIDATE_COUNT = VECTOR_HYDRATION_BATCH_SIZE + 50 +REJECTED_ROW_IDS = range(20_000_000, 20_000_005) + +_INSERT_SEARCH_ROW = """ + INSERT INTO search_index ( + id, title, content_stems, content_snippet, script_ngrams, permalink, + file_path, type, metadata, from_id, to_id, relation_type, + entity_id, category, created_at, updated_at, project_id + ) VALUES ( + :id, :title, :content_stems, :content_snippet, :script_ngrams, :permalink, + :file_path, :type, :metadata, :from_id, :to_id, :relation_type, + :entity_id, :category, :created_at, :updated_at, :project_id + ) +""" + +# Seeding 50k rows through one executemany holds a huge parameter list in the driver for +# no benefit; a few thousand at a time keeps both drivers comfortable. +_SEED_CHUNK_SIZE = 5000 + + +def _row_params( + project_id: int, + row_id: int, + name: str, + content: str, + *, + scope: str = SCOPE, + row_type: str = SearchItemType.ENTITY.value, +) -> dict[str, Any]: + now = datetime.now(timezone.utc) + return { + "id": row_id, + "title": name, + "content_stems": content, + "content_snippet": content, + "script_ngrams": "", + "permalink": f"{scope}/{row_type}/{name}", + "file_path": f"{scope}/{name}.md", + "type": row_type, + "metadata": None, + "from_id": None, + "to_id": None, + "relation_type": None, + "entity_id": row_id, + "category": None, + "created_at": now, + "updated_at": now, + "project_id": project_id, + } + + +async def _insert_rows(session_maker, rows: list[dict[str, Any]]) -> None: + async with db.scoped_session(session_maker) as session: + for start in range(0, len(rows), _SEED_CHUNK_SIZE): + await session.execute(text(_INSERT_SEARCH_ROW), rows[start : start + _SEED_CHUNK_SIZE]) + await session.commit() + + +@pytest.fixture +async def over_window_project(search_repository, session_maker) -> None: + """Seed a project whose filter match set genuinely exceeds the old scan window. + + Rows go straight into ``search_index`` rather than through entity indexing: the defect + lives in the search-row intersection, and 50k parsed notes would cost minutes to prove + the same thing. The target row is written last and numbered highest so it lands past + the old page under either backend's ordering. + """ + project_id = search_repository.project_id + filler = [ + _row_params(project_id, row_id, f"filler-{row_id}", "filler note body") + for row_id in range(FILLER_ROW_COUNT) + ] + rejected = [ + _row_params( + project_id, + row_id, + f"outside-{row_id}", + "out of scope body", + scope=REJECTED_SCOPE, + ) + for row_id in REJECTED_ROW_IDS + ] + target = [_row_params(project_id, TARGET_ROW_ID, "target", TARGET_CONTENT)] + await _insert_rows(session_maker, filler + rejected + target) + + +def _fake_embedding_provider() -> EmbeddingProvider: + return cast( + EmbeddingProvider, + type( + "EP", + (), + {"embed_query": AsyncMock(return_value=[0.0] * 384), "dimensions": 384}, + )(), + ) + + +def _semantic_repo(search_repository): + """Enable semantic retrieval on a repository the test config left keyword-only.""" + search_repository._semantic_enabled = True + search_repository._semantic_min_similarity = 0.0 + search_repository._embedding_provider = _fake_embedding_provider() + return search_repository + + +def _vector_chunks(row_ids: list[int]) -> list[dict[str, Any]]: + """One vector hit per row, ranked in the order given.""" + return [ + { + "chunk_key": f"{SearchItemType.ENTITY.value}:{row_id}:0", + "best_similarity": 0.99 - index * 0.001, + "chunk_text": TARGET_CONTENT, + "entity_id": row_id, + } + for index, row_id in enumerate(row_ids) + ] + + +@pytest.mark.slow +@pytest.mark.asyncio +async def test_filtered_vector_search_keeps_a_candidate_past_the_scan_window( + search_repository, over_window_project +): + """A highly similar row survives the filter even when it sorts past the old page.""" + repo = _semantic_repo(search_repository) + + with ( + patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), + patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), + patch.object( + repo, + "_run_vector_query", + new_callable=AsyncMock, + return_value=_vector_chunks([TARGET_ROW_ID]), + ), + ): + results = await repo.search( + search_text="the answer", + file_path_prefix=SCOPE, + retrieval_mode=SearchRetrievalMode.VECTOR, + limit=10, + ) + + assert [row.id for row in results] == [TARGET_ROW_ID] + + +@pytest.mark.slow +@pytest.mark.asyncio +async def test_filter_pass_answers_every_candidate_within_the_bind_bound( + search_repository, over_window_project +): + """A candidate pool past the bind bound is split, and each half still gets a verdict. + + Both engines cap bind parameters, so the restriction cannot be one unbounded ``IN`` + list. Splitting it must not turn into the same silent truncation it replaced: every + admitted candidate comes back, every out-of-scope one is rejected, and no single + statement carries more keys than the shared bound. + """ + repo = _semantic_repo(search_repository) + admitted = [TARGET_ROW_ID, *range(EXTRA_CANDIDATE_COUNT)] + candidates = [*admitted, *REJECTED_ROW_IDS] + assert len(candidates) > VECTOR_HYDRATION_BATCH_SIZE + + batched_key_counts: list[int] = [] + original_search = repo.search + + async def recording_search(*args, **kwargs): + if kwargs.get("candidate_keys") is not None: + batched_key_counts.append(len(kwargs["candidate_keys"])) + return await original_search(*args, **kwargs) + + with ( + patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), + patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), + patch.object( + repo, + "_run_vector_query", + new_callable=AsyncMock, + return_value=_vector_chunks(candidates), + ), + patch.object(repo, "search", recording_search), + ): + results = await repo._search_vector_only( + search_text="the answer", + permalink=None, + permalink_match=None, + title=None, + note_types=None, + after_date=None, + search_item_types=None, + categories=None, + metadata_filters=None, + file_path_prefix=SCOPE, + temporal=None, + limit=len(candidates), + offset=0, + ) + + assert {row.id for row in results} == set(admitted) + assert len(batched_key_counts) > 1 + assert max(batched_key_counts) <= VECTOR_HYDRATION_BATCH_SIZE + assert sum(batched_key_counts) == len(candidates) + + +@pytest.mark.asyncio +async def test_restriction_separates_rows_that_share_an_id(search_repository, session_maker): + """The restriction is by ``(type, id)``, since row types number independently.""" + project_id = search_repository.project_id + await _insert_rows( + session_maker, + [ + _row_params(project_id, 7, "shared-id-entity", "body"), + _row_params( + project_id, + 7, + "shared-id-observation", + "body", + row_type=SearchItemType.OBSERVATION.value, + ), + ], + ) + + rows = await search_repository.search( + limit=10, + candidate_keys=[(SearchItemType.OBSERVATION.value, 7)], + ) + + assert [(row.type, row.id) for row in rows] == [(SearchItemType.OBSERVATION.value, 7)] + + +@pytest.mark.asyncio +async def test_an_empty_candidate_set_admits_nothing(search_repository, session_maker): + """No candidates is a real state -- a search whose every hit was already dropped.""" + await _insert_rows( + session_maker, + [_row_params(search_repository.project_id, 7, "only-row", "body")], + ) + + assert await search_repository.search(limit=10, candidate_keys=[]) == [] + + +def test_restriction_groups_ids_under_one_predicate_per_row_type(): + """Type-scoped ``IN`` lists bind one parameter per key, not two.""" + params: dict[str, Any] = {} + condition = candidate_key_restriction_condition( + [("entity", 1), ("observation", 2), ("entity", 3), ("entity", 1)], + params, + ) + + assert condition == ( + "((search_index.type = :candidate_type_0 " + "AND search_index.id IN (:candidate_id_0_0, :candidate_id_0_1)) " + "OR (search_index.type = :candidate_type_1 " + "AND search_index.id IN (:candidate_id_1_0)))" + ) + assert params == { + "candidate_type_0": "entity", + "candidate_id_0_0": 1, + "candidate_id_0_1": 3, + "candidate_type_1": "observation", + "candidate_id_1_0": 2, + } + + +def test_restriction_of_no_keys_is_false_not_vacuously_true(): + """An empty ``OR`` would collapse to a predicate that admits the whole project.""" + params: dict[str, Any] = {} + + assert candidate_key_restriction_condition([], params) == "1 = 0" + assert params == {} diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index 8609324d3..9139f4859 100644 --- a/tests/repository/test_vector_pagination.py +++ b/tests/repository/test_vector_pagination.py @@ -4,6 +4,7 @@ which requires a sufficiently large candidate_limit multiplier. """ +from collections.abc import Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime @@ -12,7 +13,10 @@ import pytest -from basic_memory.repository.search_repository_base import SearchRepositoryBase +from basic_memory.repository.search_repository_base import ( + SearchIndexKey, + SearchRepositoryBase, +) from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode @@ -75,6 +79,7 @@ async def search( offset: int = 0, allow_relaxed: bool = False, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] # pragma: no cover diff --git a/tests/repository/test_vector_threshold.py b/tests/repository/test_vector_threshold.py index 711f09c22..1c6323aac 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -1,5 +1,6 @@ """Tests for semantic_min_similarity threshold filtering in vector search.""" +from collections.abc import Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime @@ -13,6 +14,7 @@ from basic_memory.repository.search_repository_base import ( SMALL_NOTE_CONTENT_LIMIT, TOP_CHUNKS_PER_RESULT, + SearchIndexKey, SearchRepositoryBase, ) from basic_memory.repository.search_trace import SearchTraceCollector @@ -79,6 +81,7 @@ async def search( offset: int = 0, allow_relaxed: bool = False, *, + candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: return [] # pragma: no cover