From fbaf5b3d52959a9d61ebb1003e437b549d693ec0 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 12:57:24 -0500 Subject: [PATCH] refactor(core): extract SearchReader and SemanticSearch from the search repository Retrieval leaves SearchRepositoryBase. SearchReader runs one prepared query over one ProjectScope in whichever mode it asks for; SemanticSearch owns vector and hybrid retrieval (adapter lookup, manifest hydration, the structured filter pass, score fusion, reranking, pagination) over a VectorRetrieval that is present or absent instead of probed with hasattr. The repository keeps what only it knows, whether semantic search is enabled and its vector tables exist, and builds a reader per call from its current state. Hydrated chunks are a typed HydratedChunk, which retires the best_distance compatibility branch and the per-backend _distance_to_similarity hooks the adapters had already replaced. PreparedSearchQuery gains has_filters, the predicate the vector path uses to decide whether to run the filter pass. Test doubles construct the pipeline directly (a FakeFts backend and a VectorRetrieval over a stubbed adapter) instead of subclassing the repository and patching its private hooks. No query behavior changes. Part of #1558. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019YW9ysxugGGBCNEGzsxtFV Signed-off-by: phernandez --- CHANGELOG.md | 12 + .../repository/postgres_search_repository.py | 9 - src/basic_memory/repository/search_query.py | 22 + src/basic_memory/repository/search_reader.py | 1136 ++++++++++++++ .../repository/search_repository_base.py | 1330 +---------------- .../repository/sqlite_search_repository.py | 9 - .../services/project_readiness.py | 2 +- .../services/retrieval_inspect.py | 9 +- .../test_multilingual_benchmark_contract.py | 2 +- test-int/semantic/test_search_diagnostics.py | 15 +- test-int/semantic/test_semantic_coverage.py | 18 +- .../repository/test_distance_to_similarity.py | 24 - tests/repository/test_hybrid_fusion.py | 314 ++-- .../test_postgres_search_repository_unit.py | 2 +- tests/repository/test_rerank_pipeline.py | 253 ++-- .../test_search_file_path_prefix.py | 23 +- tests/repository/test_search_reader.py | 283 ++++ tests/repository/test_search_trace.py | 4 +- tests/repository/test_semantic_search_base.py | 63 +- tests/repository/test_semantic_vector_sync.py | 15 - ...est_vector_filter_candidate_restriction.py | 61 +- tests/repository/test_vector_pagination.py | 208 +-- .../repository/test_vector_temporal_filter.py | 132 +- tests/repository/test_vector_threshold.py | 512 ++----- tests/services/test_project_readiness.py | 2 +- 25 files changed, 2056 insertions(+), 2404 deletions(-) create mode 100644 src/basic_memory/repository/search_reader.py delete mode 100644 tests/repository/test_distance_to_similarity.py create mode 100644 tests/repository/test_search_reader.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a523a2777..8b57ec9a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -131,6 +131,18 @@ layer with defaults, and the filter helpers both backends share move from the base into `search_filters`. No query behavior changes. +- **#1558**: Retrieval leaves `SearchRepositoryBase`. `SearchReader` runs one prepared + query over one `ProjectScope` in whichever mode it asks for, and `SemanticSearch` owns + vector and hybrid retrieval (adapter lookup, manifest hydration, the structured filter + pass, score fusion, reranking, pagination) over a `VectorRetrieval` that is present or + absent rather than probed with `hasattr`. The repository keeps what only it knows + (whether semantic search is enabled and its vector tables exist) and builds a reader + per call from its current state. Hydrated chunks are a typed `HydratedChunk`, which + retires the `best_distance` compatibility branch and the per-backend + `_distance_to_similarity` hooks the adapters had already replaced. Test doubles + construct the pipeline directly instead of subclassing the repository. No query + behavior changes. + ## v0.23.2 (2026-08-25) diff --git a/src/basic_memory/repository/postgres_search_repository.py b/src/basic_memory/repository/postgres_search_repository.py index 0cb17fcb0..06f4e57d8 100644 --- a/src/basic_memory/repository/postgres_search_repository.py +++ b/src/basic_memory/repository/postgres_search_repository.py @@ -490,15 +490,6 @@ async def _delete_stale_chunks( expected_deletions=expected_deletions, ) - @override - def _distance_to_similarity(self, distance: float) -> float: - """Convert pgvector cosine distance to cosine similarity. - - pgvector's <=> operator returns cosine distance in [0, 2], - where cos_distance = 1 - cos_similarity. - """ - return max(0.0, 1.0 - distance) - @override def _timestamp_now_expr(self) -> str: return "NOW()" diff --git a/src/basic_memory/repository/search_query.py b/src/basic_memory/repository/search_query.py index 9cd2cb2b4..59dd16182 100644 --- a/src/basic_memory/repository/search_query.py +++ b/src/basic_memory/repository/search_query.py @@ -32,6 +32,28 @@ class PreparedSearchQuery: retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS min_similarity: float | None = None + @property + def has_filters(self) -> bool: + """Whether any predicate beyond the text itself narrows the result set. + + Vector retrieval cannot evaluate these itself; when any is present it asks + the full-text pass which of its candidates the filters admit. + """ + return any( + ( + self.permalink, + self.permalink_match, + self.title, + self.note_types, + self.after_date, + self.search_item_types, + self.categories, + self.metadata_filters, + self.file_path_prefix, + self.temporal, + ) + ) + # Interrogative/function words contribute lexical noise when a strict # full-text query is relaxed: "when OR did OR a" matches loud wrong documents diff --git a/src/basic_memory/repository/search_reader.py b/src/basic_memory/repository/search_reader.py new file mode 100644 index 000000000..d5f925d54 --- /dev/null +++ b/src/basic_memory/repository/search_reader.py @@ -0,0 +1,1136 @@ +"""Scoped search retrieval shared by every route. + +``SearchReader`` runs one prepared query over one ``ProjectScope``: the engine's +full-text statement for FTS mode, and vector or hybrid retrieval through +``SemanticSearch`` when the semantic stack is present. Neither class owns indexing, +manifest writes, or table lifecycle. A repository builds a reader per call from its +current state, so a repository whose semantic stack was disabled at startup hands the +reader the matching capability set; a route reading several projects at once builds +one directly over a wider scope. +""" + +import time +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Any, assert_never + +import logfire +from loguru import logger +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from basic_memory import db +from basic_memory.repository.embedding_provider import EmbeddingProvider +from basic_memory.repository.rerank_provider import ( + RerankProvider, + build_rerank_document, + demote_tail_scores, + validate_rerank_scores, +) +from basic_memory.repository.search_filters import FtsBackend +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.search_trace import ( + BelowThreshold, + FilteredOut, + HydrationDropKey, + HydrationDropped, + MissingSearchRow, + SearchTraceCollector, + build_fts_page_stage, + build_fusion_stage, + build_rerank_stage, + build_vector_stage, + classify_hydration_drops, + read_manifest_readiness, +) +from basic_memory.repository.semantic_errors import SemanticSearchDisabledError +from basic_memory.repository.semantic_vector_index import ( + SemanticVectorIndex, + VectorKey, + VectorMatch, +) +from basic_memory.schemas.search import SearchRetrievalMode + +# --- Retrieval constants --- + +# Adapters that share the authoritative SQL database. Everything else stores vectors +# outside it and needs the stale-hit overfetch below. +BUILT_IN_VECTOR_INDEX_NAMES = frozenset({"pgvector", "sqlite-vec"}) +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 +# to keep enough unique documents in the rerank window. +RERANK_POOL_CHUNK_FANOUT = 4 +FUSION_BONUS = 0.3 +FUSION_FORMULA_VERSION = "max+0.3*min/v1" +FTS_GATE_THRESHOLD = 0.0 +TOP_CHUNKS_PER_RESULT = 5 +SMALL_NOTE_CONTENT_LIMIT = 2000 + + +# The manifest conditions under which semantic retrieval will use a stored vector. +# Vector hydration admits exactly these rows, so anything failing them is invisible +# to search: a chunk left behind by an embedding-model or vector-index change, or one +# still pending. Readiness reporting must apply the same predicate — calling such a +# row "embedded" would report an index settled that retrieval cannot answer from, +# which is the class of lie #1414 exists to remove. +# Callers bind :vector_index and :embedding_model; the scope binds its own IDs. +def current_vector_manifest_predicate(scope: ProjectScope, params: dict[str, Any]) -> str: + """SQL admitting only manifest rows retrieval can answer from, within ``scope``.""" + return ( + f"{scope.predicate('project_id', params)} " + "AND vector_index = :vector_index " + "AND embedding_model = :embedding_model " + "AND embedding_status = 'ready'" + ) + + +def parse_chunk_key(chunk_key: str) -> SearchIndexKey: + """Parse a chunk key like ``observation:5:0`` into ``(type, search_index_id)``.""" + parts = chunk_key.split(":") + return parts[0], int(parts[1]) + + +def vector_eligible(query: PreparedSearchQuery) -> bool: + """Whether a query carries text to embed and no identity filter. + + Vector and hybrid retrieval score an embedding of the query text; a permalink or + title lookup has nothing to embed and answers exactly through the full-text path. + """ + text_value = (query.search_text or "").strip() + return ( + bool(text_value) + and text_value != "*" + and not query.permalink + and not query.permalink_match + and not query.title + ) + + +def rerank_document_text(row: SearchIndexRow, max_chars: int) -> str: + """Build the document text handed to the cross-encoder for one candidate. + + Prefer the matched chunk (the most relevant passage of a large note), falling + back to the stored snippet. + """ + body = row.matched_chunk_text or row.content_snippet or "" + return build_rerank_document(row.title, body, max_chars) + + +def demote_tail(tail: list[SearchIndexRow], floor: float) -> list[SearchIndexRow]: + """Rescore un-reranked tail rows at or below the floor, preserving their order. + + The reranked pool carries [0, 1] relevance scores while the tail still holds raw + retrieval scores on a different scale ([0, 1.3] for fused hybrid). Left as is, a + tail row could outrank a reranked row numerically. Positive floors put the tail + strictly below the pool; a zero floor yields zeroes because no smaller score + exists in the public [0, 1] range. The returned pool-plus-tail sequence, rather + than a later score-only sort, owns that tie-breaking invariant. + """ + return [ + replace(row, score=score) for row, score in zip(tail, demote_tail_scores(floor, len(tail))) + ] + + +# --- Retrieval capabilities --- + + +@dataclass(frozen=True, slots=True) +class VectorRetrieval: + """The live semantic stack vector retrieval reads from. + + Present only when semantic search is enabled, a provider is configured, and the + adapter is bound; absence means the reader answers full-text queries only. + """ + + index: SemanticVectorIndex + index_name: str + embedding_provider: EmbeddingProvider + # The persisted embedding identity manifest rows are keyed by. + embedding_model: str + vector_k: int + min_similarity: float + + @property + def external(self) -> bool: + """Whether vectors live outside the SQL database that owns the manifest.""" + return self.index_name not in BUILT_IN_VECTOR_INDEX_NAMES + + +@dataclass(frozen=True, slots=True) +class Reranking: + """A configured cross-encoder and the fixed prefix it rescores.""" + + provider: RerankProvider + candidates: int + max_document_chars: int + + +@dataclass(frozen=True, slots=True) +class HydratedChunk: + """One adapter match the ready manifest confirmed retrieval may serve.""" + + entity_id: int + chunk_key: str + chunk_text: str + similarity: float + + +# --- Vector and hybrid retrieval --- + + +class SemanticSearch: + """Vector and hybrid retrieval over one scope. + + Constructed only when the semantic stack is present, so every stage reads its + adapter, provider, and thresholds directly instead of re-checking availability. + """ + + def __init__( + self, + session_maker: async_sessionmaker[AsyncSession], + scope: ProjectScope, + fts: FtsBackend, + vector: VectorRetrieval, + rerank: Reranking | None = None, + ) -> None: + self.session_maker = session_maker + self.scope = scope + self.fts = fts + self.vector = vector + self.rerank = rerank + + # --- Candidate window sizing --- + + def _active_rerank(self, query_text: str) -> Reranking | None: + """The reranker this query runs: one is configured and there is text to score.""" + return self.rerank if query_text else None + + def _rerank_candidate_limit(self, rerank: Reranking) -> int: + """Return the fixed chunk window that owns reranker-prefix membership.""" + return max(self.vector.vector_k, rerank.candidates * RERANK_POOL_CHUNK_FANOUT) + + def _candidate_limit(self, limit: int, offset: int, query_text: str) -> int: + """Size the retrieval candidate *chunk* pool for vector/hybrid search. + + ``candidate_limit`` bounds vector chunks, but many chunks of one large note + collapse to a single ``(type, id)`` row before reranking, so a chunk count does + not equal a unique-document count. When reranking is active we over-fetch by + ``RERANK_POOL_CHUNK_FANOUT`` so a few multi-chunk notes can't starve the rerank + window below ``reranker_candidates`` unique rows. This is best-effort headroom, + not a hard guarantee — a single note dominating the entire nearest-neighbour set + can still yield fewer unique rows (a pathological corpus shape). + """ + rerank = self._active_rerank(query_text) + if rerank is None: + return max(self.vector.vector_k, (limit + offset) * 10) + # Trigger: the requested window extends beyond the fixed reranked prefix. + # Why: a bounded prefix alone can under-fill large pages and hide the + # semantic pagination probe even when more matches exist. + # Outcome: keep prefix membership fixed while adding chunk headroom only + # for the untouched tail that this request must return. + tail_size = max(0, limit + offset - rerank.candidates) + return self._rerank_candidate_limit(rerank) + tail_size * 10 + + # --- Vector nearest neighbours through the ready manifest --- + + @logfire.instrument("search.vector_query", extract_args=False) + async def _run_vector_query( + self, + session: AsyncSession, + query_embedding: list[float], + candidate_limit: int, + *, + trace: SearchTraceCollector | None = None, + ) -> list[HydratedChunk]: + """Query the configured adapter and hydrate only live, ready manifest rows.""" + if trace is not None: + trace.vector = build_vector_stage( + candidate_limit=candidate_limit, + adapter_match_count=0, + hydrated_count=0, + ) + if candidate_limit <= 0: + return [] + + if not self.vector.external: + matches = await self.vector.index.search(query_embedding, limit=candidate_limit) + if trace is not None: + trace.readiness = await read_manifest_readiness( + session, + self.scope, + self.vector.index_name, + self.vector.embedding_model, + ) + return await self._hydrate_vector_matches(session, matches, trace=trace) + + scan_limit = min(candidate_limit, VECTOR_FILTER_SCAN_LIMIT) + while True: + matches = await self.vector.index.search(query_embedding, limit=scan_limit) + if trace is not None and trace.readiness is None: + trace.readiness = await read_manifest_readiness( + session, + self.scope, + self.vector.index_name, + self.vector.embedding_model, + ) + hydrated = await self._hydrate_vector_matches(session, matches, trace=trace) + if ( + len(hydrated) >= candidate_limit + or len(matches) < scan_limit + or scan_limit >= VECTOR_FILTER_SCAN_LIMIT + ): + returned = hydrated[:candidate_limit] + # Trigger: the expanded stale-hit rescan hydrated more chunks than the + # candidate window the search consumes. + # Why: chunks beyond the window never enter thresholding, fusion, or + # reranking — tracing them would invent candidates this execution + # never considered. + # Outcome: the traced stage is trimmed to the returned window. + if trace is not None and trace.vector is not None and len(hydrated) > len(returned): + # Two owners can share one parseable chunk_key (manifest uniqueness + # includes entity_id), so window membership matches by owner too. + returned_chunk_keys = {(chunk.entity_id, chunk.chunk_key) for chunk in returned} + trimmed: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} + for chunk_match in trace.vector.chunk_matches: + if (chunk_match.entity_id, chunk_match.chunk_key) in returned_chunk_keys: + trimmed.setdefault(chunk_match.key, []).append( + ( + chunk_match.chunk_key, + chunk_match.similarity, + chunk_match.entity_id, + ) + ) + # hydrated_count keeps full-scan scope so the vector stage's + # dropped count matches its hydration-drop list; the flattener + # reports the window truncation as its own candidate_window stage. + trace.vector = build_vector_stage( + previous=trace.vector, + chunk_matches=trimmed, + ) + return returned + + # Trigger: stale, pending, or wrong-model adapter hits consumed the + # requested top-k before manifest hydration. + # Why: returning early lets stale extension data crowd every live + # result out of an otherwise valid semantic search. + # Outcome: retry from the same ranked prefix with bounded geometric + # overfetch until enough live rows survive or the adapter is exhausted. + scan_limit = min(scan_limit * 2, VECTOR_FILTER_SCAN_LIMIT) + + @logfire.instrument("search.vector_manifest_hydration", extract_args=False) + async def _hydrate_vector_matches( + self, + session: AsyncSession, + matches: list[VectorMatch], + *, + trace: SearchTraceCollector | None = None, + ) -> list[HydratedChunk]: + """Resolve adapter matches through the authoritative ready manifest.""" + if not matches: + return [] + + chunks_by_key: dict[VectorKey, str] = {} + for batch_start in range(0, len(matches), VECTOR_HYDRATION_BATCH_SIZE): + batch = matches[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] + params: dict[str, Any] = { + "vector_index": self.vector.index_name, + "embedding_model": self.vector.embedding_model, + } + manifest_predicate = current_vector_manifest_predicate(self.scope, params) + predicates: list[str] = [] + for index, match in enumerate(batch): + params[f"entity_id_{index}"] = match.key.entity_id + params[f"chunk_key_{index}"] = match.key.chunk_key + predicates.append( + f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" + ) + + # Constraint: adapters may return thousands of candidates for deep pages. + # PostgreSQL and SQLite both cap bind parameters, so hydrate in fixed-size + # batches while retaining the adapter's original ranking in the final list. + result = await session.execute( + text( + "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " + "WHERE " + manifest_predicate + " " + "AND (" + " OR ".join(predicates) + ")" + ), + params, + ) + chunks_by_key.update( + { + VectorKey( + entity_id=int(row["entity_id"]), + chunk_key=str(row["chunk_key"]), + ): str(row["chunk_text"]) + for row in result.mappings().all() + } + ) + hydrated = [ + HydratedChunk( + entity_id=match.key.entity_id, + chunk_key=match.key.chunk_key, + chunk_text=chunks_by_key[match.key], + similarity=match.similarity, + ) + for match in matches + if match.key in chunks_by_key + ] + if trace is not None: + dropped_keys = [ + HydrationDropKey( + entity_id=match.key.entity_id, + chunk_key=match.key.chunk_key, + similarity=match.similarity, + configured_index=self.vector.index_name, + configured_model=self.vector.embedding_model, + ) + for match in matches + if match.key not in chunks_by_key + ] + drops = await classify_hydration_drops(session, self.scope, dropped_keys) + chunk_matches: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} + malformed_drops: list[HydrationDropped] = [] + for chunk in hydrated: + try: + key = parse_chunk_key(chunk.chunk_key) + except (ValueError, IndexError): + # A hydrated chunk with an unparseable key silently vanishes from + # retrieval; the trace must name it or the stage counts lie. + malformed_drops.append( + HydrationDropped( + entity_id=chunk.entity_id, + chunk_key=chunk.chunk_key, + similarity=chunk.similarity, + reason="malformed_key", + stored_model=None, + stored_index=None, + ) + ) + continue + chunk_matches.setdefault(key, []).append( + (chunk.chunk_key, chunk.similarity, chunk.entity_id) + ) + trace.vector = build_vector_stage( + previous=trace.vector, + adapter_match_count=len(matches), + # Malformed keys are dropped, not served — counting them as output + # would contradict the malformed_key rejection listed alongside. + hydrated_count=len(hydrated) - len(malformed_drops), + drops=(*drops, *malformed_drops), + chunk_matches=chunk_matches, + ) + return hydrated + + # --- Candidate rows and structured filters --- + + @logfire.instrument("search.fetch_candidate_rows", extract_args=False) + async def _fetch_search_index_rows_by_ids( + self, row_ids: list[int] + ) -> dict[SearchIndexKey, SearchIndexRow]: + """Fetch search_index rows by id, keyed by (type, id) to disambiguate types. + + A bare id can match one row per type (independent id sequences), so the + result must carry every matching row rather than letting one clobber another. + """ + if not row_ids: + return {} + placeholders = ",".join(f":id_{idx}" for idx in range(len(row_ids))) + params: dict[str, Any] = {f"id_{idx}": rid for idx, rid in enumerate(row_ids)} + scope_predicate = self.scope.predicate("project_id", params) + sql = f""" + SELECT + project_id, id, title, permalink, file_path, type, metadata, + from_id, to_id, relation_type, entity_id, content_snippet, + category, created_at, updated_at, 0 as score + FROM search_index + WHERE {scope_predicate} + AND id IN ({placeholders}) + """ + result: dict[SearchIndexKey, SearchIndexRow] = {} + async with db.scoped_session(self.session_maker) as session: + row_result = await session.execute(text(sql), params) + for row in row_result.fetchall(): + search_row = SearchIndexRow.from_mapping(row._asdict()) + result[(search_row.type, search_row.id)] = search_row + return result + + @logfire.instrument("search.filter_candidates", extract_args=False) + async def _filter_candidate_keys( + self, + candidate_keys: Sequence[SearchIndexKey], + query: PreparedSearchQuery, + ) -> set[SearchIndexKey]: + """Return which of ``candidate_keys`` the query's 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. + """ + filter_query = replace(query, search_text=None, retrieval_mode=SearchRetrievalMode.FTS) + 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.fts.search( + self.scope, + filter_query, + # 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 + + # --- Reranking --- + + async def _rerank_and_paginate( + self, + query_text: str, + rows: list[SearchIndexRow], + *, + offset: int, + limit: int, + stable_rows: list[SearchIndexRow] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Rerank the top candidates, then return the requested ``[offset:offset+limit]`` page. + + Trigger: a reranker is configured and there is a real query. + Why: bi-encoder/FTS ranking lands the gold document in the top-N but often + just below the top-k cutoff (#950); a cross-encoder that reads query and + document together recovers those near-misses. + Outcome: the first ``reranker_candidates`` rows are reordered by reranker + relevance (which replaces ``score``); the requested page is sliced from the + reordered list. + + Every non-empty page rescores the same fixed prefix before slicing so the + untouched tail can be demoted onto the reranker's public ``[0, 1]`` scale. + """ + page_end = offset + limit + rerank = self._active_rerank(query_text) + if rerank is None: + return rows[offset:page_end] + + # Trigger: pagination needs more rows than the fixed rerank retrieval window. + # Why: an expanded retrieval may introduce or strengthen raw candidates, but + # letting them replace the original prefix causes duplicates and skips. + # Outcome: the fixed window owns prefix membership; the expanded result only + # supplies new, de-duplicated tail rows. + pool_source = stable_rows if stable_rows is not None else rows + pool = pool_source[: rerank.candidates] + pool_keys = {(row.type, row.id) for row in pool} + tail = [row for row in rows if (row.type, row.id) not in pool_keys] + ordered_rows = pool + tail + + # Skip only when there is no prefix to calibrate or the requested page is + # empty. Even a singleton prefix or a wholly-tail page needs the prefix's + # relevance floor so raw hybrid scores cannot leak into cross-project sorting. + if not pool or offset >= len(ordered_rows): + return ordered_rows[offset:page_end] + + pre_rerank_scores = None + if trace is not None: + pre_rerank_scores = {(row.type, row.id): row.score or 0.0 for row in ordered_rows} + documents = [rerank_document_text(row, rerank.max_document_chars) for row in pool] + # A transient provider failure must surface instead of switching this page + # back to retrieval order. A prior page may already have returned reranked + # order, so degrading here can duplicate one result and omit another. + rerank_start = time.perf_counter() if trace is not None else None + with logfire.span( + "search.rerank", + candidate_count=len(pool), + document_chars=sum(map(len, documents)), + ): + scores = validate_rerank_scores( + await rerank.provider.rerank(query_text, documents), + len(pool), + ) + + order = sorted(range(len(pool)), key=lambda i: scores[i], reverse=True) + reranked = [replace(pool[i], score=scores[i]) for i in order] + logger.debug( + "Reranked candidates: pool={pool} model={model}", + pool=len(pool), + model=rerank.provider.model_name, + ) + tail_floor = reranked[-1].score or 0.0 + demoted_tail = demote_tail(tail, floor=tail_floor) + reranked_rows = reranked + demoted_tail + if trace is not None: + assert pre_rerank_scores is not None and rerank_start is not None + trace.rerank = build_rerank_stage( + provider_model=rerank.provider.model_name, + reranker_candidates=rerank.candidates, + pre_rerank_scores=pre_rerank_scores, + pool_keys=[(row.type, row.id) for row in pool], + rerank_scores={ + (pool[index].type, pool[index].id): score for index, score in enumerate(scores) + }, + post_rerank_rows=[((row.type, row.id), row.score or 0.0) for row in reranked_rows], + demoted_scores={(row.type, row.id): row.score or 0.0 for row in demoted_tail}, + tail_floor=tail_floor, + stable_pool_refetched=trace.stable_pool_refetched, + rerank_ms=(time.perf_counter() - rerank_start) * 1000, + ) + return reranked_rows[offset:page_end] + + # --- Vector-only retrieval --- + + async def vector_only( + self, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + candidate_limit: int | None = None, + apply_rerank: bool = True, + emit_observability_log: bool = True, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Run vector-only search returning chunk-level results. + + Returns individual search_index rows (entities, observations, relations) + ranked by vector similarity. Each observation or relation is a first-class + result, not collapsed into its parent entity. + + ``candidate_limit`` is supplied only by a composed retrieval stage that + already sized the shared candidate pool. + """ + query_text = (query.search_text or "").strip() + if candidate_limit is None: + candidate_limit = self._candidate_limit(limit, offset, query_text) + query_start = time.perf_counter() + embed_start = time.perf_counter() + with logfire.span("search.embed_query", query_chars=len(query_text)): + query_embedding = await self.vector.embedding_provider.embed_query(query_text) + embed_ms = (time.perf_counter() - embed_start) * 1000 + vector_query_start = time.perf_counter() + + # Constraint: vector adapters may open their own session, while the SQLite + # test/runtime pool can contain only one connection. A plain AsyncSession + # defers checkout until hydration runs after adapter search has released it. + async with self.session_maker() as session: + vector_rows = await self._run_vector_query( + session, + query_embedding, + candidate_limit, + trace=trace, + ) + vector_query_ms = (time.perf_counter() - vector_query_start) * 1000 + vector_row_count = len(vector_rows) + hydrate_ms = 0.0 + + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + effective_min_similarity=( + query.min_similarity + if query.min_similarity is not None + else self.vector.min_similarity + ), + min_similarity_source=("query" if query.min_similarity is not None else "config"), + embed_ms=embed_ms, + vector_query_ms=vector_query_ms, + ) + + def _log_vector_summary() -> None: + if not emit_observability_log: + return + + total_ms = (time.perf_counter() - query_start) * 1000 + if total_ms > 2000: + logger.warning( + "[SEMANTIC_SLOW_QUERY] Semantic query timing: scope={scope} " + "retrieval_mode={retrieval_mode} query_length={query_length} " + "candidate_limit={candidate_limit} vector_row_count={vector_row_count} " + "embed_ms={embed_ms:.2f} vector_query_ms={vector_query_ms:.2f} " + "hydrate_ms={hydrate_ms:.2f} total_ms={total_ms:.2f}", + scope=self.scope.project_ids, + retrieval_mode="vector", + query_length=len(query_text), + candidate_limit=candidate_limit, + vector_row_count=vector_row_count, + embed_ms=embed_ms, + vector_query_ms=vector_query_ms, + hydrate_ms=hydrate_ms, + total_ms=total_ms, + ) + + if not vector_rows: + _log_vector_summary() + return [] + + hydrate_start = time.perf_counter() + # Build per-search_index_row similarity scores from chunk-level results. + # Each chunk_key encodes the search_index row type and id; keep both as the + # key because different row types can share the same numeric id (#982). + # Track the best similarity per row (for ranking) and all chunks (for context). + similarity_by_si_key: dict[SearchIndexKey, float] = {} + chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} + for chunk in vector_rows: + try: + si_key = parse_chunk_key(chunk.chunk_key) + except (ValueError, IndexError): + # A chunk without a parseable key names no search row to rank. + continue + current = similarity_by_si_key.get(si_key) + if current is None or chunk.similarity > current: + similarity_by_si_key[si_key] = chunk.similarity + chunks_by_si_key.setdefault(si_key, []).append((chunk.similarity, chunk.chunk_text)) + + if not similarity_by_si_key: + hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + _log_vector_summary() + return [] + + # Filter out results below the minimum similarity threshold. + # Per-query min_similarity overrides the configured default. + effective_min_similarity = ( + query.min_similarity if query.min_similarity is not None else self.vector.min_similarity + ) + if effective_min_similarity > 0.0: + if trace is not None: + threshold_rejections = tuple( + BelowThreshold(key=key, similarity=value, threshold=effective_min_similarity) + for key, value in similarity_by_si_key.items() + if value < effective_min_similarity + ) + trace.vector = build_vector_stage( + previous=trace.vector, + threshold_rejections=threshold_rejections, + ) + similarity_by_si_key = { + k: v for k, v in similarity_by_si_key.items() if v >= effective_min_similarity + } + if not similarity_by_si_key: + hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + _log_vector_summary() + return [] + + # Fetch the actual search_index rows. Colliding (type, id) keys share one + # bare id, so deduplicate while preserving first-seen order. + si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_si_key)) + search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids) + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + missing_search_rows=tuple( + MissingSearchRow(key=key) + for key in similarity_by_si_key + if key not in search_index_rows + ), + ) + + if query.has_filters: + allowed_keys = await self._filter_candidate_keys(list(search_index_rows), query) + if trace is not None: + trace.vector = build_vector_stage( + previous=trace.vector, + filter_rejections=tuple( + FilteredOut(key=key) for key in search_index_rows if key not in allowed_keys + ), + ) + search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys} + + ranked_rows: list[SearchIndexRow] = [] + for si_key, similarity in similarity_by_si_key.items(): + row = search_index_rows.get(si_key) + if row is None: + continue + + # Small notes: return full content so the answer is always present. + # Large notes: return top-N most relevant chunks for richer context. + content_snippet = row.content_snippet or "" + if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT: + matched_chunk_text = content_snippet + else: + si_chunks = chunks_by_si_key.get(si_key, []) + si_chunks.sort(key=lambda c: c[0], reverse=True) + top_texts = [chunk_text for _, chunk_text in si_chunks[:TOP_CHUNKS_PER_RESULT]] + matched_chunk_text = "\n---\n".join(top_texts) if top_texts else None + + ranked_rows.append( + replace( + row, + score=similarity, + matched_chunk_text=matched_chunk_text, + ) + ) + + ranked_rows.sort(key=lambda item: item.score or 0.0, reverse=True) + hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 + # Rerank over the wide candidate pool, then slice to the page. Suppressed when + # hybrid calls this internally (apply_rerank=False): hybrid reranks its own + # fused result, and _rerank_and_paginate is a plain slice without a reranker. + if apply_rerank: + stable_rows = ranked_rows + rerank = self._active_rerank(query_text) + if rerank is not None: + stable_candidate_limit = self._rerank_candidate_limit(rerank) + if candidate_limit > stable_candidate_limit: + if trace is not None: + trace.stable_pool_refetched = True + stable_rows = await self.vector_only( + query, + limit=stable_candidate_limit, + offset=0, + candidate_limit=stable_candidate_limit, + apply_rerank=False, + emit_observability_log=False, + trace=None, + ) + output = await self._rerank_and_paginate( + query_text, + ranked_rows, + offset=offset, + limit=limit, + stable_rows=stable_rows, + trace=trace, + ) + else: + output = ranked_rows[offset : offset + limit] + # Vector latency owns the optional rerank stage too. Logging before the + # awaited provider call hides the feature's dominant cost and can suppress + # the slow-query warning entirely. + _log_vector_summary() + return output + + # --- Hybrid score-based fusion --- + + async def hybrid( + self, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + candidate_limit: int | None = None, + apply_rerank: bool = True, + emit_observability_log: bool = True, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Fuse FTS and vector results using score-based fusion. + + Uses the search_index (type, id) pair as the fusion key. The formula + ``max(vec, fts) + FUSION_BONUS * min(vec, fts)`` preserves + the dominant signal and rewards dual-source agreement. + """ + query_text = (query.search_text or "").strip() + rerank = self._active_rerank(query_text) + query_start = time.perf_counter() + if candidate_limit is None: + candidate_limit = self._candidate_limit(limit, offset, query_text) + fts_start = time.perf_counter() + # allow_relaxed: question-form queries rarely AND-match, and a dead FTS + # branch silently degrades hybrid to vector-only ranking. Fusion plus + # bm25 keep relaxed lexical candidates from dominating precision. + with logfire.span("search.fts", candidate_limit=candidate_limit) as fts_span: + fts_results = await self.fts.search( + self.scope, + replace(query, retrieval_mode=SearchRetrievalMode.FTS), + limit=candidate_limit, + offset=0, + allow_relaxed=True, + trace=trace, + ) + fts_span.set_attribute("result_count", len(fts_results)) + fts_ms = (time.perf_counter() - fts_start) * 1000 + vector_start = time.perf_counter() + vector_results = await self.vector_only( + query, + limit=candidate_limit, + offset=0, + # Trigger: reranking owns a bounded candidate window shared by both legs. + # Why: the disabled path historically expands the vector leg again to + # preserve recall when many vector chunks collapse into a few search rows. + # Outcome: avoid double expansion only when reranking is actually active. + candidate_limit=candidate_limit if rerank is not None else None, + apply_rerank=False, + emit_observability_log=False, + trace=trace, + ) + vector_ms = (time.perf_counter() - vector_start) * 1000 + # Trigger: with reranking disabled the vector leg expands internally and can + # hydrate more rows than the fusion window it returns. + # Why: rows cut here never fuse — left in the trace they would surface as + # candidates with no rejection and no fused rank, which the response labels + # "returned". Rows with a recorded rejection keep their chunk evidence. + # Outcome: the trace keeps rows handed to fusion (or explicitly rejected); + # the cut shows up as served-chunk shrinkage in the candidate_window stage. + if trace is not None and trace.vector is not None: + kept_row_keys = {(row.type, row.id) for row in vector_results} + kept_row_keys.update( + rejection.key + for rejection_group in ( + trace.vector.threshold_rejections, + trace.vector.filter_rejections, + trace.vector.missing_search_rows, + ) + for rejection in rejection_group + ) + if any(match.key not in kept_row_keys for match in trace.vector.chunk_matches): + fused_chunks: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} + for chunk_match in trace.vector.chunk_matches: + if chunk_match.key in kept_row_keys: + fused_chunks.setdefault(chunk_match.key, []).append( + (chunk_match.chunk_key, chunk_match.similarity, chunk_match.entity_id) + ) + trace.vector = build_vector_stage( + previous=trace.vector, + chunk_matches=fused_chunks, + ) + fusion_start = time.perf_counter() + + with logfire.span( + "search.fusion", fts_count=len(fts_results), vector_count=len(vector_results) + ) as fusion_span: + # --- Score-based fusion keyed on (type, id) --- + # A bare row id collides across row types (independent id sequences), so + # fusion must key on (type, id) or distinct rows would merge (#982). + # FTS scores are normalized to [0, 1] (BM25 is unbounded). + # Vector scores are used raw: the adapters already calibrate them to [0, 1]. + rows_by_key: dict[SearchIndexKey, SearchIndexRow] = {} + + # Normalize FTS scores to [0, 1] — handles both SQLite (negative bm25) + # and Postgres (positive ts_rank) by using absolute values + fts_abs = [abs(row.score or 0.0) for row in fts_results] + fts_max = max(fts_abs) if fts_abs else 1.0 + + fts_scores: dict[SearchIndexKey, float] = {} + fts_ranks: dict[SearchIndexKey, int] = {} + for rank, row in enumerate(fts_results): + if row.id is None: + continue + row_key = (row.type, row.id) + norm = abs(row.score or 0.0) / fts_max if fts_max > 0 else 0.0 + # Gate: FTS scores below threshold contribute zero + if norm < FTS_GATE_THRESHOLD: + norm = 0.0 + fts_scores[row_key] = norm + fts_ranks.setdefault(row_key, rank) + rows_by_key[row_key] = row + + if trace is not None: + relaxed_fallback_used = ( + trace.fts.relaxed_fallback_used if trace.fts is not None else False + ) + trace.fts = build_fts_page_stage( + [((row.type, row.id), row.score or 0.0) for row in fts_results], + normalized_scores=fts_scores, + entity_ids={(row.type, row.id): row.entity_id for row in fts_results}, + fts_max_abs=fts_max, + relaxed_fallback_used=relaxed_fallback_used, + fts_ms=fts_ms, + ) + + vec_scores: dict[SearchIndexKey, float] = {} + vec_ranks: dict[SearchIndexKey, int] = {} + for rank, row in enumerate(vector_results): + if row.id is None: + continue + row_key = (row.type, row.id) + # Trigger: no re-normalization by vec_max + # Why: vector similarity is already calibrated [0, 1]; re-normalizing + # inflates weak matches when the entire result set is mediocre + vec_scores[row_key] = row.score or 0.0 + vec_ranks.setdefault(row_key, rank) + rows_by_key[row_key] = row + + # Fuse: max(v, f) + FUSION_BONUS * min(v, f) + # Preserves the dominant signal; bonus rewards dual-source agreement. + # Output range: [0, 1.3] for dual-source, [0, 1.0] for single-source. + fused_scores: dict[SearchIndexKey, float] = {} + for row_key in fts_scores.keys() | vec_scores.keys(): + v = vec_scores.get(row_key, 0.0) + f = fts_scores.get(row_key, 0.0) + fused_scores[row_key] = max(v, f) + FUSION_BONUS * min(v, f) + + ranked = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True) + fusion_span.set_attribute("result_count", len(ranked)) + fusion_ms = (time.perf_counter() - fusion_start) * 1000 + if trace is not None: + trace.fusion = build_fusion_stage( + formula_version=FUSION_FORMULA_VERSION, + bonus=FUSION_BONUS, + fts_scores=fts_scores, + fts_ranks=fts_ranks, + vector_scores=vec_scores, + vector_ranks=vec_ranks, + ranked_scores=ranked, + fusion_ms=fusion_ms, + ) + + def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: + row_key, fused_score = entry + row = rows_by_key[row_key] + # FTS-only hits use the bounded content preview and its truncation metadata. + # Copying the full note into matched_chunk bypasses that response bound. + return replace(row, score=fused_score) + + # Rerank the top fused candidates before paginating. When reranking is active + # we materialize the whole candidate list (cheap next to a cross-encoder call) + # and hand it to the shared paginate helper; the disabled path stays cheap by + # materializing only the requested page. + if apply_rerank and rerank is not None: + candidates = [_materialize(entry) for entry in ranked] + stable_candidates = candidates + stable_candidate_limit = self._rerank_candidate_limit(rerank) + if candidate_limit > stable_candidate_limit: + if trace is not None: + trace.stable_pool_refetched = True + stable_candidates = await self.hybrid( + query, + limit=stable_candidate_limit, + offset=0, + candidate_limit=stable_candidate_limit, + apply_rerank=False, + emit_observability_log=False, + trace=None, + ) + stable_keys = {(row.type, row.id) for row in stable_candidates} + expanded_tail = [entry for entry in ranked if entry[0] not in stable_keys] + + # Trigger: deeper pages expand the FTS/vector retrieval windows. + # Why: score fusion can strengthen an existing row when its second + # signal appears later, moving it across a page already returned. + # Outcome: freeze the fixed fused universe, then order newly admitted + # rows by their earliest source rank. That rank cannot improve after a + # row first appears, so each larger window only appends to the tail. + expanded_tail.sort( + key=lambda entry: ( + min( + fts_ranks.get(entry[0], candidate_limit), + vec_ranks.get(entry[0], candidate_limit), + ), + entry[0], + ) + ) + candidates = stable_candidates + [_materialize(entry) for entry in expanded_tail] + output = await self._rerank_and_paginate( + query_text, + candidates, + offset=offset, + limit=limit, + stable_rows=stable_candidates, + trace=trace, + ) + else: + output = [_materialize(entry) for entry in ranked[offset : offset + limit]] + total_ms = (time.perf_counter() - query_start) * 1000 + if emit_observability_log and total_ms > 2500: + logger.warning( + "[SEMANTIC_SLOW_QUERY] Semantic query timing: scope={scope} " + "retrieval_mode={retrieval_mode} query_length={query_length} " + "candidate_limit={candidate_limit} fts_count={fts_count} " + "vector_count={vector_count} fts_ms={fts_ms:.2f} vector_ms={vector_ms:.2f} " + "fusion_ms={fusion_ms:.2f} total_ms={total_ms:.2f}", + scope=self.scope.project_ids, + retrieval_mode="hybrid", + query_length=len(query_text), + candidate_limit=candidate_limit, + fts_count=len(fts_results), + vector_count=len(vector_results), + fts_ms=fts_ms, + vector_ms=vector_ms, + fusion_ms=fusion_ms, + total_ms=total_ms, + ) + return output + + +# --- The reader --- + + +class SearchReader: + """Run one prepared query over one scope, whichever retrieval mode it asks for.""" + + def __init__( + self, + scope: ProjectScope, + fts: FtsBackend, + semantic: SemanticSearch | None = None, + ) -> None: + self.scope = scope + self.fts = fts + self.semantic = semantic + + def _semantic(self) -> SemanticSearch: + if self.semantic is None: + raise SemanticSearchDisabledError( + "Semantic search is disabled. Set BASIC_MEMORY_SEMANTIC_SEARCH_ENABLED=true." + ) + return self.semantic + + async def search( + self, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, + allow_relaxed: bool = False, + session: AsyncSession | None = None, + candidate_keys: Sequence[SearchIndexKey] | None = None, + trace: SearchTraceCollector | None = None, + ) -> list[SearchIndexRow]: + """Search the scope in the query's retrieval mode. + + ``candidate_keys`` restricts full-text results to those ``(type, id)`` search + rows. ``None`` searches the whole scope; an empty sequence matches nothing. + Vector and hybrid retrieval use that restriction to ask which of a known + candidate set a filter admits instead of paging the filter's whole match set + (#1431). + + ``allow_relaxed=True`` retries a zero-result strict multi-word query with + OR-joined content terms. Only the hybrid path opts in: its FTS branch otherwise + contributes nothing for question-form queries. + """ + match query.retrieval_mode: + case SearchRetrievalMode.FTS: + return await self.fts.search( + self.scope, + query, + limit=limit, + offset=offset, + allow_relaxed=allow_relaxed, + session=session, + candidate_keys=candidate_keys, + trace=trace, + ) + case SearchRetrievalMode.VECTOR: + if not vector_eligible(query): + raise ValueError( + "Vector retrieval requires a non-empty text query and does not support " + "title/permalink-only searches." + ) + return await self._semantic().vector_only( + query, limit=limit, offset=offset, trace=trace + ) + case SearchRetrievalMode.HYBRID: + if not vector_eligible(query): + raise ValueError( + "Hybrid retrieval requires a non-empty text query and does not support " + "title/permalink-only searches." + ) + return await self._semantic().hybrid(query, limit=limit, offset=offset, trace=trace) + case _: # pragma: no cover + assert_never(query.retrieval_mode) + + async def count(self, query: PreparedSearchQuery, *, allow_relaxed: bool = False) -> int: + """Count full-text matches with the same filters as ``search``.""" + if query.retrieval_mode != SearchRetrievalMode.FTS: + raise ValueError("Exact counts are only supported for full-text search retrieval.") + return await self.fts.count(self.scope, query, allow_relaxed=allow_relaxed) diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 31f97096a..15e2dd506 100644 --- a/src/basic_memory/repository/search_repository_base.py +++ b/src/basic_memory/repository/search_repository_base.py @@ -5,7 +5,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Mapping, Sequence from contextlib import asynccontextmanager -from dataclasses import dataclass, replace +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, Callable, Dict, List, Literal, Optional, cast @@ -26,31 +26,22 @@ from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) -from basic_memory.repository.rerank_provider import ( - RerankProvider, - build_rerank_document, - demote_tail_scores, - validate_rerank_scores, -) +from basic_memory.repository.rerank_provider import RerankProvider from basic_memory.repository.search_filters import FtsBackend from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + BUILT_IN_VECTOR_INDEX_NAMES, + VECTOR_HYDRATION_BATCH_SIZE, + Reranking, + SearchReader, + SemanticSearch, + VectorRetrieval, + vector_eligible, +) from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.script_ngrams import build_script_ngrams -from basic_memory.repository.search_trace import ( - BelowThreshold, - FilteredOut, - HydrationDropKey, - HydrationDropped, - MissingSearchRow, - SearchTraceCollector, - build_fts_page_stage, - build_fusion_stage, - build_rerank_stage, - build_vector_stage, - classify_hydration_drops, - read_manifest_readiness, -) +from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.repository.semantic_chunking import ( SemanticSourceRow, VectorChunkRecord, @@ -68,7 +59,6 @@ SemanticVectorIndexReconciler, VectorDeletion, VectorKey, - VectorMatch, VectorRecord, ) from basic_memory.repository.semantic_vector_sync import ( @@ -88,45 +78,10 @@ from basic_memory.temporal import TemporalFilter from basic_memory.utils import ensure_timezone_aware -# --- 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 - -# The manifest conditions under which semantic retrieval will use a stored vector. -# Vector hydration (_hydrate_vector_matches) admits exactly these rows, so anything -# failing them is invisible to search: a chunk left behind by an embedding-model or -# vector-index change, or one still pending. Readiness reporting must apply the same -# predicate — calling such a row "embedded" would report an index settled that -# retrieval cannot answer from, which is the class of lie #1414 exists to remove. -# Callers bind :vector_index and :embedding_model; the scope binds its own IDs. - - -def current_vector_manifest_predicate(scope: ProjectScope, params: dict[str, Any]) -> str: - """SQL admitting only manifest rows retrieval can answer from, within ``scope``.""" - return ( - f"{scope.predicate('project_id', params)} " - "AND vector_index = :vector_index " - "AND embedding_model = :embedding_model " - "AND embedding_status = 'ready'" - ) - +# --- Semantic sync constants --- -# 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 -# to keep enough unique documents in the rerank window. -RERANK_POOL_CHUNK_FANOUT = 4 -FUSION_BONUS = 0.3 -FUSION_FORMULA_VERSION = "max+0.3*min/v1" -FTS_GATE_THRESHOLD = 0.0 -TOP_CHUNKS_PER_RESULT = 5 -SMALL_NOTE_CONTENT_LIMIT = 2000 OVERSIZED_ENTITY_VECTOR_SHARD_SIZE = semantic_vector_sync.OVERSIZED_ENTITY_VECTOR_SHARD_SIZE _SQLITE_MAX_PREPARE_WINDOW = semantic_vector_sync.SQLITE_MAX_PREPARE_WINDOW -_BUILT_IN_VECTOR_INDEX_NAMES = frozenset({"pgvector", "sqlite-vec"}) type StoredEmbeddingStatus = Literal["pending", "ready"] @@ -216,8 +171,9 @@ class SearchRepositoryBase(ABC): This class defines the common interface that all search repositories must implement, regardless of whether they use SQLite FTS5 or Postgres tsvector for full-text search. - Shared semantic search logic (chunking, embedding orchestration, hybrid score-based fusion) - lives here. Backend-specific operations are delegated to abstract hooks. + Indexing, vector-manifest writes, and embedding orchestration live here. Reading + is delegated to ``SearchReader``, built per call from this repository's current + state so the reader sees the same semantic capability the repository has. Concrete implementations: - SQLiteSearchRepository: Uses FTS5 virtual tables with MATCH queries @@ -335,39 +291,10 @@ async def search( ) -> List[SearchIndexRow]: """Search this repository's project. - ``candidate_keys`` restricts results to those ``(type, id)`` search rows. - ``None`` searches the whole scope; 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). - - ``allow_relaxed=True`` retries a zero-result strict multi-word query with - OR-joined content terms. Only the hybrid path opts in: its FTS branch otherwise - contributes nothing for question-form queries. + The reader owns retrieval. This method owns what only the repository knows: + whether its semantic stack is enabled and its vector tables exist. See + ``SearchReader.search`` for ``candidate_keys`` and ``allow_relaxed``. """ - # --- Vector and hybrid: shared retrieval over this repository's scope --- - dispatched = await self._dispatch_retrieval_mode( - search_text=search_text, - 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=retrieval_mode, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - if dispatched is not None: - return dispatched - - # --- Full text: the engine runs the compiled statement --- query = PreparedSearchQuery( search_text=search_text, permalink=permalink, @@ -383,8 +310,8 @@ async def search( retrieval_mode=retrieval_mode, min_similarity=min_similarity, ) - return await self._fts.search( - self.scope, + reader = await self._reader_for(query) + return await reader.search( query, limit=limit, offset=offset, @@ -412,8 +339,6 @@ async def count( allow_relaxed: bool = False, ) -> int: """Count full-text matches with the same filters as ``search``.""" - if retrieval_mode != SearchRetrievalMode.FTS: - raise ValueError("Exact counts are only supported for full-text search retrieval.") query = PreparedSearchQuery( search_text=search_text, permalink=permalink, @@ -429,7 +354,50 @@ async def count( retrieval_mode=retrieval_mode, min_similarity=min_similarity, ) - return await self._fts.count(self.scope, query, allow_relaxed=allow_relaxed) + return await SearchReader(self.scope, self._fts).count(query, allow_relaxed=allow_relaxed) + + # --- Reader construction --- + + def _reranking(self) -> Reranking | None: + """The configured cross-encoder, or None when this repository does not rerank.""" + if self._rerank_provider is None: + return None + return Reranking( + provider=self._rerank_provider, + candidates=self._reranker_candidates, + max_document_chars=self._reranker_max_document_chars, + ) + + def _semantic_search(self) -> SemanticSearch: + """Vector and hybrid retrieval over this repository's live semantic stack. + + Valid once ``_ensure_vector_tables`` has bound the adapter. Read from the + current attributes on every call because the semantic flag can flip at + runtime (#711) and tests retune thresholds between searches. + """ + assert self._embedding_provider is not None + vector = VectorRetrieval( + index=self._semantic_vector_index, + index_name=self._semantic_vector_index_name, + embedding_provider=self._embedding_provider, + embedding_model=self._embedding_model_key(), + vector_k=self._semantic_vector_k, + min_similarity=self._semantic_min_similarity, + ) + return SemanticSearch(self.session_maker, self.scope, self._fts, vector, self._reranking()) + + async def _reader_for(self, query: PreparedSearchQuery) -> SearchReader: + """Build the reader for one call from this repository's current state.""" + # Trigger: the query asks for vector or hybrid retrieval and has text to embed. + # Why: whether semantic search is enabled and whether the vector tables and + # adapter exist is repository lifecycle; the reader only reads. + # Outcome: the semantic gate raises its typed error before any retrieval + # runs, and the reader receives the bound adapter. + if query.retrieval_mode != SearchRetrievalMode.FTS and vector_eligible(query): + self._assert_semantic_available() + await self._ensure_vector_tables() + return SearchReader(self.scope, self._fts, self._semantic_search()) + return SearchReader(self.scope, self._fts) # ------------------------------------------------------------------ # Abstract methods — semantic search (backend-specific DB operations) @@ -440,99 +408,6 @@ async def _ensure_vector_tables(self) -> None: """Create backend-specific vector chunk and embedding tables.""" pass - @logfire.instrument("search.vector_query", extract_args=False) - async def _run_vector_query( - self, - session: AsyncSession, - query_embedding: list[float], - candidate_limit: int, - *, - trace: SearchTraceCollector | None = None, - ) -> list[dict[str, Any]]: - """Query the configured adapter and hydrate only live, ready manifest rows.""" - if trace is not None: - trace.vector = build_vector_stage( - candidate_limit=candidate_limit, - adapter_match_count=0, - hydrated_count=0, - ) - if candidate_limit <= 0: - return [] - - external_vector_index = self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES - if not external_vector_index: - matches = await self._semantic_vector_index.search( - query_embedding, - limit=candidate_limit, - ) - if trace is not None: - trace.readiness = await read_manifest_readiness( - session, - self.scope, - self._semantic_vector_index_name, - self._embedding_model_key(), - ) - return await self._hydrate_vector_matches(session, matches, trace=trace) - - scan_limit = min(candidate_limit, VECTOR_FILTER_SCAN_LIMIT) - while True: - matches = await self._semantic_vector_index.search( - query_embedding, - limit=scan_limit, - ) - if trace is not None and trace.readiness is None: - trace.readiness = await read_manifest_readiness( - session, - self.scope, - self._semantic_vector_index_name, - self._embedding_model_key(), - ) - hydrated = await self._hydrate_vector_matches(session, matches, trace=trace) - if ( - len(hydrated) >= candidate_limit - or len(matches) < scan_limit - or scan_limit >= VECTOR_FILTER_SCAN_LIMIT - ): - returned = hydrated[:candidate_limit] - # Trigger: the expanded stale-hit rescan hydrated more chunks than the - # candidate window the search consumes. - # Why: chunks beyond the window never enter thresholding, fusion, or - # reranking — tracing them would invent candidates this execution - # never considered. - # Outcome: the traced stage is trimmed to the returned window. - if trace is not None and trace.vector is not None and len(hydrated) > len(returned): - # Two owners can share one parseable chunk_key (manifest uniqueness - # includes entity_id), so window membership matches by owner too. - returned_chunk_keys = { - (int(row["entity_id"]), str(row["chunk_key"])) for row in returned - } - trimmed: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} - for chunk_match in trace.vector.chunk_matches: - if (chunk_match.entity_id, chunk_match.chunk_key) in returned_chunk_keys: - trimmed.setdefault(chunk_match.key, []).append( - ( - chunk_match.chunk_key, - chunk_match.similarity, - chunk_match.entity_id, - ) - ) - # hydrated_count keeps full-scan scope so the vector stage's - # dropped count matches its hydration-drop list; the flattener - # reports the window truncation as its own candidate_window stage. - trace.vector = build_vector_stage( - previous=trace.vector, - chunk_matches=trimmed, - ) - return returned - - # Trigger: stale, pending, or wrong-model adapter hits consumed the - # requested top-k before manifest hydration. - # Why: returning early lets stale extension data crowd every live - # result out of an otherwise valid semantic search. - # Outcome: retry from the same ranked prefix with bounded geometric - # overfetch until enough live rows survive or the adapter is exhausted. - scan_limit = min(scan_limit * 2, VECTOR_FILTER_SCAN_LIMIT) - async def record_entity_vector_deferrals( self, *, @@ -595,110 +470,6 @@ async def record_entity_vector_deferrals( ) await session.commit() - @logfire.instrument("search.vector_manifest_hydration", extract_args=False) - async def _hydrate_vector_matches( - self, - session: AsyncSession, - matches: list[VectorMatch], - *, - trace: SearchTraceCollector | None = None, - ) -> list[dict[str, Any]]: - """Resolve adapter matches through the authoritative ready manifest.""" - if not matches: - return [] - - chunks_by_key: dict[VectorKey, str] = {} - for batch_start in range(0, len(matches), VECTOR_HYDRATION_BATCH_SIZE): - batch = matches[batch_start : batch_start + VECTOR_HYDRATION_BATCH_SIZE] - params: dict[str, Any] = { - "vector_index": self._semantic_vector_index_name, - "embedding_model": self._embedding_model_key(), - } - manifest_predicate = current_vector_manifest_predicate(self.scope, params) - predicates: list[str] = [] - for index, match in enumerate(batch): - params[f"entity_id_{index}"] = match.key.entity_id - params[f"chunk_key_{index}"] = match.key.chunk_key - predicates.append( - f"(entity_id = :entity_id_{index} AND chunk_key = :chunk_key_{index})" - ) - - # Constraint: adapters may return thousands of candidates for deep pages. - # PostgreSQL and SQLite both cap bind parameters, so hydrate in fixed-size - # batches while retaining the adapter's original ranking in the final list. - result = await session.execute( - text( - "SELECT entity_id, chunk_key, chunk_text FROM search_vector_chunks " - "WHERE " + manifest_predicate + " " - "AND (" + " OR ".join(predicates) + ")" - ), - params, - ) - chunks_by_key.update( - { - VectorKey( - entity_id=int(row["entity_id"]), - chunk_key=str(row["chunk_key"]), - ): str(row["chunk_text"]) - for row in result.mappings().all() - } - ) - hydrated = [ - { - "entity_id": match.key.entity_id, - "chunk_key": match.key.chunk_key, - "chunk_text": chunks_by_key[match.key], - "best_similarity": match.similarity, - } - for match in matches - if match.key in chunks_by_key - ] - if trace is not None: - dropped_keys = [ - HydrationDropKey( - entity_id=match.key.entity_id, - chunk_key=match.key.chunk_key, - similarity=match.similarity, - configured_index=self._semantic_vector_index_name, - configured_model=self._embedding_model_key(), - ) - for match in matches - if match.key not in chunks_by_key - ] - drops = await classify_hydration_drops(session, self.scope, dropped_keys) - chunk_matches: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} - malformed_drops: list[HydrationDropped] = [] - for row in hydrated: - try: - key = self._parse_chunk_key(str(row["chunk_key"])) - except (ValueError, IndexError): - # A hydrated chunk with an unparseable key silently vanishes from - # retrieval; the trace must name it or the stage counts lie. - malformed_drops.append( - HydrationDropped( - entity_id=int(row["entity_id"]), - chunk_key=str(row["chunk_key"]), - similarity=float(row["best_similarity"]), - reason="malformed_key", - stored_model=None, - stored_index=None, - ) - ) - continue - chunk_matches.setdefault(key, []).append( - (str(row["chunk_key"]), float(row["best_similarity"]), int(row["entity_id"])) - ) - trace.vector = build_vector_stage( - previous=trace.vector, - adapter_match_count=len(matches), - # Malformed keys are dropped, not served — counting them as output - # would contradict the malformed_key rejection listed alongside. - hydrated_count=len(hydrated) - len(malformed_drops), - drops=(*drops, *malformed_drops), - chunk_matches=chunk_matches, - ) - return hydrated - async def _write_embeddings( self, session: AsyncSession, @@ -748,7 +519,7 @@ async def _persist_embeddings( connection = await session.connection() dialect_name = connection.dialect.name external_vector_index = ( - self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + self._semantic_vector_index_name not in BUILT_IN_VECTOR_INDEX_NAMES ) lock_external_write = external_vector_index and dialect_name in {"postgresql", "sqlite"} if external_vector_index: @@ -1027,7 +798,7 @@ async def _finalize_prepared_vector_deletions( for deletion in prepared.staged_deletions ] - external_vector_index = self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES + external_vector_index = self._semantic_vector_index_name not in BUILT_IN_VECTOR_INDEX_NAMES if external_vector_index: async with db.scoped_session(self.session_maker) as session: connection = await session.connection() @@ -1076,7 +847,7 @@ async def _finalize_prepared_vector_deletions( return await self._semantic_vector_index.delete(deletions) - if self._semantic_vector_index_name in _BUILT_IN_VECTOR_INDEX_NAMES: + if self._semantic_vector_index_name in BUILT_IN_VECTOR_INDEX_NAMES: return async with db.scoped_session(self.session_maker) as session: await session.execute( @@ -1089,16 +860,6 @@ async def _finalize_prepared_vector_deletions( ) await session.commit() - @abstractmethod - def _distance_to_similarity(self, distance: float) -> float: - """Convert a backend-specific vector distance to cosine similarity in [0, 1]. - - Backend-specific implementations: - - SQLite (vec0): L2/Euclidean distance → cosine similarity via 1 - d²/2 - - Postgres (pgvector <=>): Cosine distance → cosine similarity via 1 - d - """ - pass # pragma: no cover - # ------------------------------------------------------------------ # Shared index / delete operations # ------------------------------------------------------------------ @@ -1393,7 +1154,7 @@ async def _delete_external_entity_vectors_locked( ) -> None: """Delete external vectors after the caller has acquired the project lock.""" self._assert_manifest_vector_ownership(recorded_indexes) - external_indexes = recorded_indexes - _BUILT_IN_VECTOR_INDEX_NAMES + external_indexes = recorded_indexes - BUILT_IN_VECTOR_INDEX_NAMES if not external_indexes: return @@ -1491,14 +1252,14 @@ async def _lock_external_vector_write(self, session: AsyncSession) -> None: def _uses_external_vector_index(self) -> bool: """Return whether this repository writes vectors outside the SQL backend.""" - return self._semantic_vector_index_name not in _BUILT_IN_VECTOR_INDEX_NAMES and hasattr( + return self._semantic_vector_index_name not in BUILT_IN_VECTOR_INDEX_NAMES and hasattr( self, "_semantic_vector_index" ) def _assert_manifest_vector_ownership(self, vector_index_names: Iterable[object]) -> None: """Reject cleanup that cannot reach every externally owned vector.""" recorded_indexes = frozenset(str(name) for name in vector_index_names if str(name)) - external_indexes = recorded_indexes - _BUILT_IN_VECTOR_INDEX_NAMES + external_indexes = recorded_indexes - BUILT_IN_VECTOR_INDEX_NAMES configured_index = self._semantic_vector_index_name if external_indexes and ( not hasattr(self, "_semantic_vector_index") @@ -1565,7 +1326,7 @@ async def _delete_project_vector_rows_in_session( manifest_has_embedding_status = "embedding_status" in manifest_columns configured_index = self._semantic_vector_index_name external_adapter_available = ( - configured_index not in _BUILT_IN_VECTOR_INDEX_NAMES + configured_index not in BUILT_IN_VECTOR_INDEX_NAMES and hasattr(self, "_semantic_vector_index") ) if external_adapter_available: @@ -1606,7 +1367,7 @@ async def _delete_project_vector_rows_in_session( # Outcome: fail before touching any adapter or manifest so the owner can be restored. self._assert_manifest_vector_ownership(entity_ids_by_vector_index) - builtin_indexes = frozenset(entity_ids_by_vector_index) & _BUILT_IN_VECTOR_INDEX_NAMES + builtin_indexes = frozenset(entity_ids_by_vector_index) & BUILT_IN_VECTOR_INDEX_NAMES if manifest_has_embedding_status and (not external_adapter_available or builtin_indexes): builtin_filter = "" if external_adapter_available: @@ -2055,948 +1816,3 @@ def _timestamp_now_expr(self) -> str: SQLite uses CURRENT_TIMESTAMP, Postgres uses NOW(). """ return "CURRENT_TIMESTAMP" - - # ------------------------------------------------------------------ - # Shared semantic search: retrieval mode dispatch - # ------------------------------------------------------------------ - - def _check_vector_eligible( - self, - search_text: Optional[str], - permalink: Optional[str], - permalink_match: Optional[str], - title: Optional[str], - ) -> bool: - """Check whether search_text allows vector / hybrid retrieval.""" - return ( - bool(search_text) - and bool(search_text.strip()) - and search_text.strip() != "*" - and not permalink - and not permalink_match - and not title - ) - - async def _dispatch_retrieval_mode( - self, - *, - search_text: Optional[str], - 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], - retrieval_mode: SearchRetrievalMode, - min_similarity: Optional[float] = None, - limit: int, - offset: int, - trace: SearchTraceCollector | None = None, - ) -> Optional[List[SearchIndexRow]]: - """Dispatch vector or hybrid retrieval if requested. - - Returns None when the mode is FTS so the caller should continue - with its backend-specific FTS query. - """ - mode = ( - retrieval_mode.value - if isinstance(retrieval_mode, SearchRetrievalMode) - else str(retrieval_mode) - ) - can_use_vector = self._check_vector_eligible(search_text, permalink, permalink_match, title) - search_text_value = search_text or "" - - if mode == SearchRetrievalMode.VECTOR.value: - if not can_use_vector: - raise ValueError( - "Vector retrieval requires a non-empty text query and does not support " - "title/permalink-only searches." - ) - return await self._search_vector_only( - search_text=search_text_value, - 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, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - if mode == SearchRetrievalMode.HYBRID.value: - if not can_use_vector: - raise ValueError( - "Hybrid retrieval requires a non-empty text query and does not support " - "title/permalink-only searches." - ) - return await self._search_hybrid( - search_text=search_text_value, - 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, - min_similarity=min_similarity, - limit=limit, - offset=offset, - trace=trace, - ) - - # FTS mode: return None to let the subclass handle it - return None - - # ------------------------------------------------------------------ - # Shared semantic search: vector-only retrieval - # ------------------------------------------------------------------ - - @staticmethod - def _parse_chunk_key(chunk_key: str) -> SearchIndexKey: - """Parse a chunk_key like 'observation:5:0' into (type, search_index_id).""" - parts = chunk_key.split(":") - return parts[0], int(parts[1]) - - # ------------------------------------------------------------------ - # Shared semantic search: cross-encoder reranking - # ------------------------------------------------------------------ - - def _should_rerank(self, query_text: str) -> bool: - """Return whether a configured reranker should run for this query.""" - return self._rerank_provider is not None and bool(query_text) - - def _rerank_candidate_limit(self) -> int: - """Return the fixed chunk window that owns reranker-prefix membership.""" - return max( - self._semantic_vector_k, - self._reranker_candidates * RERANK_POOL_CHUNK_FANOUT, - ) - - def _candidate_limit(self, limit: int, offset: int, query_text: str) -> int: - """Size the retrieval candidate *chunk* pool for vector/hybrid search. - - ``candidate_limit`` bounds vector chunks, but many chunks of one large note - collapse to a single ``(type, id)`` row before reranking, so a chunk count does - not equal a unique-document count. When reranking is active we over-fetch by - ``RERANK_POOL_CHUNK_FANOUT`` so a few multi-chunk notes can't starve the rerank - window below ``reranker_candidates`` unique rows. This is best-effort headroom, - not a hard guarantee — a single note dominating the entire nearest-neighbour set - can still yield fewer unique rows (a pathological corpus shape). - """ - if self._should_rerank(query_text): - # Trigger: the requested window extends beyond the fixed reranked prefix. - # Why: a bounded prefix alone can under-fill large pages and hide the - # semantic pagination probe even when more matches exist. - # Outcome: keep prefix membership fixed while adding chunk headroom only - # for the untouched tail that this request must return. - rerank_candidate_limit = self._rerank_candidate_limit() - tail_size = max(0, limit + offset - self._reranker_candidates) - return rerank_candidate_limit + tail_size * 10 - return max(self._semantic_vector_k, (limit + offset) * 10) - - def _rerank_document_text(self, row: SearchIndexRow) -> str: - """Build the document text handed to the cross-encoder for one candidate. - - Prefer the matched chunk (the most relevant passage of a large note), - falling back to the stored snippet. - """ - body = row.matched_chunk_text or row.content_snippet or "" - return build_rerank_document(row.title, body, self._reranker_max_document_chars) - - @staticmethod - def _demote_tail(tail: list[SearchIndexRow], floor: float) -> list[SearchIndexRow]: - """Rescore un-reranked tail rows at or below the floor, preserving their order. - - The reranked pool carries [0, 1] relevance scores while the tail still holds - raw retrieval scores on a different scale ([0, 1.3] for fused hybrid). Left as - is, a tail row could outrank a reranked row numerically. Positive floors put - the tail strictly below the pool; a zero floor yields zeroes because no smaller - score exists in the public [0, 1] range. The returned pool-plus-tail sequence, - rather than a later score-only sort, owns that tie-breaking invariant. - """ - return [ - replace(row, score=score) - for row, score in zip(tail, demote_tail_scores(floor, len(tail))) - ] - - async def _rerank_and_paginate( - self, - query_text: str, - rows: list[SearchIndexRow], - *, - offset: int, - limit: int, - stable_rows: list[SearchIndexRow] | None = None, - trace: SearchTraceCollector | None = None, - ) -> list[SearchIndexRow]: - """Rerank the top candidates, then return the requested ``[offset:offset+limit]`` page. - - Trigger: a reranker is configured and there is a real query. - Why: bi-encoder/FTS ranking lands the gold document in the top-N but often - just below the top-k cutoff (#950); a cross-encoder that reads query and - document together recovers those near-misses. - Outcome: the first ``reranker_candidates`` rows are reordered by reranker - relevance (which replaces ``score``); the requested page is sliced from the - reordered list. - - Every non-empty page rescores the same fixed prefix before slicing so the - untouched tail can be demoted onto the reranker's public ``[0, 1]`` scale. - """ - page_end = offset + limit - if self._rerank_provider is None or not query_text: - return rows[offset:page_end] - - # Trigger: pagination needs more rows than the fixed rerank retrieval window. - # Why: an expanded retrieval may introduce or strengthen raw candidates, but - # letting them replace the original prefix causes duplicates and skips. - # Outcome: the fixed window owns prefix membership; the expanded result only - # supplies new, de-duplicated tail rows. - pool_source = stable_rows if stable_rows is not None else rows - pool = pool_source[: self._reranker_candidates] - pool_keys = {(row.type, row.id) for row in pool} - tail = [row for row in rows if (row.type, row.id) not in pool_keys] - ordered_rows = pool + tail - - # Skip only when there is no prefix to calibrate or the requested page is - # empty. Even a singleton prefix or a wholly-tail page needs the prefix's - # relevance floor so raw hybrid scores cannot leak into cross-project sorting. - if not pool or offset >= len(ordered_rows): - return ordered_rows[offset:page_end] - - pre_rerank_scores = None - if trace is not None: - pre_rerank_scores = {(row.type, row.id): row.score or 0.0 for row in ordered_rows} - documents = [self._rerank_document_text(row) for row in pool] - # A transient provider failure must surface instead of switching this page - # back to retrieval order. A prior page may already have returned reranked - # order, so degrading here can duplicate one result and omit another. - rerank_start = time.perf_counter() if trace is not None else None - with logfire.span( - "search.rerank", - candidate_count=len(pool), - document_chars=sum(map(len, documents)), - ): - scores = validate_rerank_scores( - await self._rerank_provider.rerank(query_text, documents), - len(pool), - ) - - order = sorted(range(len(pool)), key=lambda i: scores[i], reverse=True) - reranked = [replace(pool[i], score=scores[i]) for i in order] - logger.debug( - "Reranked candidates: pool={pool} model={model}", - pool=len(pool), - model=self._rerank_provider.model_name, - ) - tail_floor = reranked[-1].score or 0.0 - demoted_tail = self._demote_tail(tail, floor=tail_floor) - reranked_rows = reranked + demoted_tail - if trace is not None: - assert pre_rerank_scores is not None and rerank_start is not None - trace.rerank = build_rerank_stage( - provider_model=self._rerank_provider.model_name, - reranker_candidates=self._reranker_candidates, - pre_rerank_scores=pre_rerank_scores, - pool_keys=[(row.type, row.id) for row in pool], - rerank_scores={ - (pool[index].type, pool[index].id): score for index, score in enumerate(scores) - }, - post_rerank_rows=[((row.type, row.id), row.score or 0.0) for row in reranked_rows], - demoted_scores={(row.type, row.id): row.score or 0.0 for row in demoted_tail}, - tail_floor=tail_floor, - stable_pool_refetched=trace.stable_pool_refetched, - rerank_ms=(time.perf_counter() - rerank_start) * 1000, - ) - return reranked_rows[offset:page_end] - - async def _search_vector_only( - self, - *, - search_text: str, - 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], - min_similarity: Optional[float] = None, - limit: int, - offset: int, - candidate_limit: int | None = None, - _emit_observability_log: bool = True, - _apply_rerank: bool = True, - trace: SearchTraceCollector | None = None, - ) -> List[SearchIndexRow]: - """Run vector-only search returning chunk-level results. - - Returns individual search_index rows (entities, observations, relations) - ranked by vector similarity. Each observation or relation is a first-class - result, not collapsed into its parent entity. - - ``candidate_limit`` is supplied only by a composed retrieval stage that - already sized the shared candidate pool. - """ - self._assert_semantic_available() - await self._ensure_vector_tables() - assert self._embedding_provider is not None - query_text = search_text.strip() - if candidate_limit is None: - candidate_limit = self._candidate_limit(limit, offset, query_text) - query_start = time.perf_counter() - embed_start = time.perf_counter() - with logfire.span("search.embed_query", query_chars=len(query_text)): - query_embedding = await self._embedding_provider.embed_query(query_text) - embed_ms = (time.perf_counter() - embed_start) * 1000 - vector_query_start = time.perf_counter() - - if hasattr(self, "_semantic_vector_index"): - # Constraint: vector adapters may open their own session, while the SQLite - # test/runtime pool can contain only one connection. A plain AsyncSession - # defers checkout until hydration runs after adapter search has released it. - async with self.session_maker() as session: - if trace is None: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - ) - else: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - else: - # Compatibility for focused test repositories that implement the - # pre-extension private query hook without configuring an adapter. - async with db.scoped_session(self.session_maker) as session: - await self._prepare_vector_session(session) - if trace is None: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - ) - else: - vector_rows = await self._run_vector_query( - session, - query_embedding, - candidate_limit, - trace=trace, - ) - vector_query_ms = (time.perf_counter() - vector_query_start) * 1000 - vector_row_count = len(vector_rows) - hydrate_ms = 0.0 - - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - effective_min_similarity=( - min_similarity if min_similarity is not None else self._semantic_min_similarity - ), - min_similarity_source=("query" if min_similarity is not None else "config"), - embed_ms=embed_ms, - vector_query_ms=vector_query_ms, - ) - - def _log_vector_summary() -> None: - if not _emit_observability_log: - return - - total_ms = (time.perf_counter() - query_start) * 1000 - if total_ms > 2000: - logger.warning( - "[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} " - "retrieval_mode={retrieval_mode} query_length={query_length} " - "candidate_limit={candidate_limit} vector_row_count={vector_row_count} " - "embed_ms={embed_ms:.2f} vector_query_ms={vector_query_ms:.2f} " - "hydrate_ms={hydrate_ms:.2f} total_ms={total_ms:.2f}", - project_id=self.project_id, - retrieval_mode="vector", - query_length=len(query_text), - candidate_limit=candidate_limit, - vector_row_count=vector_row_count, - embed_ms=embed_ms, - vector_query_ms=vector_query_ms, - hydrate_ms=hydrate_ms, - total_ms=total_ms, - ) - - if not vector_rows: - _log_vector_summary() - return [] - - hydrate_start = time.perf_counter() - # Build per-search_index_row similarity scores from chunk-level results. - # Each chunk_key encodes the search_index row type and id; keep both as the - # key because different row types can share the same numeric id (#982). - # Track the best similarity per row (for ranking) and all chunks (for context). - similarity_by_si_key: dict[SearchIndexKey, float] = {} - chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {} - for row in vector_rows: - chunk_key = row.get("chunk_key", "") - if "best_similarity" in row: - similarity = float(row["best_similarity"]) - else: - # Compatibility: private test doubles may still return native distance. - distance = float(row["best_distance"]) - similarity = self._distance_to_similarity(distance) - chunk_text = row.get("chunk_text", "") - try: - si_key = self._parse_chunk_key(chunk_key) - except (ValueError, IndexError): - # Fallback: group by entity_id for chunks without parseable keys - continue - current = similarity_by_si_key.get(si_key) - if current is None or similarity > current: - similarity_by_si_key[si_key] = similarity - chunks_by_si_key.setdefault(si_key, []).append((similarity, chunk_text)) - - if not similarity_by_si_key: - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - _log_vector_summary() - return [] - - # Filter out results below the minimum similarity threshold. - # Per-query min_similarity overrides the instance-level default. - effective_min_similarity = ( - min_similarity if min_similarity is not None else self._semantic_min_similarity - ) - if effective_min_similarity > 0.0: - if trace is not None: - threshold_rejections = tuple( - BelowThreshold(key=key, similarity=value, threshold=effective_min_similarity) - for key, value in similarity_by_si_key.items() - if value < effective_min_similarity - ) - trace.vector = build_vector_stage( - previous=trace.vector, - threshold_rejections=threshold_rejections, - ) - similarity_by_si_key = { - k: v for k, v in similarity_by_si_key.items() if v >= effective_min_similarity - } - if not similarity_by_si_key: - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - _log_vector_summary() - return [] - - # Fetch the actual search_index rows. Colliding (type, id) keys share one - # bare id, so deduplicate while preserving first-seen order. - si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_si_key)) - search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids) - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - missing_search_rows=tuple( - MissingSearchRow(key=key) - for key in similarity_by_si_key - if key not in search_index_rows - ), - ) - - # Apply optional filters if requested - filter_requested = any( - [ - permalink, - permalink_match, - title, - note_types, - after_date, - search_item_types, - categories, - metadata_filters, - file_path_prefix, - temporal, - ] - ) - - if filter_requested: - allowed_keys = await self._filter_candidate_keys( - list(search_index_rows), - 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, - ) - if trace is not None: - trace.vector = build_vector_stage( - previous=trace.vector, - filter_rejections=tuple( - FilteredOut(key=key) for key in search_index_rows if key not in allowed_keys - ), - ) - search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys} - - ranked_rows: list[SearchIndexRow] = [] - for si_key, similarity in similarity_by_si_key.items(): - row = search_index_rows.get(si_key) - if row is None: - continue - - # Small notes: return full content so the answer is always present. - # Large notes: return top-N most relevant chunks for richer context. - content_snippet = row.content_snippet or "" - if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT: - matched_chunk_text = content_snippet - else: - si_chunks = chunks_by_si_key.get(si_key, []) - si_chunks.sort(key=lambda c: c[0], reverse=True) - top_texts = [text for _, text in si_chunks[:TOP_CHUNKS_PER_RESULT]] - matched_chunk_text = "\n---\n".join(top_texts) if top_texts else None - - ranked_rows.append( - replace( - row, - score=similarity, - matched_chunk_text=matched_chunk_text, - ) - ) - - ranked_rows.sort(key=lambda item: item.score or 0.0, reverse=True) - hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 - # Rerank over the wide candidate pool, then slice to the page. Suppressed when - # hybrid calls this internally (_apply_rerank=False) — hybrid reranks its own - # fused result; _rerank_and_paginate no-ops back to a plain slice otherwise. - if _apply_rerank: - stable_rows = ranked_rows - if self._should_rerank(query_text): - stable_candidate_limit = self._rerank_candidate_limit() - if candidate_limit > stable_candidate_limit: - if trace is not None: - trace.stable_pool_refetched = True - stable_rows = await self._search_vector_only( - search_text=search_text, - 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, - min_similarity=min_similarity, - limit=stable_candidate_limit, - offset=0, - candidate_limit=stable_candidate_limit, - _emit_observability_log=False, - _apply_rerank=False, - trace=None, - ) - output = await self._rerank_and_paginate( - query_text, - ranked_rows, - offset=offset, - limit=limit, - stable_rows=stable_rows, - trace=trace, - ) - else: - output = ranked_rows[offset : offset + limit] - # Vector latency owns the optional rerank stage too. Logging before the - # awaited provider call hides the feature's dominant cost and can suppress - # the slow-query warning entirely. - _log_vector_summary() - return output - - @logfire.instrument("search.filter_candidates", extract_args=False) - 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 - - @logfire.instrument("search.fetch_candidate_rows", extract_args=False) - async def _fetch_search_index_rows_by_ids( - self, row_ids: list[int] - ) -> dict[SearchIndexKey, SearchIndexRow]: - """Fetch search_index rows by id, keyed by (type, id) to disambiguate types. - - A bare id can match one row per type (independent id sequences), so the - result must carry every matching row rather than letting one clobber another. - """ - if not row_ids: - return {} - placeholders = ",".join(f":id_{idx}" for idx in range(len(row_ids))) - params: dict[str, Any] = {f"id_{idx}": rid for idx, rid in enumerate(row_ids)} - scope_predicate = self.scope.predicate("project_id", params) - sql = f""" - SELECT - project_id, id, title, permalink, file_path, type, metadata, - from_id, to_id, relation_type, entity_id, content_snippet, - category, created_at, updated_at, 0 as score - FROM search_index - WHERE {scope_predicate} - AND id IN ({placeholders}) - """ - result: dict[SearchIndexKey, SearchIndexRow] = {} - async with db.scoped_session(self.session_maker) as session: - row_result = await session.execute(text(sql), params) - for row in row_result.fetchall(): - search_row = SearchIndexRow.from_mapping(row._asdict()) - result[(search_row.type, search_row.id)] = search_row - return result - - # ------------------------------------------------------------------ - # Shared semantic search: hybrid score-based fusion - # ------------------------------------------------------------------ - - async def _search_hybrid( - self, - *, - search_text: str, - 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], - min_similarity: Optional[float] = None, - limit: int, - offset: int, - _candidate_limit_override: int | None = None, - _apply_rerank: bool = True, - _emit_observability_log: bool = True, - trace: SearchTraceCollector | None = None, - ) -> List[SearchIndexRow]: - """Fuse FTS and vector results using score-based fusion. - - Uses the search_index (type, id) pair as the fusion key. The formula - ``max(vec, fts) + FUSION_BONUS * min(vec, fts)`` preserves - the dominant signal and rewards dual-source agreement. - """ - self._assert_semantic_available() - query_text = search_text.strip() - rerank_configured = self._should_rerank(query_text) - rerank_enabled = _apply_rerank and rerank_configured - query_start = time.perf_counter() - candidate_limit = ( - _candidate_limit_override - if _candidate_limit_override is not None - else self._candidate_limit(limit, offset, query_text) - ) - fts_start = time.perf_counter() - # allow_relaxed: question-form queries rarely AND-match, and a dead FTS - # branch silently degrades hybrid to vector-only ranking. Fusion plus - # bm25 keep relaxed lexical candidates from dominating precision. - with logfire.span("search.fts", candidate_limit=candidate_limit) as fts_span: - fts_results = await self.search( - search_text=search_text, - 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, - limit=candidate_limit, - offset=0, - allow_relaxed=True, - trace=trace, - ) - fts_span.set_attribute("result_count", len(fts_results)) - fts_ms = (time.perf_counter() - fts_start) * 1000 - vector_start = time.perf_counter() - vector_results = await self._search_vector_only( - search_text=search_text, - 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, - min_similarity=min_similarity, - limit=candidate_limit, - offset=0, - # Trigger: reranking owns a bounded candidate window shared by both legs. - # Why: the disabled path historically expands the vector leg again to - # preserve recall when many vector chunks collapse into a few search rows. - # Outcome: avoid double expansion only when reranking is actually active. - candidate_limit=candidate_limit if rerank_configured else None, - _emit_observability_log=False, - _apply_rerank=False, - trace=trace, - ) - vector_ms = (time.perf_counter() - vector_start) * 1000 - # Trigger: with reranking disabled the vector leg expands internally and can - # hydrate more rows than the fusion window it returns. - # Why: rows cut here never fuse — left in the trace they would surface as - # candidates with no rejection and no fused rank, which the response labels - # "returned". Rows with a recorded rejection keep their chunk evidence. - # Outcome: the trace keeps rows handed to fusion (or explicitly rejected); - # the cut shows up as served-chunk shrinkage in the candidate_window stage. - if trace is not None and trace.vector is not None: - kept_row_keys = {(row.type, row.id) for row in vector_results} - kept_row_keys.update( - rejection.key - for rejection_group in ( - trace.vector.threshold_rejections, - trace.vector.filter_rejections, - trace.vector.missing_search_rows, - ) - for rejection in rejection_group - ) - if any(match.key not in kept_row_keys for match in trace.vector.chunk_matches): - fused_chunks: dict[SearchIndexKey, list[tuple[str, float, int | None]]] = {} - for chunk_match in trace.vector.chunk_matches: - if chunk_match.key in kept_row_keys: - fused_chunks.setdefault(chunk_match.key, []).append( - (chunk_match.chunk_key, chunk_match.similarity, chunk_match.entity_id) - ) - trace.vector = build_vector_stage( - previous=trace.vector, - chunk_matches=fused_chunks, - ) - fusion_start = time.perf_counter() - - with logfire.span( - "search.fusion", fts_count=len(fts_results), vector_count=len(vector_results) - ) as fusion_span: - # --- Score-based fusion keyed on (type, id) --- - # A bare row id collides across row types (independent id sequences), so - # fusion must key on (type, id) or distinct rows would merge (#982). - # FTS scores are normalized to [0, 1] (BM25 is unbounded). - # Vector scores are used raw — already calibrated [0, 1] by _distance_to_similarity(). - rows_by_key: dict[SearchIndexKey, SearchIndexRow] = {} - - # Normalize FTS scores to [0, 1] — handles both SQLite (negative bm25) - # and Postgres (positive ts_rank) by using absolute values - fts_abs = [abs(row.score or 0.0) for row in fts_results] - fts_max = max(fts_abs) if fts_abs else 1.0 - - fts_scores: dict[SearchIndexKey, float] = {} - fts_ranks: dict[SearchIndexKey, int] = {} - for rank, row in enumerate(fts_results): - if row.id is None: - continue - row_key = (row.type, row.id) - norm = abs(row.score or 0.0) / fts_max if fts_max > 0 else 0.0 - # Gate: FTS scores below threshold contribute zero - if norm < FTS_GATE_THRESHOLD: - norm = 0.0 - fts_scores[row_key] = norm - fts_ranks.setdefault(row_key, rank) - rows_by_key[row_key] = row - - if trace is not None: - relaxed_fallback_used = ( - trace.fts.relaxed_fallback_used if trace.fts is not None else False - ) - trace.fts = build_fts_page_stage( - [((row.type, row.id), row.score or 0.0) for row in fts_results], - normalized_scores=fts_scores, - entity_ids={(row.type, row.id): row.entity_id for row in fts_results}, - fts_max_abs=fts_max, - relaxed_fallback_used=relaxed_fallback_used, - fts_ms=fts_ms, - ) - - vec_scores: dict[SearchIndexKey, float] = {} - vec_ranks: dict[SearchIndexKey, int] = {} - for rank, row in enumerate(vector_results): - if row.id is None: - continue - row_key = (row.type, row.id) - # Trigger: no re-normalization by vec_max - # Why: vector similarity is already calibrated [0, 1]; re-normalizing - # inflates weak matches when the entire result set is mediocre - vec_scores[row_key] = row.score or 0.0 - vec_ranks.setdefault(row_key, rank) - rows_by_key[row_key] = row - - # Fuse: max(v, f) + FUSION_BONUS * min(v, f) - # Preserves the dominant signal; bonus rewards dual-source agreement. - # Output range: [0, 1.3] for dual-source, [0, 1.0] for single-source. - fused_scores: dict[SearchIndexKey, float] = {} - for row_key in fts_scores.keys() | vec_scores.keys(): - v = vec_scores.get(row_key, 0.0) - f = fts_scores.get(row_key, 0.0) - fused_scores[row_key] = max(v, f) + FUSION_BONUS * min(v, f) - - ranked = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True) - fusion_span.set_attribute("result_count", len(ranked)) - fusion_ms = (time.perf_counter() - fusion_start) * 1000 - if trace is not None: - trace.fusion = build_fusion_stage( - formula_version=FUSION_FORMULA_VERSION, - bonus=FUSION_BONUS, - fts_scores=fts_scores, - fts_ranks=fts_ranks, - vector_scores=vec_scores, - vector_ranks=vec_ranks, - ranked_scores=ranked, - fusion_ms=fusion_ms, - ) - - def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: - row_key, fused_score = entry - row = rows_by_key[row_key] - # FTS-only hits use the bounded content preview and its truncation metadata. - # Copying the full note into matched_chunk bypasses that response bound. - return replace(row, score=fused_score) - - # Rerank the top fused candidates before paginating. When reranking is active - # we materialize the whole candidate list (cheap next to a cross-encoder call) - # and hand it to the shared paginate helper; the disabled path stays cheap by - # materializing only the requested page. - if rerank_enabled: - candidates = [_materialize(entry) for entry in ranked] - stable_candidates = candidates - stable_candidate_limit = self._rerank_candidate_limit() - if candidate_limit > stable_candidate_limit: - if trace is not None: - trace.stable_pool_refetched = True - stable_candidates = await self._search_hybrid( - search_text=search_text, - 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, - min_similarity=min_similarity, - limit=stable_candidate_limit, - offset=0, - _candidate_limit_override=stable_candidate_limit, - _apply_rerank=False, - _emit_observability_log=False, - trace=None, - ) - stable_keys = {(row.type, row.id) for row in stable_candidates} - expanded_tail = [entry for entry in ranked if entry[0] not in stable_keys] - - # Trigger: deeper pages expand the FTS/vector retrieval windows. - # Why: score fusion can strengthen an existing row when its second - # signal appears later, moving it across a page already returned. - # Outcome: freeze the fixed fused universe, then order newly admitted - # rows by their earliest source rank. That rank cannot improve after a - # row first appears, so each larger window only appends to the tail. - expanded_tail.sort( - key=lambda entry: ( - min( - fts_ranks.get(entry[0], candidate_limit), - vec_ranks.get(entry[0], candidate_limit), - ), - entry[0], - ) - ) - candidates = stable_candidates + [_materialize(entry) for entry in expanded_tail] - output = await self._rerank_and_paginate( - query_text, - candidates, - offset=offset, - limit=limit, - stable_rows=stable_candidates, - trace=trace, - ) - else: - output = [_materialize(entry) for entry in ranked[offset : offset + limit]] - total_ms = (time.perf_counter() - query_start) * 1000 - if _emit_observability_log and total_ms > 2500: - logger.warning( - "[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} " - "retrieval_mode={retrieval_mode} query_length={query_length} " - "candidate_limit={candidate_limit} fts_count={fts_count} " - "vector_count={vector_count} fts_ms={fts_ms:.2f} vector_ms={vector_ms:.2f} " - "fusion_ms={fusion_ms:.2f} total_ms={total_ms:.2f}", - project_id=self.project_id, - retrieval_mode="hybrid", - query_length=len(query_text), - candidate_limit=candidate_limit, - fts_count=len(fts_results), - vector_count=len(vector_results), - fts_ms=fts_ms, - vector_ms=vector_ms, - fusion_ms=fusion_ms, - total_ms=total_ms, - ) - return output diff --git a/src/basic_memory/repository/sqlite_search_repository.py b/src/basic_memory/repository/sqlite_search_repository.py index aa75461e7..0861023b7 100644 --- a/src/basic_memory/repository/sqlite_search_repository.py +++ b/src/basic_memory/repository/sqlite_search_repository.py @@ -407,15 +407,6 @@ async def drop_vector_tables(self) -> None: await session.commit() self._vector_tables_initialized = False - @override - def _distance_to_similarity(self, distance: float) -> float: - """Convert L2 distance to cosine similarity for normalized embeddings. - - sqlite-vec vec0 returns Euclidean (L2) distance by default. - For unit-normalized vectors: L2² = 2·(1 - cos_sim), so cos_sim = 1 - L2²/2. - """ - return max(0.0, 1.0 - (distance * distance) / 2.0) - @asynccontextmanager @override async def _prepare_entity_write_scope(self): diff --git a/src/basic_memory/services/project_readiness.py b/src/basic_memory/services/project_readiness.py index ed7bc1ca8..006bf0899 100644 --- a/src/basic_memory/services/project_readiness.py +++ b/src/basic_memory/services/project_readiness.py @@ -24,7 +24,7 @@ from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, ) -from basic_memory.repository.search_repository_base import current_vector_manifest_predicate +from basic_memory.repository.search_reader import current_vector_manifest_predicate from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_vector_index_factory import ( resolve_semantic_vector_index_name, diff --git a/src/basic_memory/services/retrieval_inspect.py b/src/basic_memory/services/retrieval_inspect.py index b57bb771e..e70c17889 100644 --- a/src/basic_memory/services/retrieval_inspect.py +++ b/src/basic_memory/services/retrieval_inspect.py @@ -16,11 +16,8 @@ from basic_memory.repository.note_content_repository import NoteContentRepository from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_repository import SearchRepository -from basic_memory.repository.search_repository_base import ( - ChunkManifestRow, - FUSION_FORMULA_VERSION, - SearchRepositoryBase, -) +from basic_memory.repository.search_reader import FUSION_FORMULA_VERSION, parse_chunk_key +from basic_memory.repository.search_repository_base import ChunkManifestRow from basic_memory.repository.search_trace import ( FinalResultEntry, QueryMeta, @@ -416,7 +413,7 @@ async def inspect_entity_chunks( inspected_chunks: list[InspectedChunk] = [] chunks_by_search_row: dict[tuple[str, int], list[InspectedChunk]] = {} for stored_row in stored_rows: - row_key = SearchRepositoryBase._parse_chunk_key(stored_row.chunk_key) + row_key = parse_chunk_key(stored_row.chunk_key) ordinal = int(stored_row.chunk_key.split(":")[2]) inspected_chunk = InspectedChunk( stored_row=stored_row, diff --git a/test-int/semantic/test_multilingual_benchmark_contract.py b/test-int/semantic/test_multilingual_benchmark_contract.py index 435562f03..461ddafa2 100644 --- a/test-int/semantic/test_multilingual_benchmark_contract.py +++ b/test-int/semantic/test_multilingual_benchmark_contract.py @@ -18,7 +18,7 @@ create_sqlite_search_vector_embeddings, ) from basic_memory.repository.semantic_chunking import split_text_into_chunks -from basic_memory.repository.search_repository_base import SMALL_NOTE_CONTENT_LIMIT +from basic_memory.repository.search_reader import SMALL_NOTE_CONTENT_LIMIT from basic_memory.schemas.search import SearchRetrievalMode from semantic.multilingual_benchmark import ( diff --git a/test-int/semantic/test_search_diagnostics.py b/test-int/semantic/test_search_diagnostics.py index 3179ed9d1..e7583fad7 100644 --- a/test-int/semantic/test_search_diagnostics.py +++ b/test-int/semantic/test_search_diagnostics.py @@ -340,20 +340,15 @@ async def test_similarity_formula_analysis(sqlite_engine_factory, tmp_path): from basic_memory import db as bm_db repo = cast(Any, service.repository) + semantic = repo._semantic_search() async with bm_db.scoped_session(repo.session_maker) as session: - await repo._prepare_vector_session(session) - vector_rows = await repo._run_vector_query( - session, - query_embedding, - candidate_limit=20, - ) + vector_rows = await semantic._run_vector_query(session, query_embedding, 20) print(f"\nQuery: '{query_text}'") print(f" {'chunk_key':<40} {'similarity':>12}") - for row in vector_rows[:10]: - similarity = float(row["best_similarity"]) - assert 0.0 <= similarity <= 1.0 - print(f" {row['chunk_key']:<40} {similarity:>12.4f}") + for chunk in vector_rows[:10]: + assert 0.0 <= chunk.similarity <= 1.0 + print(f" {chunk.chunk_key:<40} {chunk.similarity:>12.4f}") # --- Test: min_similarity threshold effectiveness --- diff --git a/test-int/semantic/test_semantic_coverage.py b/test-int/semantic/test_semantic_coverage.py index 0e4881089..d71faaca0 100644 --- a/test-int/semantic/test_semantic_coverage.py +++ b/test-int/semantic/test_semantic_coverage.py @@ -3,7 +3,7 @@ Exercises the uncovered code paths in PostgresSearchRepository: - _ensure_vector_tables (lines 258-352): pgvector extension, table creation, dimension mismatch detection -- _run_vector_query (lines 389-429): vector similarity query with cosine distance +- SemanticSearch._run_vector_query: vector similarity query with cosine distance - _write_embeddings (lines 431-458): embedding upsert into pgvector table - Metadata filters in FTS search (lines 682-745): JSONB filter operators (eq, in, contains, gt/gte/lt/lte, between) @@ -19,6 +19,7 @@ from basic_memory import db from basic_memory.config import DatabaseBackend +from basic_memory.repository.search_reader import HydratedChunk, SemanticSearch from basic_memory.schemas.search import SearchItemType, SearchQuery, SearchRetrievalMode from semantic.conftest import ( @@ -104,7 +105,7 @@ async def test_postgres_vector_table_setup_and_query(postgres_engine_factory, tm async def test_postgres_hybrid_search(postgres_engine_factory, tmp_path): """Exercise the hybrid (score-based fusion) code path on Postgres. - This covers the full _search_hybrid path including both FTS and vector + This covers the full SemanticSearch.hybrid path including both FTS and vector retrieval with score-based fusion. """ skip_if_needed(PG_FASTEMBED) @@ -118,7 +119,7 @@ async def test_postgres_hybrid_search(postgres_engine_factory, tmp_path): await seed_benchmark_notes(search_service, note_count=20) - # Hybrid search — exercises _search_hybrid score-based fusion + # Hybrid search — exercises SemanticSearch.hybrid score-based fusion results = await search_service.search( SearchQuery( text="database migration schema", @@ -158,17 +159,20 @@ async def test_postgres_hybrid_preserves_candidate_windows( repo._reranker_candidates = 100 candidate_limits: list[int] = [] - run_vector_query = repo._run_vector_query + run_vector_query = SemanticSearch._run_vector_query async def record_vector_query( + self: SemanticSearch, session: Any, query_embedding: list[float], candidate_limit: int, - ) -> list[dict[str, Any]]: + *, + trace: Any = None, + ) -> list[HydratedChunk]: candidate_limits.append(candidate_limit) - return await run_vector_query(session, query_embedding, candidate_limit) + return await run_vector_query(self, session, query_embedding, candidate_limit, trace=trace) - monkeypatch.setattr(repo, "_run_vector_query", record_vector_query) + monkeypatch.setattr(SemanticSearch, "_run_vector_query", record_vector_query) baseline_results = await search_service.search( SearchQuery( diff --git a/tests/repository/test_distance_to_similarity.py b/tests/repository/test_distance_to_similarity.py deleted file mode 100644 index faa8a27c0..000000000 --- a/tests/repository/test_distance_to_similarity.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Unit tests for backend-specific distance-to-similarity conversions.""" - -import pytest - -from basic_memory.repository.postgres_search_repository import PostgresSearchRepository -from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository - - -def test_sqlite_distance_to_similarity_formula(): - """SQLite converts L2 distance to cosine similarity for normalized vectors.""" - repo = SQLiteSearchRepository.__new__(SQLiteSearchRepository) - - assert repo._distance_to_similarity(0.0) == 1.0 - assert repo._distance_to_similarity(1.0) == pytest.approx(0.5) - assert repo._distance_to_similarity(2.0) == 0.0 - - -def test_postgres_distance_to_similarity_formula(): - """Postgres converts pgvector cosine distance to cosine similarity.""" - repo = PostgresSearchRepository.__new__(PostgresSearchRepository) - - assert repo._distance_to_similarity(0.0) == 1.0 - assert repo._distance_to_similarity(1.0) == 0.0 - assert repo._distance_to_similarity(2.0) == 0.0 diff --git a/tests/repository/test_hybrid_fusion.py b/tests/repository/test_hybrid_fusion.py index 0ad16896b..dd5e14dc4 100644 --- a/tests/repository/test_hybrid_fusion.py +++ b/tests/repository/test_hybrid_fusion.py @@ -6,26 +6,26 @@ 3. Produces zero fused score when the source score is zero """ -from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.repository.search_scope import ProjectScope -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass -from datetime import datetime -from typing import override, Any, Optional, cast +from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest +from sqlalchemy.ext.asyncio import AsyncSession 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 ( +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( FUSION_BONUS, - SearchIndexKey, - SearchRepositoryBase, + SemanticSearch, + VectorRetrieval, ) +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.search_trace import SearchTraceCollector -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex +from basic_memory.schemas.search import SearchRetrievalMode @dataclass @@ -51,159 +51,108 @@ class FakeRow: matched_chunk_text: str | None = None -class ConcreteSearchRepo(SearchRepositoryBase): - """Minimal concrete subclass for testing hybrid fusion logic.""" +class FakeFts: + """An ``FtsBackend`` that answers every pass with fixed rows and records what it was asked. - def __init__(self): - self._semantic_enabled = True - self._semantic_vector_k = 100 - self._semantic_min_similarity = 0.0 - # _search_hybrid calls _assert_semantic_available which checks this - self._embedding_provider = _fake_embedding_provider() - self._vector_dimensions = 384 - self._vector_tables_initialized = True - self.session_maker = None - self.project_id = 1 - self.scope = ProjectScope.single(1) - - @override - async def init_search_index(self): - pass # pragma: no cover + ``rows`` is the answer, or a function of the requested ``limit`` when a test needs + the lexical leg to widen with the candidate window. + """ - @override - async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: - return None # physical storage is not inspectable in this double + def __init__(self, rows: Sequence[Any] | Callable[[int], Sequence[Any]] = ()) -> None: + if isinstance(rows, Sequence): + fixed = list(rows) + self.answer: Callable[[int], Sequence[Any]] = lambda _limit: fixed + else: + self.answer = rows + self.queries: list[PreparedSearchQuery] = [] + self.calls: list[dict[str, Any]] = [] - @override async def search( self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[list[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[list[SearchItemType]] = None, - categories: Optional[list[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, + scope: ProjectScope, + query: PreparedSearchQuery, + *, + limit: int, + offset: int, allow_relaxed: bool = False, session: AsyncSession | None = None, - *, candidate_keys: Sequence[SearchIndexKey] | None = None, trace: SearchTraceCollector | None = None, ) -> list[SearchIndexRow]: - return [] # pragma: no cover - - @override - async def _ensure_vector_tables(self): - pass # pragma: no cover - - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] # pragma: no cover - - @override - async def _write_embeddings(self, session, jobs, embeddings): - pass # pragma: no cover - - @override - async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): - return [] # pragma: no cover - - @override - async def _delete_stale_chunks( + self.queries.append(query) + self.calls.append( + { + "limit": limit, + "offset": offset, + "allow_relaxed": allow_relaxed, + "candidate_keys": candidate_keys, + } + ) + return cast(list[SearchIndexRow], list(self.answer(limit))) + + async def count( self, - session, - stale_ids, - entity_id, + scope: ProjectScope, + query: PreparedSearchQuery, *, - expected_deletions=None, - ): - return [] # pragma: no cover + allow_relaxed: bool = False, + ) -> int: + return len(self.answer(0)) + + +def fake_vector_retrieval( + *, + vector_k: int = 100, + min_similarity: float = 0.0, + embed_query: AsyncMock | None = None, +) -> VectorRetrieval: + """A semantic stack whose adapter is never consulted: tests stub the neighbour stage.""" + provider = type( + "EP", + (), + { + "model_name": "fake", + "dimensions": 384, + "embed_query": embed_query or AsyncMock(return_value=[0.0] * 384), + "embed_documents": AsyncMock(return_value=[]), + "runtime_log_attrs": lambda self: {}, + }, + )() + return VectorRetrieval( + index=cast(SemanticVectorIndex, object()), + index_name="sqlite-vec", + embedding_provider=cast(EmbeddingProvider, provider), + embedding_model="fake:384", + vector_k=vector_k, + min_similarity=min_similarity, + ) - async def _update_timestamp_sql(self): - return "CURRENT_TIMESTAMP" # pragma: no cover - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) # pragma: no cover +HYBRID_QUERY = PreparedSearchQuery(search_text="test", retrieval_mode=SearchRetrievalMode.HYBRID) -def _fake_embedding_provider() -> EmbeddingProvider: - return cast( - EmbeddingProvider, - type( - "EP", - (), - { - "model_name": "fake", - "dimensions": 384, - "embed_query": AsyncMock(return_value=[0.0] * 384), - "embed_documents": AsyncMock(return_value=[]), - "runtime_log_attrs": lambda self: {}, - }, - )(), +async def fuse( + fts_results: list[Any], + vector_results: list[Any], + *, + query: PreparedSearchQuery = HYBRID_QUERY, +) -> list[SearchIndexRow]: + """Run hybrid with both legs answering fixed rows, so only fusion is under test.""" + semantic = SemanticSearch( + cast(Any, None), ProjectScope.single(1), FakeFts(fts_results), fake_vector_retrieval() ) - - -HYBRID_KWARGS: dict[str, Any] = dict( - search_text="test", - 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=None, - temporal=None, - limit=10, - offset=0, -) + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=vector_results): + return await semantic.hybrid(query, limit=10, offset=0) @pytest.mark.asyncio async def test_high_fts_score_boosts_ranking(): """FTS-only: a high normalized score should outscore a low normalized score.""" - repo = ConcreteSearchRepo() - - # Two FTS results with very different scores high_score_row = FakeRow(id=1, score=10.0, title="high") low_score_row = FakeRow(id=2, score=0.5, title="low") - fts_results = [high_score_row, low_score_row] # No vector results — isolate FTS weighting behavior - vector_results = [] - - with ( - patch.object( - repo, - "search", - new_callable=AsyncMock, - return_value=fts_results, - ), - patch.object( - repo, - "_search_vector_only", - new_callable=AsyncMock, - return_value=vector_results, - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([high_score_row, low_score_row], []) assert len(results) == 2 # After normalization: id=1 → 1.0, id=2 → 0.05 @@ -215,8 +164,6 @@ async def test_high_fts_score_boosts_ranking(): @pytest.mark.asyncio async def test_dual_source_ranks_higher_than_single(): """A result in both FTS and vector should rank above single-source results.""" - repo = ConcreteSearchRepo() - # Row 1 in both (fts=5.0→norm 1.0, vec=0.9), Row 2 FTS-only (fts=5.0→norm 1.0), # Row 3 vec-only (0.8) fts_results = [ @@ -228,13 +175,7 @@ async def test_dual_source_ranks_higher_than_single(): FakeRow(id=3, score=0.8, title="vec-only"), ] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse(fts_results, vector_results) result_ids = [r.id for r in results] # Row 1 (dual-source) should rank first, then Row 2 (FTS 1.0), then Row 3 (vec 0.8) @@ -251,19 +192,7 @@ async def test_dual_source_ranks_higher_than_single(): @pytest.mark.asyncio async def test_zero_score_produces_zero_fused(): """A zero-score FTS result with no vector match produces a zero fused score.""" - repo = ConcreteSearchRepo() - - # FTS result with score 0.0 - fts_results = [FakeRow(id=1, score=0.0, title="zero-score")] - vector_results = [] - - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([FakeRow(id=1, score=0.0, title="zero-score")], []) assert len(results) == 1 # Zero FTS score, no vector → fused = max(0, 0) + 0.3 * min(0, 0) = 0.0 @@ -277,18 +206,10 @@ async def test_cross_type_id_collision_keeps_both_results(): search_index row types have independent id sequences, so fusing on a bare row id merged unrelated rows into one result and dropped the other. """ - repo = ConcreteSearchRepo() - fts_results = [FakeRow(id=1, type="entity", score=5.0, title="entity-row")] vector_results = [FakeRow(id=1, type="relation", score=0.8, title="relation-row")] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse(fts_results, vector_results) assert {(r.type, r.id) for r in results} == {("entity", 1), ("relation", 1)} # Single-source scores must not earn the dual-source fusion bonus across types. @@ -301,21 +222,9 @@ async def test_cross_type_id_collision_keeps_both_results(): @pytest.mark.asyncio async def test_fts_only_result_does_not_copy_content_into_matched_chunk(): """FTS-only hits use the API content preview instead of a second full-note field.""" - repo = ConcreteSearchRepo() - content = "This is the full note content with the answer we need to find." - fts_results = [ - FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=content), - ] - vector_results = [] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=content)], []) assert len(results) == 1 assert results[0].matched_chunk_text is None @@ -325,20 +234,7 @@ async def test_fts_only_result_does_not_copy_content_into_matched_chunk(): @pytest.mark.asyncio async def test_fts_only_result_with_null_content_keeps_null_matched_chunk(): """FTS-only results with no content_snippet should keep matched_chunk_text as None.""" - repo = ConcreteSearchRepo() - - fts_results = [ - FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=None), - ] - vector_results = [] - - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse([FakeRow(id=1, score=5.0, title="fts-hit", content_snippet=None)], []) assert len(results) == 1 assert results[0].matched_chunk_text is None @@ -347,8 +243,6 @@ async def test_fts_only_result_with_null_content_keeps_null_matched_chunk(): @pytest.mark.asyncio async def test_dual_source_result_keeps_vector_matched_chunk(): """Dual-source results should keep matched_chunk_text from vector search, not overwrite.""" - repo = ConcreteSearchRepo() - content = "Full note content from FTS." vector_chunk = "Specific chunk matched by vector search." fts_results = [ @@ -364,14 +258,22 @@ async def test_dual_source_result_keeps_vector_matched_chunk(): ), ] - with ( - patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results), - patch.object( - repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results - ), - ): - results = await repo._search_hybrid(**HYBRID_KWARGS) + results = await fuse(fts_results, vector_results) assert len(results) == 1 - # Vector result overwrites the FTS row in rows_by_id, so matched_chunk_text is preserved + # Vector result overwrites the FTS row in rows_by_key, so matched_chunk_text is preserved assert results[0].matched_chunk_text == vector_chunk + + +@pytest.mark.asyncio +async def test_hybrid_fts_leg_runs_in_fts_mode_with_relaxation(): + """The lexical leg is the same prepared query in FTS mode, allowed to relax.""" + fts = FakeFts([FakeRow(id=1, score=5.0)]) + semantic = SemanticSearch(cast(Any, None), ProjectScope.single(1), fts, fake_vector_retrieval()) + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + assert [query.retrieval_mode for query in fts.queries] == [SearchRetrievalMode.FTS] + assert fts.queries[0].search_text == "test" + assert fts.calls[0]["allow_relaxed"] is True diff --git a/tests/repository/test_postgres_search_repository_unit.py b/tests/repository/test_postgres_search_repository_unit.py index 7ab175058..30a468e28 100644 --- a/tests/repository/test_postgres_search_repository_unit.py +++ b/tests/repository/test_postgres_search_repository_unit.py @@ -262,7 +262,7 @@ async def test_returns_empty_for_empty_embedding(self): embedding_provider=StubEmbeddingProvider(), ) session = AsyncMock() - result = await repo._run_vector_query(session, [], 10) + result = await repo._semantic_search()._run_vector_query(session, [], 10) assert result == [] diff --git a/tests/repository/test_rerank_pipeline.py b/tests/repository/test_rerank_pipeline.py index 97ca99783..3b1596325 100644 --- a/tests/repository/test_rerank_pipeline.py +++ b/tests/repository/test_rerank_pipeline.py @@ -1,7 +1,7 @@ """Rerank stage wiring in the shared search pipeline (vector + hybrid).""" from datetime import datetime, timezone -from typing import Any +from typing import Any, cast from unittest.mock import MagicMock import pytest @@ -19,16 +19,28 @@ demote_tail_scores, validate_rerank_scores, ) +from basic_memory.repository.search_filters import FtsBackend +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + RERANK_POOL_CHUNK_FANOUT, + HydratedChunk, + Reranking, + SemanticSearch, + VectorRetrieval, + rerank_document_text, +) from basic_memory.repository.search_repository import create_search_repository -from basic_memory.repository.search_repository_base import RERANK_POOL_CHUNK_FANOUT +from basic_memory.repository.search_scope import ProjectScope from basic_memory.repository.semantic_errors import ( RerankProviderContractError, RerankTransientError, SemanticDependenciesMissingError, ) +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode from basic_memory.services.entity_service import EntityService +from tests.repository.test_hybrid_fusion import FakeFts type BackendSearchRepository = SQLiteSearchRepository | PostgresSearchRepository @@ -191,69 +203,63 @@ def _row(**overrides) -> SearchIndexRow: return SearchIndexRow(**base) -def _unit_repo() -> SQLiteSearchRepository: - """A repo built without a real DB — for the pure rerank helper methods.""" - config = BasicMemoryConfig( - env="test", - projects={"test-project": "/tmp/test"}, - default_project="test-project", - database_backend=DatabaseBackend.SQLITE, - semantic_search_enabled=True, +def _reranking(provider: Any, *, candidates: int = 20, max_document_chars: int = 0) -> Reranking: + return Reranking( + provider=provider, candidates=candidates, max_document_chars=max_document_chars ) - return SQLiteSearchRepository( - MagicMock(), - project_id=1, - app_config=config, + + +def _semantic( + *, + vector_k: int = 100, + rerank: Reranking | None = None, + fts: FtsBackend | None = None, +) -> SemanticSearch: + """A pipeline built without a real DB, for the pure rerank helpers.""" + vector = VectorRetrieval( + index=cast(SemanticVectorIndex, object()), + index_name="sqlite-vec", embedding_provider=_StubEmbeddingProvider(), + embedding_model="stub:4", + vector_k=vector_k, + min_similarity=0.0, ) + return SemanticSearch(MagicMock(), ProjectScope.single(1), fts or FakeFts(), vector, rerank) # --- Pure helper behavior --- -def test_should_rerank_gating(): - repo = _unit_repo() - repo._rerank_provider = None - assert repo._should_rerank("auth") is False - repo._rerank_provider = _FakeReranker({}) - assert repo._should_rerank("") is False - assert repo._should_rerank("auth") is True +def test_rerank_is_active_only_with_a_provider_and_a_query(): + assert _semantic()._active_rerank("auth") is None + semantic = _semantic(rerank=_reranking(_FakeReranker({}))) + assert semantic._active_rerank("") is None + assert semantic._active_rerank("auth") is semantic.rerank def test_rerank_document_text_fallbacks(): - repo = _unit_repo() - assert repo._rerank_document_text(_row(title="T", matched_chunk_text="chunk")) == "chunk\nT" + assert rerank_document_text(_row(title="T", matched_chunk_text="chunk"), 0) == "chunk\nT" assert ( - repo._rerank_document_text(_row(title="T", matched_chunk_text=None, content_snippet="snip")) + rerank_document_text(_row(title="T", matched_chunk_text=None, content_snippet="snip"), 0) == "snip\nT" ) - assert ( - repo._rerank_document_text(_row(title=None, matched_chunk_text="only-body")) == "only-body" - ) - assert ( - repo._rerank_document_text(_row(title="only-title", content_snippet=None)) == "only-title" - ) - assert repo._rerank_document_text(_row(title=None, content_snippet=None)) == "" + assert rerank_document_text(_row(title=None, matched_chunk_text="only-body"), 0) == "only-body" + assert rerank_document_text(_row(title="only-title", content_snippet=None), 0) == "only-title" + assert rerank_document_text(_row(title=None, content_snippet=None), 0) == "" def test_rerank_document_text_truncation(): - repo = _unit_repo() row = _row(title="T", matched_chunk_text="x" * 500) # full text = 500 + "\nT" = 502 chars - repo._reranker_max_document_chars = 0 # disabled - assert len(repo._rerank_document_text(row)) == 502 - repo._reranker_max_document_chars = 100 # trims to the leading (most-relevant) text - trimmed = repo._rerank_document_text(row) + assert len(rerank_document_text(row, 0)) == 502 # disabled + trimmed = rerank_document_text(row, 100) # trims to the leading (most-relevant) text assert len(trimmed) == 100 and trimmed == "x" * 100 - repo._reranker_max_document_chars = 10_000 # no-op when already under the cap - assert len(repo._rerank_document_text(row)) == 502 + assert len(rerank_document_text(row, 10_000)) == 502 # no-op when already under the cap def test_rerank_document_text_cap_preserves_matched_body_with_long_title(): - repo = _unit_repo() - repo._reranker_max_document_chars = 8 row = _row(title="title-" * 20, matched_chunk_text="MATCHED body") - assert "MATCHED" in repo._rerank_document_text(row) + assert "MATCHED" in rerank_document_text(row, 8) def test_demoted_tail_scores_are_stable_as_the_tail_grows(): @@ -268,66 +274,58 @@ def test_demoted_tail_scores_are_stable_as_the_tail_grows(): def test_candidate_limit_over_fetches_chunks_for_rerank_pool(): """With reranking active, over-fetch chunks so dedup can't starve the rerank window.""" - repo = _unit_repo() - repo._semantic_vector_k = 5 - repo._reranker_candidates = 20 + plain = _semantic(vector_k=5) + assert plain._candidate_limit(limit=1, offset=0, query_text="auth") == 10 # max(5, 10) - repo._rerank_provider = None - assert repo._candidate_limit(limit=1, offset=0, query_text="auth") == 10 # max(5, 10) - - repo._rerank_provider = _FakeReranker({}) + reranked = _semantic(vector_k=5, rerank=_reranking(_FakeReranker({}), candidates=20)) assert ( - repo._candidate_limit(limit=1, offset=0, query_text="auth") == 20 * RERANK_POOL_CHUNK_FANOUT + reranked._candidate_limit(limit=1, offset=0, query_text="auth") + == 20 * RERANK_POOL_CHUNK_FANOUT ) - assert repo._candidate_limit(limit=1, offset=0, query_text="") == 10 # no query → no bump + assert reranked._candidate_limit(limit=1, offset=0, query_text="") == 10 # no query → no bump def test_candidate_limit_expands_only_for_results_beyond_rerank_pool(): """The fixed rerank window grows only enough to supply the requested tail.""" - repo = _unit_repo() - repo._semantic_vector_k = 5 - repo._reranker_candidates = 20 - repo._rerank_provider = _FakeReranker({}) + semantic = _semantic(vector_k=5, rerank=_reranking(_FakeReranker({}), candidates=20)) - first_page_limit = repo._candidate_limit(limit=10, offset=0, query_text="auth") + first_page_limit = semantic._candidate_limit(limit=10, offset=0, query_text="auth") assert first_page_limit == 20 * RERANK_POOL_CHUNK_FANOUT - assert repo._candidate_limit(limit=20, offset=0, query_text="auth") == first_page_limit - assert repo._candidate_limit(limit=10, offset=10, query_text="auth") == first_page_limit - assert repo._candidate_limit(limit=21, offset=0, query_text="auth") == 90 - assert repo._candidate_limit(limit=10, offset=19, query_text="auth") == 170 - assert repo._candidate_limit(limit=10, offset=20, query_text="auth") == 180 + assert semantic._candidate_limit(limit=20, offset=0, query_text="auth") == first_page_limit + assert semantic._candidate_limit(limit=10, offset=10, query_text="auth") == first_page_limit + assert semantic._candidate_limit(limit=21, offset=0, query_text="auth") == 90 + assert semantic._candidate_limit(limit=10, offset=19, query_text="auth") == 170 + assert semantic._candidate_limit(limit=10, offset=20, query_text="auth") == 180 # Large first pages still retrieve their untouched tail and pagination probe. - assert repo._candidate_limit(limit=101, offset=0, query_text="auth") == 890 + assert semantic._candidate_limit(limit=101, offset=0, query_text="auth") == 890 @pytest.mark.asyncio async def test_rerank_paginate_noop_paths(): - repo = _unit_repo() rows = [_row(id=1), _row(id=2)] - repo._rerank_provider = None - assert await repo._rerank_and_paginate("auth", rows, offset=0, limit=10) == rows + assert await _semantic()._rerank_and_paginate("auth", rows, offset=0, limit=10) == rows reranker = _FakeReranker({}) - repo._rerank_provider = reranker - assert await repo._rerank_and_paginate("", rows, offset=0, limit=10) == rows - assert await repo._rerank_and_paginate("auth", [], offset=0, limit=10) == [] - assert await repo._rerank_and_paginate("auth", rows, offset=2, limit=10) == [] + semantic = _semantic(rerank=_reranking(reranker)) + assert await semantic._rerank_and_paginate("", rows, offset=0, limit=10) == rows + assert await semantic._rerank_and_paginate("auth", [], offset=0, limit=10) == [] + assert await semantic._rerank_and_paginate("auth", rows, offset=2, limit=10) == [] assert reranker.calls == 0 @pytest.mark.asyncio async def test_rerank_paginate_reorders_rescore_and_demotes_tail(): - repo = _unit_repo() - repo._rerank_provider = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9, "Charlie": 0.5}) - repo._reranker_candidates = 2 + semantic = _semantic( + rerank=_reranking(_FakeReranker({"Alpha": 0.1, "Bravo": 0.9, "Charlie": 0.5}), candidates=2) + ) rows = [ _row(id=1, title="Alpha"), _row(id=2, title="Bravo"), _row(id=3, title="Charlie"), # past the pool ] - result = await repo._rerank_and_paginate("auth", rows, offset=0, limit=3) + result = await semantic._rerank_and_paginate("auth", rows, offset=0, limit=3) assert [r.title for r in result] == ["Bravo", "Alpha", "Charlie"] assert result[0].score == 0.9 # reranker relevance replaces the prior score @@ -343,16 +341,16 @@ async def test_rerank_paginate_reorders_rescore_and_demotes_tail(): @pytest.mark.asyncio async def test_rerank_paginate_preserves_pool_before_tail_at_zero_floor(): - repo = _unit_repo() - repo._rerank_provider = _FakeReranker({"Alpha": 0.0, "Bravo": 0.9}) - repo._reranker_candidates = 2 + semantic = _semantic( + rerank=_reranking(_FakeReranker({"Alpha": 0.0, "Bravo": 0.9}), candidates=2) + ) rows = [ _row(id=1, title="Alpha"), _row(id=2, title="Bravo"), _row(id=3, title="Charlie"), ] - result = await repo._rerank_and_paginate("auth", rows, offset=0, limit=3) + result = await semantic._rerank_and_paginate("auth", rows, offset=0, limit=3) assert [row.title for row in result] == ["Bravo", "Alpha", "Charlie"] assert [row.score for row in result] == [0.9, 0.0, 0.0] @@ -361,17 +359,15 @@ async def test_rerank_paginate_preserves_pool_before_tail_at_zero_floor(): @pytest.mark.asyncio async def test_rerank_paginate_scores_singleton_prefix_and_demotes_tail(): """A one-row prefix still calibrates scores before cross-project merging.""" - repo = _unit_repo() reranker = _FakeReranker({"Only": 0.4}) - repo._rerank_provider = reranker - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(reranker, candidates=2)) stable_rows = [_row(id=1, title="Only", score=0.5)] expanded_rows = [ stable_rows[0], _row(id=2, title="Tail", score=1.3), ] - result = await repo._rerank_and_paginate( + result = await semantic._rerank_and_paginate( "auth", expanded_rows, offset=0, @@ -388,10 +384,8 @@ async def test_rerank_paginate_scores_singleton_prefix_and_demotes_tail(): @pytest.mark.asyncio async def test_rerank_paginate_calibrates_tail_scores_on_deep_page(): """Deep pages rescore the fixed prefix before returning its calibrated tail.""" - repo = _unit_repo() reranker = _FakeReranker({"n1": 0.9, "n2": 0.8}) - repo._rerank_provider = reranker - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(reranker, candidates=2)) stable_rows = [_row(id=1, title="n1"), _row(id=2, title="n2")] expanded_rows = [ _row(id=3, title="newly strengthened"), @@ -401,7 +395,7 @@ async def test_rerank_paginate_calibrates_tail_scores_on_deep_page(): _row(id=5, title="n5"), ] - result = await repo._rerank_and_paginate( + result = await semantic._rerank_and_paginate( "auth", expanded_rows, offset=2, @@ -417,10 +411,8 @@ async def test_rerank_paginate_calibrates_tail_scores_on_deep_page(): @pytest.mark.asyncio async def test_rerank_paginate_keeps_expanded_candidates_out_of_stable_prefix(): """A larger tail retrieval cannot replace candidates in the reranked prefix.""" - repo = _unit_repo() reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9, "Charlie": 1.0}) - repo._rerank_provider = reranker - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(reranker, candidates=2)) stable_rows = [_row(id=1, title="Alpha"), _row(id=2, title="Bravo")] expanded_rows = [ _row(id=3, title="Charlie"), @@ -429,14 +421,14 @@ async def test_rerank_paginate_keeps_expanded_candidates_out_of_stable_prefix(): _row(id=4, title="Delta"), ] - result = await repo._rerank_and_paginate( + result = await semantic._rerank_and_paginate( "auth", expanded_rows, offset=0, limit=3, stable_rows=stable_rows, ) - expanded_result = await repo._rerank_and_paginate( + expanded_result = await semantic._rerank_and_paginate( "auth", expanded_rows + [_row(id=5, title="Echo"), _row(id=6, title="Foxtrot")], offset=0, @@ -453,38 +445,32 @@ async def test_rerank_paginate_keeps_expanded_candidates_out_of_stable_prefix(): @pytest.mark.asyncio async def test_rerank_paginate_surfaces_transient_provider_error(): """Transient failures must not silently replace reranked order with retrieval order.""" - repo = _unit_repo() - repo._rerank_provider = _ExplodingReranker() - repo._reranker_candidates = 20 + semantic = _semantic(rerank=_reranking(_ExplodingReranker(), candidates=20)) rows = [_row(id=1, title="A"), _row(id=2, title="B")] with pytest.raises(RerankTransientError, match="backend unreachable"): - await repo._rerank_and_paginate("auth", rows, offset=0, limit=10) + await semantic._rerank_and_paginate("auth", rows, offset=0, limit=10) @pytest.mark.asyncio async def test_rerank_paginate_does_not_duplicate_results_when_later_page_is_transient(): """A later page fails instead of changing order and repeating an earlier result.""" - repo = _unit_repo() - repo._rerank_provider = _SucceedsThenTransientReranker() - repo._reranker_candidates = 2 + semantic = _semantic(rerank=_reranking(_SucceedsThenTransientReranker(), candidates=2)) rows = [_row(id=1, title="A"), _row(id=2, title="B")] - first_page = await repo._rerank_and_paginate("auth", rows, offset=0, limit=1) + first_page = await semantic._rerank_and_paginate("auth", rows, offset=0, limit=1) assert [row.id for row in first_page] == [2] with pytest.raises(RerankTransientError, match="backend unreachable"): - await repo._rerank_and_paginate("auth", rows, offset=1, limit=1) + await semantic._rerank_and_paginate("auth", rows, offset=1, limit=1) @pytest.mark.asyncio async def test_rerank_paginate_misaligned_scores_raise(): """A length mismatch is a provider bug — fail fast, don't degrade.""" - repo = _unit_repo() - repo._rerank_provider = _BadReranker() - repo._reranker_candidates = 20 + semantic = _semantic(rerank=_reranking(_BadReranker(), candidates=20)) with pytest.raises(RerankProviderContractError, match="Reranker returned 0 scores"): - await repo._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) + await semantic._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) @pytest.mark.asyncio @@ -498,11 +484,9 @@ async def test_rerank_paginate_misaligned_scores_raise(): ) async def test_rerank_paginate_surfaces_permanent_faults(exc): """Permanent faults (contract break, missing deps) propagate — not silently degraded.""" - repo = _unit_repo() - repo._rerank_provider = _PermanentFaultReranker(exc) - repo._reranker_candidates = 20 + semantic = _semantic(rerank=_reranking(_PermanentFaultReranker(exc), candidates=20)) with pytest.raises(type(exc)): - await repo._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) + await semantic._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) # --- End-to-end through both repository backends --- @@ -601,17 +585,20 @@ async def test_vector_search_expands_tail_from_stable_rerank_pool( rerank_search_repository._rerank_provider = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) candidate_limits: list[int] = [] - run_vector_query = rerank_search_repository._run_vector_query + run_vector_query = SemanticSearch._run_vector_query async def record_vector_query( + self: SemanticSearch, session: Any, query_embedding: list[float], candidate_limit: int, - ) -> list[dict[str, Any]]: + *, + trace: Any = None, + ) -> list[HydratedChunk]: candidate_limits.append(candidate_limit) - return await run_vector_query(session, query_embedding, candidate_limit) + return await run_vector_query(self, session, query_embedding, candidate_limit, trace=trace) - monkeypatch.setattr(rerank_search_repository, "_run_vector_query", record_vector_query) + monkeypatch.setattr(SemanticSearch, "_run_vector_query", record_vector_query) results = await rerank_search_repository.search( search_text="auth session token", @@ -637,12 +624,12 @@ async def slow_rerank(query: str, documents: list[str]) -> list[float]: rerank_search_repository._rerank_provider = reranker monkeypatch.setattr(reranker, "rerank", slow_rerank) monkeypatch.setattr( - "basic_memory.repository.search_repository_base.time.perf_counter", + "basic_memory.repository.search_reader.time.perf_counter", lambda: clock["now"], ) warning = MagicMock() monkeypatch.setattr( - "basic_memory.repository.search_repository_base.logger.warning", + "basic_memory.repository.search_reader.logger.warning", warning, ) @@ -751,17 +738,20 @@ async def test_hybrid_search_preserves_candidate_windows( rerank_search_repository._rerank_provider = None candidate_limits: list[int] = [] - run_vector_query = rerank_search_repository._run_vector_query + run_vector_query = SemanticSearch._run_vector_query async def record_vector_query( + self: SemanticSearch, session: Any, query_embedding: list[float], candidate_limit: int, - ) -> list[dict[str, Any]]: + *, + trace: Any = None, + ) -> list[HydratedChunk]: candidate_limits.append(candidate_limit) - return await run_vector_query(session, query_embedding, candidate_limit) + return await run_vector_query(self, session, query_embedding, candidate_limit, trace=trace) - monkeypatch.setattr(rerank_search_repository, "_run_vector_query", record_vector_query) + monkeypatch.setattr(SemanticSearch, "_run_vector_query", record_vector_query) baseline_results = await rerank_search_repository.search( search_text="auth session token", @@ -801,17 +791,12 @@ async def record_vector_query( @pytest.mark.asyncio async def test_hybrid_search_keeps_deep_tail_stable_as_candidate_window_grows(monkeypatch): """Late dual-source evidence cannot move a row across an earlier tail page.""" - repo = _unit_repo() - repo._semantic_vector_k = 2 - repo._reranker_candidates = 2 reranker = _FakeReranker({"Alpha": 0.9, "Bravo": 0.8}) - repo._rerank_provider = reranker charlie = _row(id=3, title="Charlie") delta = _row(id=4, title="Delta") - async def fake_fts_search(*args, limit: int, **kwargs) -> list[SearchIndexRow]: - assert kwargs["retrieval_mode"] == SearchRetrievalMode.FTS + def fts_window(limit: int) -> list[SearchIndexRow]: if limit <= 8: return [ _row(id=1, title="Alpha", score=10.0), @@ -824,8 +809,12 @@ async def fake_fts_search(*args, limit: int, **kwargs) -> list[SearchIndexRow]: _row(id=4, title="Delta", score=7.0), ] - async def fake_vector_search(**kwargs) -> list[SearchIndexRow]: + fts = FakeFts(fts_window) + semantic = _semantic(vector_k=2, rerank=_reranking(reranker, candidates=2), fts=fts) + + async def fake_vector_search(query: PreparedSearchQuery, **kwargs: Any) -> list[SearchIndexRow]: candidate_limit = kwargs["candidate_limit"] + assert candidate_limit is not None if candidate_limit <= 8: return [ _row(id=1, title="Alpha", score=1.0), @@ -844,22 +833,11 @@ async def fake_vector_search(**kwargs) -> list[SearchIndexRow]: _row(id=4, title="Delta", score=0.7), ] - monkeypatch.setattr(repo, "search", fake_fts_search) - monkeypatch.setattr(repo, "_search_vector_only", fake_vector_search) + monkeypatch.setattr(semantic, "vector_only", fake_vector_search) async def deep_page(offset: int) -> list[SearchIndexRow]: - return await repo._search_hybrid( - search_text="auth", - 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=None, - temporal=None, + return await semantic.hybrid( + PreparedSearchQuery(search_text="auth", retrieval_mode=SearchRetrievalMode.HYBRID), limit=1, offset=offset, ) @@ -870,6 +848,7 @@ async def deep_page(offset: int) -> list[SearchIndexRow]: assert [row.id for row in first_tail_page] == [charlie.id] assert [row.id for row in second_tail_page] == [delta.id] assert reranker.calls == 2 + assert all(query.retrieval_mode == SearchRetrievalMode.FTS for query in fts.queries) @pytest.mark.asyncio diff --git a/tests/repository/test_search_file_path_prefix.py b/tests/repository/test_search_file_path_prefix.py index afe4be233..b668218d2 100644 --- a/tests/repository/test_search_file_path_prefix.py +++ b/tests/repository/test_search_file_path_prefix.py @@ -10,6 +10,7 @@ from datetime import datetime, timezone from types import SimpleNamespace +from typing import cast from unittest.mock import AsyncMock import pytest @@ -18,6 +19,8 @@ from basic_memory.models.knowledge import Entity from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_filters import file_path_prefix_condition +from basic_memory.repository.search_reader import HydratedChunk, SemanticSearch +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.schemas.search import ( SearchItemType, SearchRetrievalMode, @@ -305,18 +308,24 @@ async def test_semantic_retrieval_honors_the_scope( ), ) monkeypatch.setattr(search_repository, "_ensure_vector_tables", AsyncMock()) - monkeypatch.setattr(search_repository, "_prepare_vector_session", AsyncMock()) + # The nearest-neighbour stage is stubbed, so the adapter is never consulted. monkeypatch.setattr( search_repository, + "_semantic_vector_index", + cast(SemanticVectorIndex, object()), + raising=False, + ) + monkeypatch.setattr( + SemanticSearch, "_run_vector_query", AsyncMock( return_value=[ - { - "entity_id": row_id, - "chunk_key": f"entity:{row_id}:0", - "chunk_text": "subtree scope fixture", - "best_similarity": 0.9, - } + HydratedChunk( + entity_id=row_id, + chunk_key=f"entity:{row_id}:0", + chunk_text="subtree scope fixture", + similarity=0.9, + ) for row_id in seeded_paths.values() ] ), diff --git a/tests/repository/test_search_reader.py b/tests/repository/test_search_reader.py new file mode 100644 index 000000000..ed6da6e62 --- /dev/null +++ b/tests/repository/test_search_reader.py @@ -0,0 +1,283 @@ +"""SearchReader dispatch, and the SemanticSearch edges the pipeline tests do not reach.""" + +from dataclasses import replace +from itertools import count +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( + HydratedChunk, + Reranking, + SearchReader, + SemanticSearch, + parse_chunk_key, + vector_eligible, +) +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.repository.search_trace import ManifestReadiness, SearchTraceCollector +from basic_memory.repository.semantic_errors import SemanticSearchDisabledError +from basic_memory.schemas.search import SearchRetrievalMode +from tests.repository.test_hybrid_fusion import ( + HYBRID_QUERY, + FakeFts, + FakeRow, + fake_vector_retrieval, +) +from tests.repository.test_vector_threshold import run_vector_only, vector_semantic + +SCOPE = ProjectScope.single(1) +SEMANTIC_MODES = [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID] + + +def _semantic(fts: FakeFts | None = None) -> SemanticSearch: + return SemanticSearch(cast(Any, None), SCOPE, fts or FakeFts(), fake_vector_retrieval()) + + +# --- Eligibility and keys --- + + +@pytest.mark.parametrize( + ("query", "eligible"), + [ + (PreparedSearchQuery(search_text="auth"), True), + (PreparedSearchQuery(search_text=" auth "), True), + (PreparedSearchQuery(search_text=None), False), + (PreparedSearchQuery(search_text=" "), False), + (PreparedSearchQuery(search_text="*"), False), + (PreparedSearchQuery(search_text="auth", permalink="specs/auth"), False), + (PreparedSearchQuery(search_text="auth", permalink_match="specs/*"), False), + (PreparedSearchQuery(search_text="auth", title="Auth"), False), + ], +) +def test_vector_eligible_requires_text_and_no_identity_filter(query, eligible): + assert vector_eligible(query) is eligible + + +def test_parse_chunk_key_reads_type_and_row_id(): + assert parse_chunk_key("observation:5:0") == ("observation", 5) + with pytest.raises(ValueError): + parse_chunk_key("entity:not-a-number:0") + with pytest.raises(IndexError): + parse_chunk_key("garbage") + + +# --- SearchReader dispatch --- + + +@pytest.mark.asyncio +async def test_fts_mode_hands_the_engine_the_query_and_every_option(): + fts = FakeFts([FakeRow(id=1)]) + reader = SearchReader(SCOPE, fts) + query = PreparedSearchQuery(search_text="auth") + + rows = await reader.search( + query, + limit=5, + offset=2, + allow_relaxed=True, + candidate_keys=[("entity", 1)], + ) + + assert [row.id for row in rows] == [1] + assert fts.queries == [query] + assert fts.calls == [ + {"limit": 5, "offset": 2, "allow_relaxed": True, "candidate_keys": [("entity", 1)]} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", SEMANTIC_MODES) +async def test_semantic_modes_reject_queries_with_nothing_to_embed(mode): + reader = SearchReader(SCOPE, FakeFts(), _semantic()) + + with pytest.raises(ValueError, match="requires a non-empty text query"): + await reader.search( + PreparedSearchQuery(title="Auth", retrieval_mode=mode), limit=10, offset=0 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", SEMANTIC_MODES) +async def test_semantic_modes_without_a_semantic_stack_raise_disabled(mode): + """A reader built without the semantic stack answers full-text queries only.""" + reader = SearchReader(SCOPE, FakeFts()) + + with pytest.raises(SemanticSearchDisabledError): + await reader.search( + PreparedSearchQuery(search_text="auth", retrieval_mode=mode), limit=10, offset=0 + ) + + +@pytest.mark.asyncio +async def test_semantic_modes_route_to_the_pipeline(): + semantic = _semantic() + reader = SearchReader(SCOPE, FakeFts(), semantic) + vector_query = PreparedSearchQuery( + search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR + ) + + with patch.object( + semantic, "vector_only", new_callable=AsyncMock, return_value=[FakeRow(id=1)] + ) as vector_only: + rows = await reader.search(vector_query, limit=3, offset=1) + assert [row.id for row in rows] == [1] + vector_only.assert_awaited_once_with(vector_query, limit=3, offset=1, trace=None) + + with patch.object( + semantic, "hybrid", new_callable=AsyncMock, return_value=[FakeRow(id=2)] + ) as hybrid: + rows = await reader.search(HYBRID_QUERY, limit=3, offset=1) + assert [row.id for row in rows] == [2] + hybrid.assert_awaited_once_with(HYBRID_QUERY, limit=3, offset=1, trace=None) + + +@pytest.mark.asyncio +async def test_count_is_full_text_only(): + reader = SearchReader(SCOPE, FakeFts([FakeRow(id=1), FakeRow(id=2)])) + + assert await reader.count(PreparedSearchQuery(search_text="auth"), allow_relaxed=True) == 2 + with pytest.raises(ValueError, match="Exact counts are only supported"): + await reader.count( + PreparedSearchQuery(search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR) + ) + + +# --- SemanticSearch edges --- + + +@pytest.mark.asyncio +async def test_empty_candidate_window_skips_the_adapter(): + adapter = MagicMock() + adapter.search = AsyncMock(side_effect=AssertionError("adapter must not be consulted")) + vector = replace(fake_vector_retrieval(), index=adapter) + semantic = SemanticSearch(cast(Any, None), SCOPE, FakeFts(), vector) + + assert await semantic._run_vector_query(AsyncMock(), [0.1], 0) == [] + + +@pytest.mark.asyncio +async def test_no_row_ids_means_no_row_fetch(): + session_maker = MagicMock(side_effect=AssertionError("no session for an empty fetch")) + semantic = SemanticSearch(session_maker, SCOPE, FakeFts(), fake_vector_retrieval()) + + assert await semantic._fetch_search_index_rows_by_ids([]) == {} + + +@pytest.mark.asyncio +async def test_vector_only_with_no_parseable_chunk_keys_returns_nothing(): + semantic = vector_semantic() + fetch_rows = AsyncMock() + rows = [HydratedChunk(entity_id=0, chunk_key="garbage", chunk_text="bad", similarity=0.95)] + + assert await run_vector_only(semantic, rows, fetch_rows) == [] + fetch_rows.assert_not_called() + + +@pytest.mark.asyncio +async def test_hybrid_skips_rows_without_an_id_on_both_legs(): + fts = FakeFts([FakeRow(id=None, score=2.0), FakeRow(id=1, score=5.0)]) + semantic = _semantic(fts) + vector_results = [FakeRow(id=None, score=0.9), FakeRow(id=2, score=0.8)] + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=vector_results): + rows = await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + assert [(row.type, row.id) for row in rows] == [("entity", 1), ("entity", 2)] + + +@pytest.mark.asyncio +async def test_hybrid_slow_query_warning_names_the_scope(monkeypatch): + semantic = _semantic(FakeFts([FakeRow(id=1, score=5.0)])) + clock = count(0.0, 3.0) + monkeypatch.setattr( + "basic_memory.repository.search_reader.time.perf_counter", lambda: next(clock) + ) + warning = MagicMock() + monkeypatch.setattr("basic_memory.repository.search_reader.logger.warning", warning) + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + warning.assert_called_once() + assert warning.call_args.args[0].startswith("[SEMANTIC_SLOW_QUERY]") + assert warning.call_args.kwargs["retrieval_mode"] == "hybrid" + assert warning.call_args.kwargs["scope"] == (1,) + + +@pytest.mark.asyncio +async def test_built_in_adapter_reads_manifest_readiness_when_tracing(monkeypatch): + """A traced built-in lookup reports how much of the manifest retrieval can answer from.""" + adapter = MagicMock() + adapter.search = AsyncMock(return_value=[]) + semantic = SemanticSearch( + cast(Any, None), SCOPE, FakeFts(), replace(fake_vector_retrieval(), index=adapter) + ) + readiness = ManifestReadiness( + configured_index="sqlite-vec", + configured_model="fake:384", + ready_rows=3, + pending_rows=1, + other_identity_rows=0, + ) + read_readiness = AsyncMock(return_value=readiness) + monkeypatch.setattr( + "basic_memory.repository.search_reader.read_manifest_readiness", read_readiness + ) + trace = SearchTraceCollector() + session = AsyncMock() + + assert await semantic._run_vector_query(session, [0.1], 5, trace=trace) == [] + + assert trace.readiness is readiness + read_readiness.assert_awaited_once_with(session, SCOPE, "sqlite-vec", "fake:384") + adapter.search.assert_awaited_once_with([0.1], limit=5) + + +@pytest.mark.asyncio +async def test_fts_gate_zeroes_weak_lexical_scores(monkeypatch): + """Below the gate a lexical hit contributes nothing to fusion instead of a sliver.""" + monkeypatch.setattr("basic_memory.repository.search_reader.FTS_GATE_THRESHOLD", 0.5) + semantic = _semantic(FakeFts([FakeRow(id=1, score=10.0), FakeRow(id=2, score=1.0)])) + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + rows = await semantic.hybrid(HYBRID_QUERY, limit=10, offset=0) + + assert [(row.id, row.score) for row in rows] == [(1, 1.0), (2, 0.0)] + + +class _PrefixReranker: + """Scores documents by position so the rerank stage is deterministic.""" + + model_name = "prefix" + + async def rerank(self, query: str, documents: list[str]) -> list[float]: + return [1.0 - index * 0.1 for index in range(len(documents))] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +@pytest.mark.asyncio +async def test_hybrid_trace_records_the_stable_pool_refetch(): + """A page past the fixed rerank prefix refetches the stable pool, and the trace says so.""" + rows = [ + FakeRow(id=index, score=10.0 - index, title=f"n{index}", entity_id=index) + for index in range(1, 6) + ] + semantic = SemanticSearch( + cast(Any, None), + SCOPE, + FakeFts(rows), + replace(fake_vector_retrieval(), vector_k=2), + Reranking(provider=_PrefixReranker(), candidates=2, max_document_chars=0), + ) + trace = SearchTraceCollector() + + with patch.object(semantic, "vector_only", new_callable=AsyncMock, return_value=[]): + page = await semantic.hybrid(HYBRID_QUERY, limit=1, offset=3, trace=trace) + + assert trace.stable_pool_refetched is True + assert [row.id for row in page] == [4] diff --git a/tests/repository/test_search_trace.py b/tests/repository/test_search_trace.py index c6b913bcb..89efb107d 100644 --- a/tests/repository/test_search_trace.py +++ b/tests/repository/test_search_trace.py @@ -14,7 +14,7 @@ from basic_memory.config import BasicMemoryConfig, DatabaseBackend 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 FUSION_BONUS +from basic_memory.repository.search_reader import FUSION_BONUS from basic_memory.repository.search_trace import ( BelowThreshold, FilteredOut, @@ -751,7 +751,7 @@ async def test_classify_hydration_drop_observes_pending_to_ready_transition( ) await session.commit() async with db.scoped_session(session_maker) as session: - assert await repository._hydrate_vector_matches(session, [match]) == [] + assert await repository._semantic_search()._hydrate_vector_matches(session, [match]) == [] async with db.scoped_session(session_maker) as session: await session.execute( diff --git a/tests/repository/test_semantic_search_base.py b/tests/repository/test_semantic_search_base.py index d9ae052f1..4a5857aca 100644 --- a/tests/repository/test_semantic_search_base.py +++ b/tests/repository/test_semantic_search_base.py @@ -8,16 +8,21 @@ from contextlib import asynccontextmanager from datetime import datetime from types import SimpleNamespace -from typing import override, Any +from typing import override, Any, cast from unittest.mock import AsyncMock, Mock import pytest import basic_memory.repository.search_repository_base as search_repository_base_module +from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider -from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_index_row import SearchIndexKey, SearchIndexRow +from basic_memory.repository.search_reader import ( + HydratedChunk, + SemanticSearch, + VectorRetrieval, +) from basic_memory.repository.search_repository_base import ( - SearchIndexKey, SearchRepositoryBase, _PreparedEntityVectorSync, ) @@ -37,6 +42,7 @@ from basic_memory.repository.semantic_vector_sync import PendingEmbeddingJob from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode from basic_memory.temporal import TemporalFilter +from tests.repository.test_hybrid_fusion import FakeFts # --- Helpers --- @@ -67,6 +73,7 @@ def __init__(self): self.session_maker = None self.project_id = 1 self.scope = ProjectScope.single(1) + self._fts = FakeFts() @override async def init_search_index(self): @@ -113,17 +120,6 @@ async def search( async def _ensure_vector_tables(self): pass - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] - @override async def _write_embeddings(self, session, jobs, embeddings): pass @@ -146,10 +142,6 @@ async def _delete_stale_chunks( async def _update_timestamp_sql(self): return "CURRENT_TIMESTAMP" - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - class _RecordingVectorIndex: """Protocol-complete adapter that records generation-safe upserts.""" @@ -185,12 +177,25 @@ async def search( return [] +def _semantic_search(*, index_name: str, adapter: Any = None) -> SemanticSearch: + """The vector pipeline over an adapter the test controls.""" + vector = VectorRetrieval( + index=adapter if adapter is not None else _RecordingVectorIndex(), + index_name=index_name, + embedding_provider=cast( + EmbeddingProvider, SimpleNamespace(model_name="stub", dimensions=4) + ), + embedding_model="stub:4", + vector_k=100, + min_similarity=0.0, + ) + return SemanticSearch(cast(Any, None), ProjectScope.single(1), FakeFts(), vector) + + @pytest.mark.asyncio async def test_vector_match_hydration_batches_large_adapter_results() -> None: """Deep vector pages must not create an unbounded SQL bind-parameter list.""" - repo = _ConcreteRepo() - repo._semantic_vector_index_name = "milvus" - repo._embedding_provider = SimpleNamespace(model_name="stub", dimensions=4) + semantic = _semantic_search(index_name="milvus") matches = [ VectorMatch( key=VectorKey(entity_id=entity_id, chunk_key=f"entity:{entity_id}:0"), @@ -214,11 +219,11 @@ def hydrated_batch(_statement, params): session.execute.side_effect = hydrated_batch - hydrated = await repo._hydrate_vector_matches(session, matches) + hydrated = await semantic._hydrate_vector_matches(session, matches) assert session.execute.await_count == 3 assert max(len(call.args[1]) for call in session.execute.await_args_list) == 503 - assert [row["entity_id"] for row in hydrated] == list(range(600)) + assert [chunk.entity_id for chunk in hydrated] == list(range(600)) @pytest.mark.asyncio @@ -226,8 +231,6 @@ async def test_external_vector_query_overfetches_past_stale_adapter_hits( monkeypatch: pytest.MonkeyPatch, ) -> None: """Stale top-k extension hits must not crowd live manifest rows out.""" - repo = _ConcreteRepo() - repo._semantic_vector_index_name = "milvus" def matches(count: int) -> list[VectorMatch]: return [ @@ -239,15 +242,15 @@ def matches(count: int) -> list[VectorMatch]: ] adapter: Any = SimpleNamespace(search=AsyncMock(side_effect=[matches(2), matches(4)])) - repo._semantic_vector_index = adapter + semantic = _semantic_search(index_name="milvus", adapter=adapter) live_rows = [ - {"entity_id": 2, "chunk_key": "entity:2:0", "best_similarity": 0.9}, - {"entity_id": 3, "chunk_key": "entity:3:0", "best_similarity": 0.8}, + HydratedChunk(entity_id=2, chunk_key="entity:2:0", chunk_text="two", similarity=0.9), + HydratedChunk(entity_id=3, chunk_key="entity:3:0", chunk_text="three", similarity=0.8), ] hydrate = AsyncMock(side_effect=[[], live_rows]) - monkeypatch.setattr(repo, "_hydrate_vector_matches", hydrate) + monkeypatch.setattr(semantic, "_hydrate_vector_matches", hydrate) - result = await SearchRepositoryBase._run_vector_query(repo, AsyncMock(), [0.1], 2) + result = await semantic._run_vector_query(AsyncMock(), [0.1], 2) assert result == live_rows assert [call.kwargs["limit"] for call in adapter.search.await_args_list] == [2, 4] diff --git a/tests/repository/test_semantic_vector_sync.py b/tests/repository/test_semantic_vector_sync.py index fb1c12318..0e1003424 100644 --- a/tests/repository/test_semantic_vector_sync.py +++ b/tests/repository/test_semantic_vector_sync.py @@ -78,17 +78,6 @@ async def search( async def _ensure_vector_tables(self): pass - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] - @override async def _write_embeddings(self, session, jobs, embeddings): pass @@ -108,10 +97,6 @@ async def _delete_stale_chunks( ): return [] - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - def _pending_job( entity_id: int = 1, diff --git a/tests/repository/test_vector_filter_candidate_restriction.py b/tests/repository/test_vector_filter_candidate_restriction.py index a40413862..8407a037a 100644 --- a/tests/repository/test_vector_filter_candidate_restriction.py +++ b/tests/repository/test_vector_filter_candidate_restriction.py @@ -30,10 +30,14 @@ from basic_memory import db from basic_memory.repository.embedding_provider import EmbeddingProvider from basic_memory.repository.search_filters import candidate_key_restriction_condition -from basic_memory.repository.search_repository_base import ( +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( VECTOR_FILTER_SCAN_LIMIT, VECTOR_HYDRATION_BATCH_SIZE, + HydratedChunk, + SemanticSearch, ) +from basic_memory.repository.semantic_vector_index import SemanticVectorIndex from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode # The scope the admitted rows share, plus a sibling scope the filter must reject. @@ -138,7 +142,11 @@ def _fake_embedding_provider() -> EmbeddingProvider: type( "EP", (), - {"embed_query": AsyncMock(return_value=[0.0] * 384), "dimensions": 384}, + { + "embed_query": AsyncMock(return_value=[0.0] * 384), + "dimensions": 384, + "model_name": "stub", + }, )(), ) @@ -148,18 +156,20 @@ def _semantic_repo(search_repository): search_repository._semantic_enabled = True search_repository._semantic_min_similarity = 0.0 search_repository._embedding_provider = _fake_embedding_provider() + # The nearest-neighbour stage is stubbed below, so the adapter is never consulted. + search_repository._semantic_vector_index = cast(SemanticVectorIndex, object()) return search_repository -def _vector_chunks(row_ids: list[int]) -> list[dict[str, Any]]: +def _vector_chunks(row_ids: list[int]) -> list[HydratedChunk]: """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, - } + HydratedChunk( + entity_id=row_id, + chunk_key=f"{SearchItemType.ENTITY.value}:{row_id}:0", + chunk_text=TARGET_CONTENT, + similarity=0.99 - index * 0.001, + ) for index, row_id in enumerate(row_ids) ] @@ -174,9 +184,8 @@ async def test_filtered_vector_search_keeps_a_candidate_past_the_scan_window( with ( patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), patch.object( - repo, + SemanticSearch, "_run_vector_query", new_callable=AsyncMock, return_value=_vector_chunks([TARGET_ROW_ID]), @@ -210,36 +219,28 @@ async def test_filter_pass_answers_every_candidate_within_the_bind_bound( assert len(candidates) > VECTOR_HYDRATION_BATCH_SIZE batched_key_counts: list[int] = [] - original_search = repo.search + original_search = repo._fts.search - async def recording_search(*args, **kwargs): + async def recording_search(scope, query, **kwargs): if kwargs.get("candidate_keys") is not None: batched_key_counts.append(len(kwargs["candidate_keys"])) - return await original_search(*args, **kwargs) + return await original_search(scope, query, **kwargs) with ( - patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock), - patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock), patch.object( - repo, + SemanticSearch, "_run_vector_query", new_callable=AsyncMock, return_value=_vector_chunks(candidates), ), - patch.object(repo, "search", recording_search), + patch.object(repo._fts, "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, + results = await repo._semantic_search().vector_only( + PreparedSearchQuery( + search_text="the answer", + file_path_prefix=SCOPE, + retrieval_mode=SearchRetrievalMode.VECTOR, + ), limit=len(candidates), offset=0, ) diff --git a/tests/repository/test_vector_pagination.py b/tests/repository/test_vector_pagination.py index f2141ff05..ed79d7ce3 100644 --- a/tests/repository/test_vector_pagination.py +++ b/tests/repository/test_vector_pagination.py @@ -4,207 +4,43 @@ which requires a sufficiently large candidate_limit multiplier. """ -from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.repository.search_scope import ProjectScope -from collections.abc import Sequence -from contextlib import asynccontextmanager -from dataclasses import dataclass -from datetime import datetime -from typing import override, Any -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock import pytest -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 -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.search_reader import HydratedChunk +from tests.repository.test_vector_threshold import FakeRow, run_vector_only, vector_semantic -@dataclass -class FakeRow: - """Minimal stand-in for SearchIndexRow in pagination tests.""" - - id: int - type: str = "entity" - score: float = 0.0 - matched_chunk_text: str | None = None - content_snippet: str | None = None - - -class ConcreteSearchRepo(SearchRepositoryBase): - """Minimal concrete subclass for testing base class pagination logic.""" - - def __init__(self): - self._semantic_enabled = True - self._semantic_vector_k = 100 - self._semantic_min_similarity = 0.0 - self._embedding_provider = None - self._vector_dimensions = 384 - self._vector_tables_initialized = True - self.session_maker = None - self.project_id = 1 - self.scope = ProjectScope.single(1) - - @override - async def init_search_index(self): - pass # pragma: no cover - - @override - async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: - return None # physical storage is not inspectable in this double - - @override - async def search( - self, - search_text: str | None = None, - permalink: str | None = None, - permalink_match: str | None = None, - title: str | None = None, - note_types: list[str] | None = None, - after_date: datetime | None = None, - search_item_types: list[SearchItemType] | None = None, - categories: list[str] | None = None, - metadata_filters: dict[str, Any] | None = None, - file_path_prefix: str | None = None, - temporal: TemporalFilter | None = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: float | None = None, - limit: int = 10, - offset: int = 0, - allow_relaxed: bool = False, - session: AsyncSession | None = None, - *, - candidate_keys: Sequence[SearchIndexKey] | None = None, - trace: SearchTraceCollector | None = None, - ) -> list[SearchIndexRow]: - return [] # pragma: no cover - - @override - async def _ensure_vector_tables(self): - pass # pragma: no cover - - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] # pragma: no cover - - @override - async def _write_embeddings(self, session, jobs, embeddings): - pass # pragma: no cover - - @override - async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): - return [] # pragma: no cover - - @override - async def _delete_stale_chunks( - self, - session, - stale_ids, - entity_id, - *, - expected_deletions=None, - ): - return [] # pragma: no cover - - async def _update_timestamp_sql(self): - return "CURRENT_TIMESTAMP" # pragma: no cover - - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - - -@asynccontextmanager -async def fake_scoped_session(session_maker): - yield AsyncMock() - - -class _EmbeddingProvider: - dimensions = 384 - model_name = "stub" - - async def embed_query(self, text: str) -> list[float]: - return [0.0] * self.dimensions - - async def embed_documents(self, texts: list[str]) -> list[list[float]]: - return [[0.0] * self.dimensions for _ in texts] - - def runtime_log_attrs(self) -> dict[str, object]: - return {} - - -def _make_descending_vector_rows(count: int) -> list[dict[str, Any]]: - """Build vector rows with scores descending from ~1.0 to ~0.5.""" - rows = [] - for i in range(count): - # Similarity decreases linearly: 0.95, 0.94, 0.93, ... - similarity = 0.95 - (i * 0.01) - distance = (1.0 / similarity) - 1.0 - rows.append( - { - "chunk_key": f"entity:{i}:0", - "best_distance": distance, - "chunk_text": f"chunk text {i}", - } +def _make_descending_vector_rows(count: int) -> list[HydratedChunk]: + """Build vector rows with similarity descending from 0.95 in steps of 0.01.""" + return [ + HydratedChunk( + entity_id=index, + chunk_key=f"entity:{index}:0", + chunk_text=f"chunk text {index}", + similarity=0.95 - (index * 0.01), ) - return rows + for index in range(count) + ] @pytest.mark.asyncio async def test_page1_scores_gte_page2_scores(): """Page 1 minimum score must be >= page 2 maximum score.""" - repo = ConcreteSearchRepo() - + semantic = vector_semantic() # 20 results with descending scores fake_rows = _make_descending_vector_rows(20) - - repo._embedding_provider = _EmbeddingProvider() - fake_index_rows = {("entity", i): FakeRow(id=i) for i in range(20)} - async def run_page(offset, limit): - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", - fake_scoped_session, - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value=fake_index_rows, - ), - ): - return await repo._search_vector_only( - search_text="test", - 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=None, - temporal=None, - limit=limit, - offset=offset, - ) + async def run_page(offset: int, limit: int): + return await run_vector_only( + semantic, + fake_rows, + AsyncMock(return_value=fake_index_rows), + limit=limit, + offset=offset, + ) page1 = await run_page(offset=0, limit=10) page2 = await run_page(offset=10, limit=10) diff --git a/tests/repository/test_vector_temporal_filter.py b/tests/repository/test_vector_temporal_filter.py index 0270d7815..0cf4a73cd 100644 --- a/tests/repository/test_vector_temporal_filter.py +++ b/tests/repository/test_vector_temporal_filter.py @@ -10,130 +10,88 @@ pin both halves at the seam rather than trusting the call sites to stay in step. """ -from typing import Any +from dataclasses import replace +from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest +from basic_memory.repository.search_reader import SemanticSearch +from basic_memory.repository.search_scope import ProjectScope from basic_memory.temporal import TemporalFilter, TimeKind, parse_point from tests.repository.test_hybrid_fusion import ( - HYBRID_KWARGS, - ConcreteSearchRepo as HybridSearchRepo, + HYBRID_QUERY, + FakeFts, FakeRow as HybridFakeRow, + fake_vector_retrieval, ) from tests.repository.test_vector_threshold import ( - COMMON_SEARCH_KWARGS, - ConcreteSearchRepo as VectorSearchRepo, + VECTOR_QUERY, FakeRow, - _fake_embedding_provider, _make_vector_rows, - fake_scoped_session, + run_vector_only, + vector_semantic, ) TEMPORAL = TemporalFilter(kind=TimeKind.EFFECTIVE, at=parse_point("2026-07-28")) -def _vector_kwargs(**overrides: Any) -> dict[str, Any]: - return {**COMMON_SEARCH_KWARGS, **overrides} - - -def _hybrid_kwargs(**overrides: Any) -> dict[str, Any]: - return {**HYBRID_KWARGS, **overrides} - - -def _forwarded_temporal(leg: AsyncMock) -> Any: - """The `temporal` argument one retrieval leg was actually called with.""" - assert leg.await_args is not None, "leg was never awaited" - return leg.await_args.kwargs["temporal"] - - @pytest.mark.asyncio async def test_temporal_filter_applies_in_vector_mode(): """A valid-time filter narrows the vector candidate set, and is forwarded verbatim.""" - repo = VectorSearchRepo() - repo._semantic_min_similarity = 0.0 - repo._embedding_provider = _fake_embedding_provider(AsyncMock(return_value=[0.0] * 384)) - # The embedding neighbourhood offers three entities; only entity 1 asserts a range # covering the queried date, so the FTS intersection pass returns just that one. - filter_pass = AsyncMock(return_value=[FakeRow(id=1)]) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=_make_vector_rows([0.9, 0.8, 0.7]), - ), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(3)}, - ), - patch.object(repo, "search", filter_pass), - ): - results = await repo._search_vector_only(**_vector_kwargs(temporal=TEMPORAL)) + filter_pass = FakeFts([FakeRow(id=1)]) + semantic = vector_semantic(fts=filter_pass) + + results = await run_vector_only( + semantic, + _make_vector_rows([0.9, 0.8, 0.7]), + AsyncMock(return_value={("entity", i): FakeRow(id=i) for i in range(3)}), + query=replace(VECTOR_QUERY, temporal=TEMPORAL), + ) assert [row.id for row in results] == [1] # Counted as a requested filter... - filter_pass.assert_awaited_once() - # ...and forwarded unchanged, so the intersection asks the same question. - assert _forwarded_temporal(filter_pass) is TEMPORAL + assert len(filter_pass.queries) == 1 + # ...and forwarded unchanged, so the intersection asks the same question, + # about the candidates themselves rather than a page of the whole match set. + assert filter_pass.queries[0].temporal is TEMPORAL + assert filter_pass.queries[0].search_text is None + assert filter_pass.calls[0]["candidate_keys"] == [("entity", 0), ("entity", 1), ("entity", 2)] @pytest.mark.asyncio async def test_vector_mode_without_a_temporal_filter_runs_no_intersection_pass(): """An unfiltered semantic search must not pay for a filter pass it does not need.""" - repo = VectorSearchRepo() - repo._semantic_min_similarity = 0.0 - repo._embedding_provider = _fake_embedding_provider(AsyncMock(return_value=[0.0] * 384)) - filter_pass = AsyncMock(return_value=[]) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=_make_vector_rows([0.9]), - ), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0)}, - ), - patch.object(repo, "search", filter_pass), - ): - results = await repo._search_vector_only(**_vector_kwargs()) + filter_pass = FakeFts([]) + semantic = vector_semantic(fts=filter_pass) + + results = await run_vector_only( + semantic, + _make_vector_rows([0.9]), + AsyncMock(return_value={("entity", 0): FakeRow(id=0)}), + ) assert [row.id for row in results] == [0] - filter_pass.assert_not_awaited() + assert filter_pass.queries == [] @pytest.mark.asyncio async def test_temporal_filter_applies_in_hybrid_mode(): """Hybrid fuses two legs; both must ask the same valid-time question.""" - repo = HybridSearchRepo() - fts_leg = AsyncMock(return_value=[HybridFakeRow(id=1, score=5.0, title="dated")]) + fts_leg = FakeFts([HybridFakeRow(id=1, score=5.0, title="dated")]) + semantic = SemanticSearch( + cast(Any, None), ProjectScope.single(1), fts_leg, fake_vector_retrieval() + ) vector_leg = AsyncMock(return_value=[HybridFakeRow(id=1, score=0.9, title="dated")]) - with ( - patch.object(repo, "search", fts_leg), - patch.object(repo, "_search_vector_only", vector_leg), - ): - results = await repo._search_hybrid(**_hybrid_kwargs(temporal=TEMPORAL)) + with patch.object(semantic, "vector_only", vector_leg): + results = await semantic.hybrid( + replace(HYBRID_QUERY, temporal=TEMPORAL), limit=10, offset=0 + ) assert [row.id for row in results] == [1] - assert _forwarded_temporal(fts_leg) is TEMPORAL - assert _forwarded_temporal(vector_leg) is TEMPORAL + assert fts_leg.queries[0].temporal is TEMPORAL + assert vector_leg.await_args is not None, "vector leg was never awaited" + assert vector_leg.await_args.args[0].temporal is TEMPORAL diff --git a/tests/repository/test_vector_threshold.py b/tests/repository/test_vector_threshold.py index fdfb6d8a7..4314176d5 100644 --- a/tests/repository/test_vector_threshold.py +++ b/tests/repository/test_vector_threshold.py @@ -1,27 +1,22 @@ """Tests for semantic_min_similarity threshold filtering in vector search.""" -from sqlalchemy.ext.asyncio import AsyncSession -from basic_memory.repository.search_scope import ProjectScope -from collections.abc import Sequence from contextlib import asynccontextmanager -from dataclasses import dataclass -from datetime import datetime -from typing import override, Any, Optional, cast +from dataclasses import dataclass, replace +from typing import Any from unittest.mock import AsyncMock, patch import pytest -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 ( +from basic_memory.repository.search_query import PreparedSearchQuery +from basic_memory.repository.search_reader import ( SMALL_NOTE_CONTENT_LIMIT, TOP_CHUNKS_PER_RESULT, - SearchIndexKey, - SearchRepositoryBase, + HydratedChunk, + SemanticSearch, ) -from basic_memory.repository.search_trace import SearchTraceCollector -from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode -from basic_memory.temporal import TemporalFilter +from basic_memory.repository.search_scope import ProjectScope +from basic_memory.schemas.search import SearchRetrievalMode +from tests.repository.test_hybrid_fusion import FakeFts, fake_vector_retrieval @dataclass @@ -35,174 +30,71 @@ class FakeRow: content_snippet: str | None = None -class ConcreteSearchRepo(SearchRepositoryBase): - """Minimal concrete subclass for testing base class threshold logic.""" - - def __init__(self): - # Skip super().__init__ — we only need the attributes under test - self._semantic_enabled = True - self._semantic_vector_k = 100 - self._semantic_min_similarity = 0.0 - self._embedding_provider = None - self._vector_dimensions = 384 - self._vector_tables_initialized = True - self.session_maker = None - self.project_id = 1 - self.scope = ProjectScope.single(1) - - # --- Abstract method stubs (not exercised by these tests) --- - - @override - async def init_search_index(self): - pass # pragma: no cover - - @override - async def get_entity_physical_chunk_keys(self, entity_id: int) -> set[str] | None: - return None # physical storage is not inspectable in this double - - @override - async def search( - self, - search_text: Optional[str] = None, - permalink: Optional[str] = None, - permalink_match: Optional[str] = None, - title: Optional[str] = None, - note_types: Optional[list[str]] = None, - after_date: Optional[datetime] = None, - search_item_types: Optional[list[SearchItemType]] = None, - categories: Optional[list[str]] = None, - metadata_filters: Optional[dict[str, Any]] = None, - file_path_prefix: Optional[str] = None, - temporal: Optional[TemporalFilter] = None, - retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS, - min_similarity: Optional[float] = None, - limit: int = 10, - offset: int = 0, - allow_relaxed: bool = False, - session: AsyncSession | None = None, - *, - candidate_keys: Sequence[SearchIndexKey] | None = None, - trace: SearchTraceCollector | None = None, - ) -> list[SearchIndexRow]: - return [] # pragma: no cover - - @override - async def _ensure_vector_tables(self): - pass # pragma: no cover - - @override - async def _run_vector_query( - self, - session, - query_embedding, - candidate_limit, - *, - trace: SearchTraceCollector | None = None, - ): - return [] # pragma: no cover - - @override - async def _write_embeddings(self, session, jobs, embeddings): - pass # pragma: no cover - - @override - async def _delete_entity_chunks(self, session, entity_id, *, expected_deletions=None): - return [] # pragma: no cover - - @override - async def _delete_stale_chunks( - self, - session, - stale_ids, - entity_id, - *, - expected_deletions=None, - ): - return [] # pragma: no cover - - async def _update_timestamp_sql(self): - return "CURRENT_TIMESTAMP" # pragma: no cover - - @override - def _distance_to_similarity(self, distance: float) -> float: - return 1.0 / (1.0 + max(distance, 0.0)) - - -def _make_vector_rows(scores: list[float]) -> list[dict[str, Any]]: - """Build fake vector query rows with controlled distances. - - Distance = (1/score) - 1 inverts the similarity formula: - similarity = 1 / (1 + distance) - """ - rows = [] - for i, score in enumerate(scores): - distance = (1.0 / score) - 1.0 - rows.append( - { - "chunk_key": f"entity:{i}:0", - "best_distance": distance, - "chunk_text": f"chunk text for entity:{i}:0", - } +def _make_vector_rows(scores: list[float]) -> list[HydratedChunk]: + """One hydrated chunk per search row, ranked at the given similarity.""" + return [ + HydratedChunk( + entity_id=index, + chunk_key=f"entity:{index}:0", + chunk_text=f"chunk text for entity:{index}:0", + similarity=score, ) - return rows + for index, score in enumerate(scores) + ] -def _fake_embedding_provider(mock_embed: AsyncMock) -> EmbeddingProvider: - return cast( - EmbeddingProvider, - type("EP", (), {"embed_query": mock_embed, "dimensions": 384})(), - ) +def fake_session_maker() -> Any: + """A session factory for the hydration step; the stubbed stages never touch it.""" + @asynccontextmanager + async def session(): + yield AsyncMock() -@asynccontextmanager -async def fake_scoped_session(session_maker): - """Fake scoped_session that yields a mock session object.""" - yield AsyncMock() - - -COMMON_SEARCH_KWARGS: dict[str, Any] = dict( - search_text="test", - 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=None, - temporal=None, - limit=10, - offset=0, -) + return session -@pytest.mark.asyncio -async def test_threshold_zero_returns_all(): - """With threshold=0.0 (default), all results pass through.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 +def vector_semantic(*, min_similarity: float = 0.0, fts: FakeFts | None = None) -> SemanticSearch: + """A vector pipeline with a stubbed adapter, ready for its neighbour stage to be patched.""" + return SemanticSearch( + fake_session_maker(), + ProjectScope.single(1), + fts or FakeFts(), + fake_vector_retrieval(min_similarity=min_similarity), + ) + - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) +VECTOR_QUERY = PreparedSearchQuery(search_text="test", retrieval_mode=SearchRetrievalMode.VECTOR) - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) +async def run_vector_only( + semantic: SemanticSearch, + vector_rows: list[HydratedChunk], + fetch_rows: AsyncMock, + *, + query: PreparedSearchQuery = VECTOR_QUERY, + limit: int = 10, + offset: int = 0, +) -> list[Any]: + """Run vector-only search with the neighbour and row-fetch stages stubbed.""" with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(3)}, + semantic, "_run_vector_query", new_callable=AsyncMock, return_value=vector_rows ), + patch.object(semantic, "_fetch_search_index_rows_by_ids", fetch_rows), ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + return await semantic.vector_only(query, limit=limit, offset=offset) + + +def _index_rows(count: int) -> AsyncMock: + return AsyncMock(return_value={("entity", i): FakeRow(id=i) for i in range(count)}) + + +@pytest.mark.asyncio +async def test_threshold_zero_returns_all(): + """With threshold=0.0 (default), all results pass through.""" + semantic = vector_semantic(min_similarity=0.0) + + results = await run_vector_only(semantic, _make_vector_rows([0.9, 0.5, 0.3]), _index_rows(3)) assert len(results) == 3 @@ -210,129 +102,56 @@ async def test_threshold_zero_returns_all(): @pytest.mark.asyncio async def test_threshold_filters_low_scores(): """Results below the threshold are excluded.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.6 - - # Scores: 0.9 (pass), 0.5 (fail), 0.3 (fail) - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) + semantic = vector_semantic(min_similarity=0.6) - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - # Only entity_0 (score=0.9) passes the threshold; the fetch only gets id 0 - return_value={("entity", 0): FakeRow(id=0)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + # Scores: 0.9 (pass), 0.5 (fail), 0.3 (fail). Only entity_0 reaches the row fetch. + results = await run_vector_only(semantic, _make_vector_rows([0.9, 0.5, 0.3]), _index_rows(1)) - # Only the 0.9 result passes the 0.6 threshold assert len(results) == 1 @pytest.mark.asyncio async def test_threshold_returns_empty_when_all_below(): """All results below threshold → empty list, no DB fetch.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.8 - - # All scores below 0.8 - fake_rows = _make_vector_rows([0.5, 0.4, 0.3]) + semantic = vector_semantic(min_similarity=0.8) + fetch_rows = AsyncMock() - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - - mock_fetch = AsyncMock() - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object(repo, "_fetch_search_index_rows_by_ids", mock_fetch), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only(semantic, _make_vector_rows([0.5, 0.4, 0.3]), fetch_rows) assert results == [] # Should short-circuit before fetching search_index rows - mock_fetch.assert_not_called() + fetch_rows.assert_not_called() @pytest.mark.asyncio -async def test_per_query_min_similarity_overrides_instance_default(): - """Per-query min_similarity takes precedence over instance-level default.""" - repo = ConcreteSearchRepo() - # Instance default would filter out 0.5 and 0.3 - repo._semantic_min_similarity = 0.6 - - # Scores: 0.9, 0.5, 0.3 - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(3)}, - ), - ): - # Override to 0.0 → all results pass through despite instance default of 0.6 - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS, min_similarity=0.0) +async def test_per_query_min_similarity_overrides_configured_default(): + """Per-query min_similarity takes precedence over the configured default.""" + # The configured default would filter out 0.5 and 0.3 + semantic = vector_semantic(min_similarity=0.6) + + # Override to 0.0 → all results pass through despite the configured 0.6 + results = await run_vector_only( + semantic, + _make_vector_rows([0.9, 0.5, 0.3]), + _index_rows(3), + query=replace(VECTOR_QUERY, min_similarity=0.0), + ) assert len(results) == 3 @pytest.mark.asyncio async def test_per_query_min_similarity_tightens_threshold(): - """Per-query min_similarity=0.8 filters more aggressively than instance default.""" - repo = ConcreteSearchRepo() - # Instance default is permissive - repo._semantic_min_similarity = 0.0 - - # Scores: 0.9, 0.5, 0.3 - fake_rows = _make_vector_rows([0.9, 0.5, 0.3]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - # Only id=0 (score=0.9) will be fetched after filtering - return_value={("entity", 0): FakeRow(id=0)}, - ), - ): - # Override to 0.8 → only score=0.9 passes - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS, min_similarity=0.8) + """Per-query min_similarity=0.8 filters more aggressively than the configured default.""" + semantic = vector_semantic(min_similarity=0.0) + + # Override to 0.8 → only score=0.9 passes + results = await run_vector_only( + semantic, + _make_vector_rows([0.9, 0.5, 0.3]), + _index_rows(1), + query=replace(VECTOR_QUERY, min_similarity=0.8), + ) assert len(results) == 1 assert results[0].id == 0 @@ -341,29 +160,9 @@ async def test_per_query_min_similarity_tightens_threshold(): @pytest.mark.asyncio async def test_matched_chunk_text_populated_on_vector_results(): """Vector search results carry the matched chunk text from the best-matching chunk.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - fake_rows = _make_vector_rows([0.9, 0.7]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) + semantic = vector_semantic() - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", i): FakeRow(id=i) for i in range(2)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only(semantic, _make_vector_rows([0.9, 0.7]), _index_rows(2)) assert len(results) == 2 # Results are sorted by score descending, so id=0 (0.9) first, id=1 (0.7) second @@ -372,56 +171,32 @@ async def test_matched_chunk_text_populated_on_vector_results(): assert results[1].matched_chunk_text == "chunk text for entity:1:0" -def _make_multi_chunk_vector_rows(si_id: int, scores: list[float]) -> list[dict[str, Any]]: - """Build multiple fake vector chunks for a single search_index row. - - Each chunk gets a unique chunk_index within the same si_id. - Distance = (1/score) - 1 inverts the similarity formula. - """ - rows = [] - for chunk_idx, score in enumerate(scores): - distance = (1.0 / score) - 1.0 - rows.append( - { - "chunk_key": f"entity:{si_id}:{chunk_idx}", - "best_distance": distance, - "chunk_text": f"chunk-{chunk_idx} (sim={score})", - } +def _make_multi_chunk_vector_rows(si_id: int, scores: list[float]) -> list[HydratedChunk]: + """Several chunks of one search row, each at its own similarity.""" + return [ + HydratedChunk( + entity_id=si_id, + chunk_key=f"entity:{si_id}:{chunk_index}", + chunk_text=f"chunk-{chunk_index} (sim={score})", + similarity=score, ) - return rows + for chunk_index, score in enumerate(scores) + ] @pytest.mark.asyncio async def test_top_n_chunks_joined_in_matched_chunk_text(): """Large note with 7 chunks: top 5 by similarity are joined with separator.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - # 7 chunks for entity 0, with varying similarities + semantic = vector_semantic() chunk_scores = [0.6, 0.9, 0.4, 0.8, 0.75, 0.3, 0.85] - fake_rows = _make_multi_chunk_vector_rows(si_id=0, scores=chunk_scores) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - # content_snippet exceeds SMALL_NOTE_CONTENT_LIMIT → top-N chunks path large_content = "x" * (SMALL_NOTE_CONTENT_LIMIT + 1) - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only( + semantic, + _make_multi_chunk_vector_rows(si_id=0, scores=chunk_scores), + AsyncMock(return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}), + ) assert len(results) == 1 text = results[0].matched_chunk_text @@ -440,32 +215,15 @@ async def test_top_n_chunks_joined_in_matched_chunk_text(): @pytest.mark.asyncio async def test_small_note_returns_full_content_as_matched_chunk(): """Small note (content_snippet under limit) returns full content instead of chunks.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - fake_rows = _make_vector_rows([0.9]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - + semantic = vector_semantic() small_content = "This is a short note with all the important details." assert len(small_content) <= SMALL_NOTE_CONTENT_LIMIT - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0, content_snippet=small_content)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only( + semantic, + _make_vector_rows([0.9]), + AsyncMock(return_value={("entity", 0): FakeRow(id=0, content_snippet=small_content)}), + ) assert len(results) == 1 # Full content returned instead of the chunk text @@ -475,33 +233,31 @@ async def test_small_note_returns_full_content_as_matched_chunk(): @pytest.mark.asyncio async def test_large_note_returns_chunks_not_full_content(): """Large note (content_snippet over limit) returns top-N chunks, not full content.""" - repo = ConcreteSearchRepo() - repo._semantic_min_similarity = 0.0 - - fake_rows = _make_vector_rows([0.9]) - - mock_embed = AsyncMock(return_value=[0.0] * 384) - repo._embedding_provider = _fake_embedding_provider(mock_embed) - + semantic = vector_semantic() large_content = "x" * (SMALL_NOTE_CONTENT_LIMIT + 500) - with ( - patch( - "basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session - ), - 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=fake_rows), - patch.object( - repo, - "_fetch_search_index_rows_by_ids", - new_callable=AsyncMock, - return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}, - ), - ): - results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS) + results = await run_vector_only( + semantic, + _make_vector_rows([0.9]), + AsyncMock(return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)}), + ) assert len(results) == 1 # Should use chunk text, not the full content assert results[0].matched_chunk_text == "chunk text for entity:0:0" assert results[0].matched_chunk_text != large_content + + +@pytest.mark.asyncio +async def test_unparseable_chunk_key_names_no_search_row(): + """A chunk whose key does not spell a search row is skipped rather than ranked.""" + semantic = vector_semantic() + rows = [ + HydratedChunk(entity_id=0, chunk_key="entity:0:0", chunk_text="good", similarity=0.9), + HydratedChunk(entity_id=0, chunk_key="garbage", chunk_text="bad", similarity=0.95), + ] + + results = await run_vector_only(semantic, rows, _index_rows(1)) + + assert [row.id for row in results] == [0] + assert results[0].matched_chunk_text == "good" diff --git a/tests/services/test_project_readiness.py b/tests/services/test_project_readiness.py index a7aebf5d7..277831cd9 100644 --- a/tests/services/test_project_readiness.py +++ b/tests/services/test_project_readiness.py @@ -16,7 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from basic_memory.repository.search_repository import create_search_repository -from basic_memory.repository.search_repository_base import VECTOR_HYDRATION_BATCH_SIZE +from basic_memory.repository.search_reader import VECTOR_HYDRATION_BATCH_SIZE from basic_memory.repository.embedding_provider_factory import ( configured_embedding_provider_identity, )