From 1815370c40b66eb9d6667a519fdcc3a81cee3da0 Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 01:25:03 -0500 Subject: [PATCH 1/2] fix(core): reuse retrieval for stable rerank pagination Signed-off-by: phernandez --- .../repository/search_repository_base.py | 203 +++++------ .../repository/semantic_vector_index.py | 3 + .../repository/sqlite_vec_index.py | 18 +- test-int/semantic/test_semantic_coverage.py | 10 +- test-int/test_multi_project_search.py | 6 +- test-int/test_stable_rerank_pagination.py | 80 +++++ tests/repository/test_rerank_pipeline.py | 4 +- tests/repository/test_search_trace.py | 2 +- .../test_stable_rerank_pagination.py | 321 ++++++++++++++++++ 9 files changed, 539 insertions(+), 108 deletions(-) create mode 100644 test-int/test_stable_rerank_pagination.py create mode 100644 tests/repository/test_stable_rerank_pagination.py diff --git a/src/basic_memory/repository/search_repository_base.py b/src/basic_memory/repository/search_repository_base.py index 7e824a86e..bd16ce192 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, field, replace from datetime import datetime, timezone from typing import Any, Callable, Dict, List, Literal, Optional, cast @@ -77,6 +77,7 @@ StagedVectorDeletion as _StagedVectorDeletion, VectorChunkState, ) +from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex from basic_memory.runtime.storage import RUNTIME_MARKDOWN_CONTENT_TYPE from basic_memory.runtime.vector_sync import VectorSyncBatchResult from basic_memory.schemas.search import ( @@ -128,6 +129,14 @@ type StoredEmbeddingStatus = Literal["pending", "ready"] +@dataclass(slots=True) +class _VectorCandidateWindow: + """Request-local prefix rows and original chunk ranks needed by hybrid fusion.""" + + stable_rows: list[SearchIndexRow] = field(default_factory=list) + source_ranks: dict[SearchIndexKey, int] = field(default_factory=dict) + + @dataclass(frozen=True, slots=True) class ChunkManifestRow: """One persisted vector-chunk manifest row.""" @@ -580,10 +589,18 @@ async def _run_vector_query( 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 ( + isinstance(self._semantic_vector_index, SQLiteVecIndex) + and self._rerank_provider is not None + ): + matches = await self._semantic_vector_index.search( + query_embedding, limit=candidate_limit, stable_prefix=True + ) + else: + matches = await self._semantic_vector_index.search( + query_embedding, + limit=candidate_limit, + ) if trace is not None: trace.readiness = await read_manifest_readiness( session, @@ -768,8 +785,11 @@ async def _hydrate_vector_matches( "chunk_key": match.key.chunk_key, "chunk_text": chunks_by_key[match.key], "best_similarity": match.similarity, + "candidate_rank": match.candidate_rank + if match.candidate_rank is not None + else rank, } - for match in matches + for rank, match in enumerate(matches) if match.key in chunks_by_key ] if trace is not None: @@ -2461,6 +2481,7 @@ async def _search_vector_only( candidate_limit: int | None = None, _emit_observability_log: bool = True, _apply_rerank: bool = True, + _candidate_window: _VectorCandidateWindow | None = None, trace: SearchTraceCollector | None = None, ) -> List[SearchIndexRow]: """Run vector-only search returning chunk-level results. @@ -2470,7 +2491,9 @@ async def _search_vector_only( result, not collapsed into its parent entity. ``candidate_limit`` is supplied only by a composed retrieval stage that - already sized the shared candidate pool. + already sized the shared candidate pool. ``_candidate_window`` collects the + fixed prefix and original chunk ranks for hybrid fusion, without another + retrieval or hydration pass. It is owned by this request, never the repository. """ self._assert_semantic_available() await self._ensure_vector_tables() @@ -2570,7 +2593,9 @@ def _log_vector_summary() -> None: # 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: + stable_chunks: dict[SearchIndexKey, list[tuple[float, str]]] = {} + stable_limit = self._rerank_candidate_limit() if self._should_rerank(query_text) else 0 + for rank, row in enumerate(vector_rows): chunk_key = row.get("chunk_key", "") if "best_similarity" in row: similarity = float(row["best_similarity"]) @@ -2588,6 +2613,20 @@ def _log_vector_summary() -> None: 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)) + # Built-in indexes limit adapter matches before manifest hydration. A + # dropped match must still consume its original prefix slot. External + # adapters instead fill the requested window with live hydrated chunks. + prefix_rank = ( + int(row.get("candidate_rank", rank)) + if self._semantic_vector_index_name in _BUILT_IN_VECTOR_INDEX_NAMES + else rank + ) + if prefix_rank < stable_limit: + stable_chunks.setdefault(si_key, []).append((similarity, chunk_text)) + if _candidate_window is not None: + _candidate_window.source_ranks[si_key] = min( + _candidate_window.source_ranks.get(si_key, prefix_rank), prefix_rank + ) if not similarity_by_si_key: hydrate_ms = (time.perf_counter() - hydrate_start) * 1000 @@ -2671,63 +2710,40 @@ def _log_vector_summary() -> None: ) 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, + def materialize_chunks( + chunks: dict[SearchIndexKey, list[tuple[float, str]]], + ) -> list[SearchIndexRow]: + ranked_rows: list[SearchIndexRow] = [] + for si_key, si_chunks in chunks.items(): + row = search_index_rows.get(si_key) + if row is None or si_key not in similarity_by_si_key: + continue + si_chunks.sort(key=lambda chunk: chunk[0], reverse=True) + similarity = si_chunks[0][0] + if effective_min_similarity > 0.0 and similarity < effective_min_similarity: + continue + content_snippet = row.content_snippet or "" + if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT: + matched_chunk_text = content_snippet + else: + matched_chunk_text = "\n---\n".join( + text for _, text in si_chunks[:TOP_CHUNKS_PER_RESULT] + ) + 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) + return ranked_rows - ranked_rows.sort(key=lambda item: item.score or 0.0, reverse=True) + ranked_rows = materialize_chunks(chunks_by_si_key) + stable_rows = materialize_chunks(stable_chunks) if stable_limit else ranked_rows + if _candidate_window is not None: + _candidate_window.stable_rows.extend(stable_rows) 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, @@ -2852,9 +2868,6 @@ async def _search_hybrid( 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. @@ -2866,13 +2879,8 @@ async def _search_hybrid( 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) - ) + 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 @@ -2899,6 +2907,7 @@ async def _search_hybrid( fts_span.set_attribute("result_count", len(fts_results)) fts_ms = (time.perf_counter() - fts_start) * 1000 vector_start = time.perf_counter() + vector_window = _VectorCandidateWindow() vector_results = await self._search_vector_only( search_text=search_text, permalink=permalink, @@ -2921,6 +2930,7 @@ async def _search_hybrid( candidate_limit=candidate_limit if rerank_configured else None, _emit_observability_log=False, _apply_rerank=False, + _candidate_window=vector_window if rerank_configured else None, trace=trace, ) vector_ms = (time.perf_counter() - vector_start) * 1000 @@ -3019,7 +3029,7 @@ async def _search_hybrid( 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) + ranked = sorted(fused_scores.items(), key=lambda item: (-item[1], item[0])) fusion_span.set_attribute("result_count", len(ranked)) fusion_ms = (time.perf_counter() - fusion_start) * 1000 if trace is not None: @@ -3045,33 +3055,33 @@ def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: # 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: + if rerank_configured: 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, - ) + # Reconstruct the fixed fusion universe from each source's original + # prefix. New second-source evidence may affect the tail, but cannot + # change prefix scores, membership, or the text sent to the reranker. + stable_fts = fts_results[:stable_candidate_limit] + # Both backends order lexical results by score, so expansion keeps + # the same normalization maximum and already-gated prefix scores. + stable_fts_scores = { + (row.type, row.id): fts_scores[(row.type, row.id)] for row in stable_fts + } + stable_vectors = {(row.type, row.id): row for row in vector_window.stable_rows} + stable_sources = {(row.type, row.id): row for row in stable_fts} | stable_vectors + stable_scores: dict[SearchIndexKey, float] = {} + for key in stable_sources: + f = stable_fts_scores.get(key, 0.0) + v = (stable_vectors[key].score or 0.0) if key in stable_vectors else 0.0 + stable_scores[key] = max(v, f) + FUSION_BONUS * min(v, f) + stable_candidates = [ + replace(stable_sources[key], score=score) + for key, score in sorted( + stable_scores.items(), key=lambda item: (-item[1], item[0]) + )[:stable_candidate_limit] + ] 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] @@ -3079,13 +3089,16 @@ def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: # 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. + # rows by their earliest source rank, using chunk rank before vector + # collapse. A later chunk may be only the second unique document; + # its collapsed rank would wrongly move it ahead of an earlier page. expanded_tail.sort( key=lambda entry: ( min( fts_ranks.get(entry[0], candidate_limit), - vec_ranks.get(entry[0], candidate_limit), + vector_window.source_ranks.get(entry[0], candidate_limit) + if entry[0] in vec_scores + else candidate_limit, ), entry[0], ) @@ -3102,7 +3115,7 @@ def _materialize(entry: tuple[SearchIndexKey, float]) -> SearchIndexRow: 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: + if total_ms > 2500: logger.warning( "[SEMANTIC_SLOW_QUERY] Semantic query timing: project_id={project_id} " "retrieval_mode={retrieval_mode} query_length={query_length} " diff --git a/src/basic_memory/repository/semantic_vector_index.py b/src/basic_memory/repository/semantic_vector_index.py index e869fb0a8..8a1e922cb 100644 --- a/src/basic_memory/repository/semantic_vector_index.py +++ b/src/basic_memory/repository/semantic_vector_index.py @@ -53,6 +53,9 @@ class VectorMatch: key: VectorKey similarity: float + # Zero-based position before adapter-side filtering, when filtering consumes + # top-k slots (sqlite-vec). Other adapters use their returned match order. + candidate_rank: int | None = None @runtime_checkable diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index c3344c31c..c2d3b1e79 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -327,25 +327,34 @@ async def search( query: Sequence[float], *, limit: int, + stable_prefix: bool = False, ) -> list[VectorMatch]: if not query or limit <= 0: return [] validate_query_dimensions(self.scope, query) await self.initialize() vector_k = min(limit, SQLITE_VEC_MAX_K) + # vec0's equal-distance membership changes with k. Reranking needs one + # repeatable universe to reconstruct smaller prefixes without a second query. + # Reuse the adapter's existing ceiling; only the requested slots survive SQL. + query_k = str(SQLITE_VEC_MAX_K) if stable_prefix else ":vector_k" async with db.scoped_session(self._session_maker) as session: await self._ensure_loaded(session) result = await session.execute( text( "WITH vector_matches AS MATERIALIZED (" " SELECT rowid, distance, source_hash FROM search_vector_embeddings " - " WHERE embedding MATCH :query AND k = :vector_k" - ") " - "SELECT c.entity_id, c.chunk_key, vector_matches.distance " - "FROM vector_matches " + f" WHERE embedding MATCH :query AND k = {query_k}" + "), ranked_matches AS MATERIALIZED (" + " SELECT *, ROW_NUMBER() OVER (ORDER BY distance, rowid) - 1 AS candidate_rank" + " FROM vector_matches) " + "SELECT c.entity_id, c.chunk_key, vector_matches.distance, " + "vector_matches.candidate_rank " + "FROM ranked_matches AS vector_matches " "JOIN search_vector_chunks c ON c.id = vector_matches.rowid " "AND c.source_hash = vector_matches.source_hash " "WHERE c.project_id = :project_id " + "AND vector_matches.candidate_rank < :vector_k " "AND c.vector_index = 'sqlite-vec' " "AND c.embedding_status = 'ready' " "AND c.embedding_model = :embedding_identity " @@ -366,6 +375,7 @@ async def search( entity_id=int(row["entity_id"]), chunk_key=str(row["chunk_key"]), ), + candidate_rank=int(row["candidate_rank"]), similarity=max( 0.0, min(1.0, 1.0 - (float(row["distance"]) ** 2) / 2.0), diff --git a/test-int/semantic/test_semantic_coverage.py b/test-int/semantic/test_semantic_coverage.py index 0e4881089..920445965 100644 --- a/test-int/semantic/test_semantic_coverage.py +++ b/test-int/semantic/test_semantic_coverage.py @@ -211,7 +211,7 @@ async def record_vector_query( assert growing_prefix_results assert reranker.calls == 2 - assert candidate_limits == [90, 80] + assert candidate_limits == [90] candidate_limits.clear() large_page_with_probe = await search_service.search( @@ -226,7 +226,7 @@ async def record_vector_query( assert len(large_page_with_probe) == 101 assert reranker.calls == 3 - assert candidate_limits == [890, 80] + assert candidate_limits == [890] # A larger retrieval window must extend the same sequence rather than # reordering rows already exposed by an earlier deep page. @@ -243,7 +243,7 @@ async def record_vector_query( ) assert len(stable_window) == 40 - assert candidate_limits == [280, 80] + assert candidate_limits == [280] candidate_limits.clear() third_page = await search_service.search( @@ -256,7 +256,7 @@ async def record_vector_query( limit=10, offset=20, ) - assert candidate_limits == [180, 80] + assert candidate_limits == [180] candidate_limits.clear() fourth_page = await search_service.search( @@ -269,7 +269,7 @@ async def record_vector_query( limit=10, offset=30, ) - assert candidate_limits == [280, 80] + assert candidate_limits == [280] assert [row.permalink for row in third_page] == [row.permalink for row in stable_window[20:30]] assert [row.permalink for row in fourth_page] == [row.permalink for row in stable_window[30:40]] diff --git a/test-int/test_multi_project_search.py b/test-int/test_multi_project_search.py index a3c8d487d..81cff7999 100644 --- a/test-int/test_multi_project_search.py +++ b/test-int/test_multi_project_search.py @@ -193,9 +193,11 @@ def record(_conn, _cursor, statement, _params, _context, _many): @pytest.mark.asyncio @pytest.mark.parametrize("mode", list(SearchRetrievalMode)) +@pytest.mark.parametrize("reranker_enabled", [False, True]) async def test_global_pagination_is_independent_of_page_size( - corpus: Corpus, mode: SearchRetrievalMode + corpus: Corpus, mode: SearchRetrievalMode, reranker_enabled: bool ) -> None: + corpus.config.reranker_enabled = reranker_enabled repo = corpus.repository([project.id for project in corpus.projects[:2]]) query = SearchService.prepare_query(SearchQuery(text="nebula", retrieval_mode=mode)) assert query is not None @@ -209,6 +211,8 @@ async def test_global_pagination_is_independent_of_page_size( (r.project_id, r.type, r.id, r.score) for r in complete ] assert await repo.search(query, offset=100) == [] + # This database-scoped reader never invokes the single-project rerank pipeline. + assert corpus.provider.query_calls == (5 if mode != SearchRetrievalMode.FTS else 0) @pytest.mark.asyncio diff --git a/test-int/test_stable_rerank_pagination.py b/test-int/test_stable_rerank_pagination.py new file mode 100644 index 000000000..311e93a2f --- /dev/null +++ b/test-int/test_stable_rerank_pagination.py @@ -0,0 +1,80 @@ +"""Real API/MCP pagination over the same single-pass semantic repository.""" + +import json + +import pytest +from fastapi import FastAPI +from fastmcp import Client, FastMCP +from httpx import AsyncClient +from sqlalchemy import func, update +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from basic_memory import db + +from basic_memory.deps.services import get_search_service_v2_external +from basic_memory.models import Project +from basic_memory.services.search_service import SearchService +from tests.repository.test_rerank_pipeline import ( + BackendSearchRepository, + _FakeReranker, + rerank_search_repository as rerank_search_repository, +) +from tests.repository.test_stable_rerank_pagination import ( + pagination_repository as pagination_repository, +) + + +@pytest.fixture +def session_maker( + engine_factory: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], +) -> async_sessionmaker[AsyncSession]: + return engine_factory[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["vector", "hybrid"]) +async def test_api_and_mcp_keep_probe_and_page_boundaries( + pagination_repository: BackendSearchRepository, + search_service: SearchService, + app: FastAPI, + client: AsyncClient, + test_project: Project, + mcp_server: FastMCP, + mode: str, +) -> None: + v2_project_url = f"/v2/projects/{test_project.external_id}" + repo = pagination_repository + repo._rerank_provider = _FakeReranker({"Note 00": 0.8, "Note 01": 0.9}) + async with db.scoped_session(repo.session_maker) as session: + await session.execute( + update(Project).where(Project.id == test_project.id).values(last_indexed_at=func.now()) + ) + await session.commit() + search_service.repository = repo + app.dependency_overrides[get_search_service_v2_external] = lambda: search_service + query = {"text": "auth session token", "retrieval_mode": mode, "min_similarity": 0.5} + response = await client.request( + "QUERY", f"{v2_project_url}/search/", json=query, params={"page_size": 100} + ) + assert response.status_code == 200, response.text + expected = response.json()["results"] + pages = [] + async with Client(mcp_server) as mcp: + for page in range(1, (len(expected) + 4) // 5 + 2): + result = await mcp.call_tool( + "search_notes", + { + "project": test_project.name, + "query": "auth session token", + "search_type": mode, + "min_similarity": 0.5, + "page": page, + "page_size": 5, + "output_format": "json", + }, + ) + payload = json.loads(result.content[0].text) + assert payload["has_more"] == (page * 5 < len(expected)) + assert payload["total_is_exact"] is False + pages.extend(payload["results"]) + assert [row["permalink"] for row in pages] == [row["permalink"] for row in expected] diff --git a/tests/repository/test_rerank_pipeline.py b/tests/repository/test_rerank_pipeline.py index 97ca99783..368321964 100644 --- a/tests/repository/test_rerank_pipeline.py +++ b/tests/repository/test_rerank_pipeline.py @@ -620,7 +620,7 @@ async def record_vector_query( ) assert [row.permalink for row in results] == ["specs/bravo", "specs/alpha"] - assert candidate_limits == [18, 8] + assert candidate_limits == [18] @pytest.mark.asyncio @@ -795,7 +795,7 @@ async def record_vector_query( assert growing_prefix_results assert reranker.calls == 2 - assert candidate_limits == [90, 80] + assert candidate_limits == [90] @pytest.mark.asyncio diff --git a/tests/repository/test_search_trace.py b/tests/repository/test_search_trace.py index aa6f95a89..19b7aabb9 100644 --- a/tests/repository/test_search_trace.py +++ b/tests/repository/test_search_trace.py @@ -847,7 +847,7 @@ async def test_fts_vector_hybrid_and_rerank_trace_variants( ) assert isinstance(reranked_trace, VectorQueryTrace) assert reranked_trace.rerank is not None - assert reranked_trace.rerank.stable_pool_refetched is True + assert reranked_trace.rerank.stable_pool_refetched is False assert reranked_trace.rerank.entries[0].key == ("entity", 3) alpha_rerank = next( entry for entry in reranked_trace.rerank.entries if entry.key == ("entity", 1) diff --git a/tests/repository/test_stable_rerank_pagination.py b/tests/repository/test_stable_rerank_pagination.py new file mode 100644 index 000000000..0e87ce9e4 --- /dev/null +++ b/tests/repository/test_stable_rerank_pagination.py @@ -0,0 +1,321 @@ +"""Single-pass pagination through real SQLite/pgvector retrieval and public boundaries.""" + +from collections import Counter +from typing import Any + +import pytest +from logfire.testing import CaptureLogfire, capfire as capfire +from sqlalchemy import bindparam, text +from sqlalchemy.ext.asyncio import AsyncSession + +from basic_memory import db +from basic_memory.config import DatabaseBackend +from basic_memory.models import Entity +from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.search_trace import SearchTraceCollector +from basic_memory.schemas.search import SearchRetrievalMode +from tests.repository.test_rerank_pipeline import ( + BackendSearchRepository, + _FakeReranker, + _entity_row, + rerank_search_repository as rerank_search_repository, +) + + +@pytest.fixture +async def pagination_repository( + rerank_search_repository: BackendSearchRepository, +) -> BackendSearchRepository: + repo = rerank_search_repository + repo._semantic_vector_k = 8 + repo._reranker_candidates = 2 + repo._reranker_max_document_chars = 2000 + await repo.init_search_index() + for index in range(32): + # Distinct chunk passages straddle the fixed prefix. Some rows have only + # lexical evidence, others only vector evidence, and one is filter-rejected. + content = "auth session token " + ("deep " if index % 2 else "") + if index < 2: + content = "\n\n".join(f"{content}passage {n} " + "detail " * 160 for n in range(3)) + if index == 30: + content = "oauth related concepts without lexical terms" + row = _entity_row( + project_id=repo.project_id, + row_id=700 + index, + title=f"Note {index:02d}", + permalink=f"{'excluded' if index == 31 else 'notes'}/{index:02d}", + content=content, + ) + async with db.scoped_session(repo.session_maker) as session: + session.add( + Entity( + id=row.id, + project_id=repo.project_id, + title=row.title, + note_type="spec", + content_type="text/markdown", + permalink=row.permalink, + file_path=row.file_path, + entity_metadata={"status": "active"}, + ) + ) + await session.commit() + await repo.index_item(row) + if index != 29: + await repo.sync_entity_vectors(row.id) + return repo + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID]) +@pytest.mark.parametrize("enabled", [False, True]) +async def test_one_pass_stable_pages_and_fixed_reranker_documents( + pagination_repository: BackendSearchRepository, + mode: SearchRetrievalMode, + enabled: bool, + capfire: CaptureLogfire, +) -> None: + repo = pagination_repository + reranker = _FakeReranker({f"Note {i:02d}": (i + 1) / 33 for i in range(32)}) + repo._rerank_provider = reranker if enabled else None + if not enabled: + # The legacy disabled path has a best-effort chunk window. Keep this corpus + # within it while testing unchanged ranking and one retrieval on later pages. + repo._semantic_vector_k = 100 + + async def page(limit: int, offset: int = 0) -> list[SearchIndexRow]: + capfire.exporter.clear() + collector = SearchTraceCollector() + calls_before = reranker.calls + rows = await repo.search( + search_text="auth session token", + retrieval_mode=mode, + file_path_prefix="notes", + min_similarity=0.5, + limit=limit, + offset=offset, + trace=collector, + ) + names = Counter(span["name"] for span in capfire.exporter.exported_spans_as_dict()) + assert names["search.embed_query"] == names["search.vector_query"] == 1 + assert names["search.fts"] == int(mode == SearchRetrievalMode.HYBRID) + assert reranker.calls - calls_before == int(enabled and bool(rows)) + assert collector.stable_pool_refetched is False + assert all(row.project_id == repo.project_id for row in rows) + assert all(row.file_path.startswith("notes/") for row in rows) + return rows + + complete = await page(40) + assert len(complete) == (31 if mode == SearchRetrievalMode.HYBRID else 30) + expected = [(row.type, row.id) for row in complete] + for size in (1, 4, 13): + actual = [ + row for offset in range(0, len(complete), size) for row in await page(size, offset) + ] + assert [(row.type, row.id) for row in actual] == expected + assert len({(row.type, row.id) for row in actual}) == len(expected) + assert [row.score for row in actual] == [row.score for row in complete] + assert await page(4, 100) == [] + if enabled: + assert reranker.document_batches + assert all(batch == reranker.document_batches[0] for batch in reranker.document_batches) + assert all(len(doc) <= 2000 for doc in reranker.document_batches[0]) + + +@pytest.mark.asyncio +async def test_hybrid_tail_orders_by_chunk_admission_before_collapse( + pagination_repository: BackendSearchRepository, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A late vector-only row cannot jump ahead of previously returned FTS tail rows.""" + repo = pagination_repository + repo._rerank_provider = _FakeReranker({"Note": 0.8}) + # Model a multi-chunk result crowding the raw window. The second unique row + # appears only after 100 chunks, so its document rank (1) is not its admission rank. + chunks = [ + { + "chunk_key": f"entity:700:{i}", + "chunk_text": f"auth passage {i}", + "best_similarity": 1.0 - i / 1000, + } + for i in range(100) + ] + [{"chunk_key": "entity:730:0", "chunk_text": "oauth", "best_similarity": 0.8}] + + async def vector_query( + session: AsyncSession, + embedding: list[float], + candidate_limit: int, + **kwargs: Any, + ) -> list[dict[str, Any]]: + return chunks[:candidate_limit] + + monkeypatch.setattr(repo, "_run_vector_query", vector_query) + complete = await repo.search( + search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=40 + ) + assert len(complete) == 32 + pages = [ + row + for offset in range(len(complete)) + for row in await repo.search( + search_text="auth session token", + retrieval_mode=SearchRetrievalMode.HYBRID, + limit=1, + offset=offset, + ) + ] + assert [row.id for row in pages] == [row.id for row in complete] + assert pages[-1].id == 730 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID]) +async def test_fixed_prefix_keeps_its_original_matched_passages( + pagination_repository: BackendSearchRepository, + monkeypatch: pytest.MonkeyPatch, + mode: SearchRetrievalMode, +) -> None: + repo = pagination_repository + # Use a vector-only query for the hybrid case too: empty FTS must not change + # the fixed vector documents. A long note uses matched passages for reranking. + row = _entity_row( + project_id=repo.project_id, + row_id=730, + title="Original passage", + permalink="notes/30", + content="body " * 1000, + ) + await repo.index_item(row) + chunks = [ + { + "chunk_key": f"entity:700:{i}", + "chunk_text": f"first note {i}", + "best_similarity": 1.0 - i / 1000, + } + for i in range(7) + ] + [ + {"chunk_key": "entity:730:0", "chunk_text": "ORIGINAL PASSAGE", "best_similarity": 0.9}, + {"chunk_key": "entity:730:1", "chunk_text": "LATER PASSAGE", "best_similarity": 0.8}, + ] + + async def vector_query( + session: AsyncSession, + embedding: list[float], + candidate_limit: int, + **kwargs: Any, + ) -> list[dict[str, Any]]: + return chunks[:candidate_limit] + + monkeypatch.setattr(repo, "_run_vector_query", vector_query) + reranker = _FakeReranker({"Original passage": 0.9, "Note 00": 0.8}) + repo._rerank_provider = reranker + for size in (2, 20): + await repo.search(search_text="unmatchedquery", retrieval_mode=mode, limit=size) + assert reranker.document_batches[0] == reranker.document_batches[1] + assert any("ORIGINAL PASSAGE" in doc for doc in reranker.document_batches[0]) + assert all("LATER PASSAGE" not in doc for batch in reranker.document_batches for doc in batch) + + +@pytest.mark.asyncio +async def test_dropped_adapter_slots_do_not_refill_fixed_prefix( + pagination_repository: BackendSearchRepository, + monkeypatch: pytest.MonkeyPatch, +) -> None: + repo = pagination_repository + reranker = _FakeReranker({"Note 00": 0.8, "Note 30": 0.9}) + repo._rerank_provider = reranker + + async def vector_query( + session: AsyncSession, + embedding: list[float], + candidate_limit: int, + **kwargs: Any, + ) -> list[dict[str, Any]]: + return [ + { + "chunk_key": "entity:700:0", + "chunk_text": "kept", + "best_similarity": 1.0, + "candidate_rank": 0, + }, + { + "chunk_key": "entity:730:0", + "chunk_text": "later", + "best_similarity": 0.9, + "candidate_rank": 8, + }, + ] + + monkeypatch.setattr(repo, "_run_vector_query", vector_query) + rows = await repo.search(search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR, limit=3) + assert [row.id for row in rows] == [700, 730] + assert len(reranker.document_batches) == 1 + assert len(reranker.document_batches[0]) == 1 + assert "Note 00" in reranker.document_batches[0][0] + + +@pytest.mark.asyncio +async def test_prefix_threshold_uses_only_its_own_chunk_evidence( + pagination_repository: BackendSearchRepository, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep private vector-hook compatibility when its evidence is not score-ordered.""" + repo = pagination_repository + reranker = _FakeReranker({"Note": 0.8}) + repo._rerank_provider = reranker + chunks = [ + {"chunk_key": "entity:730:0", "chunk_text": "weak", "best_similarity": 0.4}, + *[ + {"chunk_key": f"entity:700:{i}", "chunk_text": "strong", "best_similarity": 0.9} + for i in range(7) + ], + {"chunk_key": "entity:730:1", "chunk_text": "later strong", "best_similarity": 0.8}, + ] + + async def vector_query( + session: AsyncSession, + embedding: list[float], + candidate_limit: int, + **kwargs: Any, + ) -> list[dict[str, Any]]: + return chunks[:candidate_limit] + + monkeypatch.setattr(repo, "_run_vector_query", vector_query) + rows = await repo.search( + search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR, limit=3, min_similarity=0.5 + ) + assert [row.id for row in rows] == [700, 730] + assert len(reranker.document_batches[0]) == 1 + + +@pytest.mark.asyncio +async def test_pending_manifest_slots_preserve_backend_prefix_membership( + pagination_repository: BackendSearchRepository, +) -> None: + repo = pagination_repository + reranker = _FakeReranker({"Note": 0.8}) + repo._rerank_provider = reranker + async with repo.session_maker() as session: + matches = await repo._run_vector_query(session, [1.0, 0.0, 0.0, 0.0], 8) + assert len(matches) == 8 + async with db.scoped_session(repo.session_maker) as session: + await session.execute( + text( + "UPDATE search_vector_chunks SET embedding_status = 'pending' " + "WHERE project_id = :project_id AND chunk_key IN :chunk_keys" + ).bindparams(bindparam("chunk_keys", expanding=True)), + { + "project_id": repo.project_id, + "chunk_keys": [row["chunk_key"] for row in matches[:7]], + }, + ) + await session.commit() + for size in (2, 20): + rows = await repo.search( + search_text="auth", retrieval_mode=SearchRetrievalMode.VECTOR, limit=size + ) + assert rows + assert reranker.document_batches[0] == reranker.document_batches[1] + # SQLite filters after nearest-neighbor selection; Postgres filters before it. + expected = 1 if repo._app_config.database_backend == DatabaseBackend.SQLITE else 2 + assert len(reranker.document_batches[0]) == expected From 2e1b9a05d2efcd65418df20b1ac99ed3d514691e Mon Sep 17 00:00:00 2001 From: phernandez Date: Tue, 15 Sep 2026 01:33:50 -0500 Subject: [PATCH 2/2] test(core): share rerank fixtures across standalone suites Signed-off-by: phernandez --- pyproject.toml | 4 + test-int/test_stable_rerank_pagination.py | 4 +- tests/repository/test_rerank_pipeline.py | 121 +----------- .../test_stable_rerank_pagination.py | 48 +---- tests/semantic_search_helpers.py | 172 ++++++++++++++++++ 5 files changed, 187 insertions(+), 162 deletions(-) create mode 100644 tests/semantic_search_helpers.py diff --git a/pyproject.toml b/pyproject.toml index 24936db0a..73624e453 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -192,6 +192,10 @@ style = "pep440" bump = true fallback-version = "0.0.0" +[tool.ty.environment] +# Shared test helpers use the same import root as pytest's pythonpath above. +extra-paths = ["tests"] + [tool.ty.rules] all = "error" diff --git a/test-int/test_stable_rerank_pagination.py b/test-int/test_stable_rerank_pagination.py index 311e93a2f..b5bb54b0c 100644 --- a/test-int/test_stable_rerank_pagination.py +++ b/test-int/test_stable_rerank_pagination.py @@ -14,12 +14,10 @@ from basic_memory.deps.services import get_search_service_v2_external from basic_memory.models import Project from basic_memory.services.search_service import SearchService -from tests.repository.test_rerank_pipeline import ( +from semantic_search_helpers import ( BackendSearchRepository, _FakeReranker, rerank_search_repository as rerank_search_repository, -) -from tests.repository.test_stable_rerank_pagination import ( pagination_repository as pagination_repository, ) diff --git a/tests/repository/test_rerank_pipeline.py b/tests/repository/test_rerank_pipeline.py index 368321964..a09799a09 100644 --- a/tests/repository/test_rerank_pipeline.py +++ b/tests/repository/test_rerank_pipeline.py @@ -29,66 +29,14 @@ 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 - -type BackendSearchRepository = SQLiteSearchRepository | PostgresSearchRepository - - -class _StubEmbeddingProvider: - """Deterministic embeddings that give the two auth notes DIFFERENT similarity. - - An "auth" doc containing "deep" is tilted slightly off the query axis, so vector - retrieval ranks the plain-auth note strictly above it. That makes the pre-rerank - baseline a real ordering (not a tie), so a rerank that promotes the lower note is - a genuine "recover the below-cutoff doc" scenario (#950), not a coin flip. - """ - - model_name = "stub" - dimensions = 4 - - async def embed_query(self, text: str) -> list[float]: - return self._vectorize(text) - - async def embed_documents(self, texts: list[str]) -> list[list[float]]: - return [self._vectorize(t) for t in texts] - - def runtime_log_attrs(self) -> dict[str, Any]: - return {} - - @staticmethod - def _vectorize(text: str) -> list[float]: - lowered = text.lower() - if "auth" not in lowered: - return [0.0, 0.0, 0.0, 1.0] - # Unit vectors; cos with the query axis [1,0,0,0] is 1.0 vs 0.9. - if "deep" in lowered: - return [0.9, 0.4358898943540674, 0.0, 0.0] - return [1.0, 0.0, 0.0, 0.0] - - -class _FakeReranker: - """Scores a document by the marker substring it contains; records call count.""" - - model_name = "fake-reranker" - - def __init__(self, score_by_marker: dict[str, float]): - self.score_by_marker = score_by_marker - self.calls = 0 - self.document_batches: list[list[str]] = [] - - async def rerank(self, query: str, documents: list[str]) -> list[float]: - self.calls += 1 - self.document_batches.append(documents) - scores = [] - for doc in documents: - score = 0.0 - for marker, value in self.score_by_marker.items(): - if marker in doc: - score = value - scores.append(score) - return scores - - def runtime_log_attrs(self) -> dict[str, Any]: - return {} +from semantic_search_helpers import ( + BackendSearchRepository, + _StubEmbeddingProvider, + _FakeReranker, + _entity_row, + _semantic_search_repository, + rerank_search_repository as rerank_search_repository, +) class _BadReranker: @@ -156,24 +104,6 @@ def test_validate_rerank_scores_rejects_non_numeric_value(): validate_rerank_scores(["not-a-score"], expected_count=1) -def _entity_row(*, project_id: int, row_id: int, title: str, permalink: str, content: str): - now = datetime.now(timezone.utc) - return SearchIndexRow( - project_id=project_id, - id=row_id, - type=SearchItemType.ENTITY.value, - title=title, - permalink=permalink, - file_path=f"{permalink}.md", - metadata={"note_type": "spec"}, - entity_id=row_id, - content_stems=content, - content_snippet=content, - created_at=now, - updated_at=now, - ) - - def _row(**overrides) -> SearchIndexRow: now = datetime.now(timezone.utc) base: dict[str, Any] = dict( @@ -508,41 +438,6 @@ async def test_rerank_paginate_surfaces_permanent_faults(exc): # --- End-to-end through both repository backends --- -def _semantic_search_repository( - session_maker: Any, - project_id: int, - app_config: BasicMemoryConfig, - **config_updates: object, -) -> BackendSearchRepository: - config = app_config.model_copy( - update={ - "semantic_search_enabled": True, - "semantic_min_similarity": 0.0, - **config_updates, - } - ) - repository_type = ( - PostgresSearchRepository - if config.database_backend == DatabaseBackend.POSTGRES - else SQLiteSearchRepository - ) - return repository_type( - session_maker, - project_id=project_id, - app_config=config, - embedding_provider=_StubEmbeddingProvider(), - ) - - -@pytest.fixture -def rerank_search_repository( - session_maker: Any, - test_project: Any, - app_config: BasicMemoryConfig, -) -> BackendSearchRepository: - return _semantic_search_repository(session_maker, test_project.id, app_config) - - async def _index_two_auth_notes(repo: BackendSearchRepository) -> None: await repo.init_search_index() await repo.bulk_index_items( diff --git a/tests/repository/test_stable_rerank_pagination.py b/tests/repository/test_stable_rerank_pagination.py index 0e87ce9e4..1b65376c7 100644 --- a/tests/repository/test_stable_rerank_pagination.py +++ b/tests/repository/test_stable_rerank_pagination.py @@ -10,62 +10,18 @@ from basic_memory import db from basic_memory.config import DatabaseBackend -from basic_memory.models import Entity from basic_memory.repository.search_index_row import SearchIndexRow from basic_memory.repository.search_trace import SearchTraceCollector from basic_memory.schemas.search import SearchRetrievalMode -from tests.repository.test_rerank_pipeline import ( +from semantic_search_helpers import ( BackendSearchRepository, _FakeReranker, _entity_row, + pagination_repository as pagination_repository, rerank_search_repository as rerank_search_repository, ) -@pytest.fixture -async def pagination_repository( - rerank_search_repository: BackendSearchRepository, -) -> BackendSearchRepository: - repo = rerank_search_repository - repo._semantic_vector_k = 8 - repo._reranker_candidates = 2 - repo._reranker_max_document_chars = 2000 - await repo.init_search_index() - for index in range(32): - # Distinct chunk passages straddle the fixed prefix. Some rows have only - # lexical evidence, others only vector evidence, and one is filter-rejected. - content = "auth session token " + ("deep " if index % 2 else "") - if index < 2: - content = "\n\n".join(f"{content}passage {n} " + "detail " * 160 for n in range(3)) - if index == 30: - content = "oauth related concepts without lexical terms" - row = _entity_row( - project_id=repo.project_id, - row_id=700 + index, - title=f"Note {index:02d}", - permalink=f"{'excluded' if index == 31 else 'notes'}/{index:02d}", - content=content, - ) - async with db.scoped_session(repo.session_maker) as session: - session.add( - Entity( - id=row.id, - project_id=repo.project_id, - title=row.title, - note_type="spec", - content_type="text/markdown", - permalink=row.permalink, - file_path=row.file_path, - entity_metadata={"status": "active"}, - ) - ) - await session.commit() - await repo.index_item(row) - if index != 29: - await repo.sync_entity_vectors(row.id) - return repo - - @pytest.mark.asyncio @pytest.mark.parametrize("mode", [SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID]) @pytest.mark.parametrize("enabled", [False, True]) diff --git a/tests/semantic_search_helpers.py b/tests/semantic_search_helpers.py new file mode 100644 index 000000000..02f8aed03 --- /dev/null +++ b/tests/semantic_search_helpers.py @@ -0,0 +1,172 @@ +"""Deterministic semantic providers and corpus fixtures shared by unit/integration suites.""" + +from datetime import datetime, timezone +from typing import Any + +import pytest + +from basic_memory import db +from basic_memory.config import BasicMemoryConfig, DatabaseBackend +from basic_memory.models import Entity +from basic_memory.repository.postgres_search_repository import PostgresSearchRepository +from basic_memory.repository.search_index_row import SearchIndexRow +from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository +from basic_memory.schemas.search import SearchItemType + + +type BackendSearchRepository = SQLiteSearchRepository | PostgresSearchRepository + + +class _StubEmbeddingProvider: + """Deterministic embeddings that give the two auth notes DIFFERENT similarity. + + An "auth" doc containing "deep" is tilted slightly off the query axis, so vector + retrieval ranks the plain-auth note strictly above it. That makes the pre-rerank + baseline a real ordering (not a tie), so a rerank that promotes the lower note is + a genuine "recover the below-cutoff doc" scenario (#950), not a coin flip. + """ + + model_name = "stub" + dimensions = 4 + + async def embed_query(self, text: str) -> list[float]: + return self._vectorize(text) + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + return [self._vectorize(t) for t in texts] + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + @staticmethod + def _vectorize(text: str) -> list[float]: + lowered = text.lower() + if "auth" not in lowered: + return [0.0, 0.0, 0.0, 1.0] + # Unit vectors; cos with the query axis [1,0,0,0] is 1.0 vs 0.9. + if "deep" in lowered: + return [0.9, 0.4358898943540674, 0.0, 0.0] + return [1.0, 0.0, 0.0, 0.0] + + +class _FakeReranker: + """Scores a document by the marker substring it contains; records call count.""" + + model_name = "fake-reranker" + + def __init__(self, score_by_marker: dict[str, float]): + self.score_by_marker = score_by_marker + self.calls = 0 + self.document_batches: list[list[str]] = [] + + async def rerank(self, query: str, documents: list[str]) -> list[float]: + self.calls += 1 + self.document_batches.append(documents) + scores = [] + for doc in documents: + score = 0.0 + for marker, value in self.score_by_marker.items(): + if marker in doc: + score = value + scores.append(score) + return scores + + def runtime_log_attrs(self) -> dict[str, Any]: + return {} + + +def _entity_row(*, project_id: int, row_id: int, title: str, permalink: str, content: str): + now = datetime.now(timezone.utc) + return SearchIndexRow( + project_id=project_id, + id=row_id, + type=SearchItemType.ENTITY.value, + title=title, + permalink=permalink, + file_path=f"{permalink}.md", + metadata={"note_type": "spec"}, + entity_id=row_id, + content_stems=content, + content_snippet=content, + created_at=now, + updated_at=now, + ) + + +def _semantic_search_repository( + session_maker: Any, + project_id: int, + app_config: BasicMemoryConfig, + **config_updates: object, +) -> BackendSearchRepository: + config = app_config.model_copy( + update={ + "semantic_search_enabled": True, + "semantic_min_similarity": 0.0, + **config_updates, + } + ) + repository_type = ( + PostgresSearchRepository + if config.database_backend == DatabaseBackend.POSTGRES + else SQLiteSearchRepository + ) + return repository_type( + session_maker, + project_id=project_id, + app_config=config, + embedding_provider=_StubEmbeddingProvider(), + ) + + +@pytest.fixture +def rerank_search_repository( + session_maker: Any, + test_project: Any, + app_config: BasicMemoryConfig, +) -> BackendSearchRepository: + return _semantic_search_repository(session_maker, test_project.id, app_config) + + +@pytest.fixture +async def pagination_repository( + rerank_search_repository: BackendSearchRepository, +) -> BackendSearchRepository: + repo = rerank_search_repository + repo._semantic_vector_k = 8 + repo._reranker_candidates = 2 + repo._reranker_max_document_chars = 2000 + await repo.init_search_index() + for index in range(32): + # Distinct chunk passages straddle the fixed prefix. Some rows have only + # lexical evidence, others only vector evidence, and one is filter-rejected. + content = "auth session token " + ("deep " if index % 2 else "") + if index < 2: + content = "\n\n".join(f"{content}passage {n} " + "detail " * 160 for n in range(3)) + if index == 30: + content = "oauth related concepts without lexical terms" + row = _entity_row( + project_id=repo.project_id, + row_id=700 + index, + title=f"Note {index:02d}", + permalink=f"{'excluded' if index == 31 else 'notes'}/{index:02d}", + content=content, + ) + async with db.scoped_session(repo.session_maker) as session: + session.add( + Entity( + id=row.id, + project_id=repo.project_id, + title=row.title, + note_type="spec", + content_type="text/markdown", + permalink=row.permalink, + file_path=row.file_path, + entity_metadata={"status": "active"}, + ) + ) + await session.commit() + await repo.index_item(row) + if index != 29: + await repo.sync_entity_vectors(row.id) + return repo