From 9a5fc0688943d96cc11a43938fa2c5b20aca7e85 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 10 Aug 2026 22:58:52 -0500 Subject: [PATCH 1/3] fix(core): align reranker defaults and timeout config Signed-off-by: phernandez --- docs/semantic-search.md | 1 + src/basic_memory/config_models.py | 6 +++++ .../repository/fastembed_rerank_provider.py | 3 ++- .../repository/rerank_provider_factory.py | 16 +++++++++----- .../test_fastembed_rerank_provider.py | 7 ++++++ .../test_rerank_provider_factory.py | 22 +++++++++++++++++++ 6 files changed, 48 insertions(+), 7 deletions(-) diff --git a/docs/semantic-search.md b/docs/semantic-search.md index bce2af0fc..e41b39057 100644 --- a/docs/semantic-search.md +++ b/docs/semantic-search.md @@ -394,6 +394,7 @@ All settings use the `BASIC_MEMORY_` environment prefix: | `reranker_model` | `BASIC_MEMORY_RERANKER_MODEL` | `jinaai/jina-reranker-v1-tiny-en` | Model identifier. LiteLLM requires explicit `provider/model` routing. | | `reranker_candidates` | `BASIC_MEMORY_RERANKER_CANDIDATES` | `20` | Number of leading retrieval results rescored on every page. Larger values can improve recall but increase latency and provider usage. | | `reranker_max_document_chars` | `BASIC_MEMORY_RERANKER_MAX_DOCUMENT_CHARS` | `0` | Maximum characters sent per candidate. `0` sends the full matched text; a positive cap bounds latency and request size. | +| `reranker_timeout` | `BASIC_MEMORY_RERANKER_TIMEOUT` | `30.0` | Maximum seconds for each LiteLLM rerank request. FastEmbed runs locally and ignores this setting. | | `reranker_api_base` | `BASIC_MEMORY_RERANKER_API_BASE` | Unset | Optional custom endpoint for the LiteLLM provider. | | `reranker_api_key` | `BASIC_MEMORY_RERANKER_API_KEY` | Unset | Optional credential passed directly to LiteLLM. When unset, LiteLLM resolves provider credentials from its normal environment variables. | diff --git a/src/basic_memory/config_models.py b/src/basic_memory/config_models.py index 2211f3620..1c2d8b2ca 100644 --- a/src/basic_memory/config_models.py +++ b/src/basic_memory/config_models.py @@ -450,6 +450,12 @@ def __init__(self, **data: Any) -> None: ... "most-relevant matched chunk leads the text, so a modest cap keeps most of the signal.", ge=0, ) + reranker_timeout: float = Field( + default=30.0, + description="Maximum seconds allowed for each LiteLLM rerank request. " + "FastEmbed runs locally and ignores this setting.", + gt=0, + ) reranker_api_base: str | None = Field( default=None, description="Optional custom API base URL for the litellm reranker provider " diff --git a/src/basic_memory/repository/fastembed_rerank_provider.py b/src/basic_memory/repository/fastembed_rerank_provider.py index 039620713..c77202c34 100644 --- a/src/basic_memory/repository/fastembed_rerank_provider.py +++ b/src/basic_memory/repository/fastembed_rerank_provider.py @@ -10,6 +10,7 @@ from loguru import logger from requests import exceptions as requests_exceptions +from basic_memory.config_models import DEFAULT_FASTEMBED_RERANK_MODEL from basic_memory.repository.rerank_provider import validate_rerank_scores from basic_memory.repository.semantic_errors import ( RerankProviderContractError, @@ -68,7 +69,7 @@ class FastEmbedRerankProvider: def __init__( self, - model_name: str = "Xenova/ms-marco-MiniLM-L-6-v2", + model_name: str = DEFAULT_FASTEMBED_RERANK_MODEL, *, cache_dir: str | None = None, threads: int | None = None, diff --git a/src/basic_memory/repository/rerank_provider_factory.py b/src/basic_memory/repository/rerank_provider_factory.py index a9eeceb55..a453a343d 100644 --- a/src/basic_memory/repository/rerank_provider_factory.py +++ b/src/basic_memory/repository/rerank_provider_factory.py @@ -20,12 +20,12 @@ from basic_memory.repository.rerank_provider import RerankProvider # Key on the fields that change the loaded provider's identity: provider, model, -# (for the litellm path) the endpoint/key routing, and the resolved cache dir. The -# cache dir matters because the fastembed provider is constructed with it — omitting -# it (as an earlier version did) lets two configs with different cache dirs share one -# singleton pointing at the wrong directory, the #741/#872 class of bug the embedding -# factory guards against. CPU-derived thread counts stay out (they drift per call). -type RerankCacheKey = tuple[str, str, str | None, str | None, str] +# (for the litellm path) the endpoint/key routing and timeout, and the resolved cache +# dir. The cache dir matters because the fastembed provider is constructed with it — +# omitting it (as an earlier version did) lets two configs with different cache dirs +# share one singleton pointing at the wrong directory, the #741/#872 class of bug the +# embedding factory guards against. CPU-derived thread counts stay out (they drift per call). +type RerankCacheKey = tuple[str, str, str | None, str | None, float | None, str] _RERANK_PROVIDER_CACHE: dict[RerankCacheKey, RerankProvider] = {} _RERANK_PROVIDER_CACHE_LOCK = Lock() @@ -35,14 +35,17 @@ def _rerank_cache_key(app_config: BasicMemoryConfig) -> RerankCacheKey: provider_name = app_config.reranker_provider.strip().lower() api_base_digest = None api_key_digest = None + timeout = None if provider_name == "litellm": api_base_digest = _sensitive_value_digest(app_config.reranker_api_base) api_key_digest = _sensitive_value_digest(app_config.reranker_api_key) + timeout = app_config.reranker_timeout return ( provider_name, app_config.reranker_model, api_base_digest, api_key_digest, + timeout, _resolve_cache_dir(app_config), ) @@ -91,6 +94,7 @@ def create_rerank_provider(app_config: BasicMemoryConfig) -> RerankProvider | No model_name=app_config.reranker_model, api_key=app_config.reranker_api_key, api_base=app_config.reranker_api_base, + timeout=app_config.reranker_timeout, ) else: raise ValueError(f"Unsupported reranker provider: {provider_name}") diff --git a/tests/repository/test_fastembed_rerank_provider.py b/tests/repository/test_fastembed_rerank_provider.py index c3f8daec6..027778ce4 100644 --- a/tests/repository/test_fastembed_rerank_provider.py +++ b/tests/repository/test_fastembed_rerank_provider.py @@ -8,6 +8,7 @@ import pytest from requests import Response, exceptions as requests_exceptions +from basic_memory.config_models import DEFAULT_FASTEMBED_RERANK_MODEL from basic_memory.repository.fastembed_rerank_provider import FastEmbedRerankProvider from basic_memory.repository.semantic_errors import ( RerankProviderContractError, @@ -49,6 +50,12 @@ def _http_error(status_code: int) -> requests_exceptions.HTTPError: return requests_exceptions.HTTPError(f"HTTP {status_code}", response=response) +def test_constructor_uses_configured_default_model(): + provider = FastEmbedRerankProvider() + + assert provider.model_name == DEFAULT_FASTEMBED_RERANK_MODEL + + @pytest.mark.asyncio async def test_lazy_loads_once_and_reuses_model(monkeypatch): _install_stub(monkeypatch) diff --git a/tests/repository/test_rerank_provider_factory.py b/tests/repository/test_rerank_provider_factory.py index 62edeac95..3be01b9c6 100644 --- a/tests/repository/test_rerank_provider_factory.py +++ b/tests/repository/test_rerank_provider_factory.py @@ -88,12 +88,14 @@ def test_litellm_provider_selected_with_routing(): reranker_model="cohere/rerank-v3.5", reranker_api_key="secret", reranker_api_base="https://rerank.example", + reranker_timeout=12.5, ) provider = create_rerank_provider(config) assert isinstance(provider, LiteLLMRerankProvider) assert provider.model_name == "cohere/rerank-v3.5" assert provider._api_key == "secret" assert provider._api_base == "https://rerank.example" + assert provider._timeout == 12.5 def test_unsupported_provider_raises(): @@ -140,12 +142,32 @@ def test_distinct_cache_dir_does_not_collide(): assert b.cache_dir == "/tmp/rr-b" +def test_distinct_litellm_timeout_does_not_collide(): + """Two hosted configs differing only in timeout need distinct provider instances.""" + common = { + "reranker_enabled": True, + "reranker_provider": "litellm", + "reranker_model": "cohere/rerank-v3.5", + } + + fast_timeout = create_rerank_provider(_config(**common, reranker_timeout=5.0)) + slow_timeout = create_rerank_provider(_config(**common, reranker_timeout=45.0)) + + assert fast_timeout is not slow_timeout + + def test_reranker_enabled_requires_semantic_search(): """Config rejects reranking without semantic search rather than silently no-op'ing.""" with pytest.raises(ValidationError, match="requires semantic_search_enabled"): _config(reranker_enabled=True, semantic_search_enabled=False) +@pytest.mark.parametrize("timeout", [0, -1.0]) +def test_reranker_timeout_must_be_positive(timeout): + with pytest.raises(ValidationError, match="reranker_timeout"): + _config(reranker_timeout=timeout) + + def test_litellm_provider_rejects_default_fastembed_model(): """Selecting litellm without overriding the model is a footgun; reject it at config.""" with pytest.raises(ValidationError, match="requires an explicit reranker_model"): From 932247be447cadae301dbc809474747fac1ef1dc Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 10 Aug 2026 23:00:57 -0500 Subject: [PATCH 2/3] fix(mcp): preserve retryable search outages Signed-off-by: phernandez --- src/basic_memory/mcp/tools/chatgpt_tools.py | 14 +++++++++++- tests/mcp/tools/test_chatgpt_tools.py | 25 ++++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/src/basic_memory/mcp/tools/chatgpt_tools.py b/src/basic_memory/mcp/tools/chatgpt_tools.py index f539a54d0..e804fd60a 100644 --- a/src/basic_memory/mcp/tools/chatgpt_tools.py +++ b/src/basic_memory/mcp/tools/chatgpt_tools.py @@ -14,7 +14,7 @@ from basic_memory.mcp.client_info import is_openai_mcp_client from basic_memory.mcp.server import mcp from basic_memory.mcp.tools.read_note import read_note -from basic_memory.mcp.tools.search import search_notes +from basic_memory.mcp.tools.search import _SERVICE_UNAVAILABLE_HEADING, search_notes from basic_memory.schemas.search import SearchResponse, SearchResult @@ -182,6 +182,18 @@ async def search( ) if isinstance(results, str): + # Trigger: search_notes translated an API 503 into its retryable outage response. + # Why: OpenAI clients need to distinguish a temporary provider failure from an + # internal adapter error before deciding whether to retry. + # Outcome: preserve the retry signal in the Actions-compatible error payload. + if results.startswith(_SERVICE_UNAVAILABLE_HEADING): + return _text_content( + { + "results": [], + "error": "Search temporarily unavailable", + "error_message": "Search temporarily unavailable, retry shortly", + } + ) logger.warning(f"Search failed with error: {results[:100]}...") search_results = { "results": [], diff --git a/tests/mcp/tools/test_chatgpt_tools.py b/tests/mcp/tools/test_chatgpt_tools.py index 2ada010f4..3d0cb9b71 100644 --- a/tests/mcp/tools/test_chatgpt_tools.py +++ b/tests/mcp/tools/test_chatgpt_tools.py @@ -79,6 +79,29 @@ async def fake_search_notes_fn(*args, **kwargs): assert "error_details" in content +@pytest.mark.asyncio +async def test_search_retryable_outage_returns_explicit_retry_message( + monkeypatch, client, test_project, context_state +): + """The search_notes 503 shape remains retryable through the ChatGPT adapter.""" + import basic_memory.mcp.tools.chatgpt_tools as chatgpt_tools + + async def fake_search_notes_fn(*args, **kwargs): + return f"{chatgpt_tools._SERVICE_UNAVAILABLE_HEADING}\n\nReranker temporarily unavailable" + + monkeypatch.setattr(chatgpt_tools, "search_notes", fake_search_notes_fn) + + context = await _openai_mcp_context(context_state) + result = await chatgpt_tools.search("retryable query", context=context) + + content = json.loads(result[0]["text"]) + assert content == { + "results": [], + "error": "Search temporarily unavailable", + "error_message": "Search temporarily unavailable, retry shortly", + } + + @pytest.mark.asyncio async def test_search_uses_dynamic_default_search_type( monkeypatch, client, test_project, context_state @@ -315,7 +338,7 @@ async def boom(*args, **kwargs): assert isinstance(result, list) content = json.loads(result[0]["text"]) assert content["error"] == "Internal search error" - assert "error_message" in content + assert content["error_message"] == "boom" @pytest.mark.asyncio From c4fd2a84f652a2135e96b1fad8468aa5241f8a56 Mon Sep 17 00:00:00 2001 From: phernandez Date: Mon, 10 Aug 2026 23:30:56 -0500 Subject: [PATCH 3/3] test(core): close reranker phase-1 coverage gaps Signed-off-by: phernandez --- tests/repository/test_rerank_pipeline.py | 357 +++++++++++++++++------ 1 file changed, 267 insertions(+), 90 deletions(-) diff --git a/tests/repository/test_rerank_pipeline.py b/tests/repository/test_rerank_pipeline.py index 76b805954..dfc8a61d0 100644 --- a/tests/repository/test_rerank_pipeline.py +++ b/tests/repository/test_rerank_pipeline.py @@ -6,9 +6,18 @@ import pytest +import basic_memory.repository.postgres_search_repository as postgres_search_repository_module +import basic_memory.repository.search_repository as search_repository_module +import basic_memory.repository.sqlite_search_repository as sqlite_search_repository_module 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.rerank_provider import demote_tail_scores, validate_rerank_scores +from basic_memory.repository.rerank_provider import ( + build_rerank_document, + demote_tail_scores, + validate_rerank_scores, +) +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.semantic_errors import ( RerankProviderContractError, @@ -18,6 +27,8 @@ from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode +type BackendSearchRepository = SQLiteSearchRepository | PostgresSearchRepository + class _StubEmbeddingProvider: """Deterministic embeddings that give the two auth notes DIFFERENT similarity. @@ -59,9 +70,11 @@ class _FakeReranker: 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 @@ -90,7 +103,11 @@ class _ExplodingReranker: model_name = "boom" + def __init__(self) -> None: + self.calls = 0 + async def rerank(self, query: str, documents: list[str]) -> list[float]: + self.calls += 1 raise RerankTransientError("cross-encoder backend unreachable") def runtime_log_attrs(self) -> dict[str, Any]: @@ -485,22 +502,45 @@ async def test_rerank_paginate_surfaces_permanent_faults(exc): await repo._rerank_and_paginate("auth", [_row(id=1), _row(id=2)], offset=0, limit=10) -# --- End-to-end through the SQLite repo --- +# --- 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(), + ) -def _enable_semantic(repo: SQLiteSearchRepository) -> None: - try: - import sqlite_vec # noqa: F401 - except ImportError: # pragma: no cover - pytest.skip("sqlite-vec dependency is required for vector search tests.") - repo._semantic_enabled = True - repo._embedding_provider = _StubEmbeddingProvider() - repo._vector_dimensions = 4 - repo._vector_tables_initialized = False - repo._semantic_min_similarity = 0.0 +@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: SQLiteSearchRepository) -> None: + +async def _index_two_auth_notes(repo: BackendSearchRepository) -> None: await repo.init_search_index() await repo.bulk_index_items( [ @@ -525,17 +565,13 @@ async def _index_two_auth_notes(repo: SQLiteSearchRepository) -> None: @pytest.mark.asyncio -async def test_vector_search_applies_reranker(search_repository): - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) +async def test_vector_search_applies_reranker(rerank_search_repository): + await _index_two_auth_notes(rerank_search_repository) # Promote the note that vector similarity alone leaves tied/second. reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) - search_repository._rerank_provider = reranker + rerank_search_repository._rerank_provider = reranker - results = await search_repository.search( + results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.VECTOR, limit=5, @@ -552,21 +588,17 @@ async def test_vector_search_applies_reranker(search_repository): @pytest.mark.asyncio async def test_vector_search_expands_tail_from_stable_rerank_pool( - search_repository, + rerank_search_repository, monkeypatch, ): """Vector retrieval keeps its fixed prefix when a request also needs tail rows.""" - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) - search_repository._semantic_vector_k = 5 - search_repository._reranker_candidates = 2 - search_repository._rerank_provider = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) + await _index_two_auth_notes(rerank_search_repository) + rerank_search_repository._semantic_vector_k = 5 + rerank_search_repository._reranker_candidates = 2 + rerank_search_repository._rerank_provider = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) candidate_limits: list[int] = [] - run_vector_query = search_repository._run_vector_query + run_vector_query = rerank_search_repository._run_vector_query async def record_vector_query( session: Any, @@ -576,9 +608,9 @@ async def record_vector_query( candidate_limits.append(candidate_limit) return await run_vector_query(session, query_embedding, candidate_limit) - monkeypatch.setattr(search_repository, "_run_vector_query", record_vector_query) + monkeypatch.setattr(rerank_search_repository, "_run_vector_query", record_vector_query) - results = await search_repository.search( + results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.VECTOR, limit=3, @@ -589,12 +621,8 @@ async def record_vector_query( @pytest.mark.asyncio -async def test_vector_slow_query_timing_includes_reranker(search_repository, monkeypatch): - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) +async def test_vector_slow_query_timing_includes_reranker(rerank_search_repository, monkeypatch): + await _index_two_auth_notes(rerank_search_repository) clock = {"now": 0.0} reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) original_rerank = reranker.rerank @@ -603,7 +631,7 @@ async def slow_rerank(query: str, documents: list[str]) -> list[float]: clock["now"] = 3.0 return await original_rerank(query, documents) - search_repository._rerank_provider = reranker + rerank_search_repository._rerank_provider = reranker monkeypatch.setattr(reranker, "rerank", slow_rerank) monkeypatch.setattr( "basic_memory.repository.search_repository_base.time.perf_counter", @@ -615,7 +643,7 @@ async def slow_rerank(query: str, documents: list[str]) -> list[float]: warning, ) - await search_repository.search( + await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.VECTOR, limit=5, @@ -627,17 +655,13 @@ async def slow_rerank(query: str, documents: list[str]) -> list[float]: @pytest.mark.asyncio -async def test_hybrid_search_reranks_once(search_repository): +async def test_hybrid_search_reranks_once(rerank_search_repository): """Hybrid reranks the fused result exactly once — not again inside its vector leg.""" - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) + await _index_two_auth_notes(rerank_search_repository) reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) - search_repository._rerank_provider = reranker + rerank_search_repository._rerank_provider = reranker - results = await search_repository.search( + results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=5, @@ -652,21 +676,17 @@ async def test_hybrid_search_reranks_once(search_repository): @pytest.mark.asyncio async def test_hybrid_search_preserves_candidate_windows( - search_repository, + rerank_search_repository, monkeypatch, ): """Hybrid preserves legacy recall unless reranking owns the shared candidate pool.""" - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) - search_repository._semantic_vector_k = 100 - search_repository._reranker_candidates = 100 - search_repository._rerank_provider = None + await _index_two_auth_notes(rerank_search_repository) + rerank_search_repository._semantic_vector_k = 100 + rerank_search_repository._reranker_candidates = 100 + rerank_search_repository._rerank_provider = None candidate_limits: list[int] = [] - run_vector_query = search_repository._run_vector_query + run_vector_query = rerank_search_repository._run_vector_query async def record_vector_query( session: Any, @@ -676,9 +696,9 @@ async def record_vector_query( candidate_limits.append(candidate_limit) return await run_vector_query(session, query_embedding, candidate_limit) - monkeypatch.setattr(search_repository, "_run_vector_query", record_vector_query) + monkeypatch.setattr(rerank_search_repository, "_run_vector_query", record_vector_query) - baseline_results = await search_repository.search( + baseline_results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=10, @@ -688,11 +708,11 @@ async def record_vector_query( assert candidate_limits == [1000] candidate_limits.clear() - search_repository._semantic_vector_k = 5 - search_repository._reranker_candidates = 20 + rerank_search_repository._semantic_vector_k = 5 + rerank_search_repository._reranker_candidates = 20 reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) - search_repository._rerank_provider = reranker - reranked_results = await search_repository.search( + rerank_search_repository._rerank_provider = reranker + reranked_results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=11, @@ -702,7 +722,7 @@ async def record_vector_query( assert candidate_limits == [80] candidate_limits.clear() - growing_prefix_results = await search_repository.search( + growing_prefix_results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=21, @@ -786,17 +806,13 @@ async def deep_page(offset: int) -> list[SearchIndexRow]: @pytest.mark.asyncio -async def test_hybrid_search_surfaces_transient_reranker_error(search_repository): +async def test_hybrid_search_surfaces_transient_reranker_error(rerank_search_repository): """Hybrid search must not replace reranked order with raw order during an outage.""" - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) - search_repository._rerank_provider = _ExplodingReranker() + await _index_two_auth_notes(rerank_search_repository) + rerank_search_repository._rerank_provider = _ExplodingReranker() with pytest.raises(RerankTransientError, match="backend unreachable"): - await search_repository.search( + await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=5, @@ -804,19 +820,15 @@ async def test_hybrid_search_surfaces_transient_reranker_error(search_repository @pytest.mark.asyncio -async def test_hybrid_search_propagates_contract_error(search_repository): +async def test_hybrid_search_propagates_contract_error(rerank_search_repository): """A provider-contract break (e.g. incomplete rerank response) must surface, not hide.""" - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) - search_repository._rerank_provider = _PermanentFaultReranker( + await _index_two_auth_notes(rerank_search_repository) + rerank_search_repository._rerank_provider = _PermanentFaultReranker( RerankProviderContractError("incomplete rerank response") ) with pytest.raises(RerankProviderContractError): - await search_repository.search( + await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.HYBRID, limit=5, @@ -824,18 +836,183 @@ async def test_hybrid_search_propagates_contract_error(search_repository): @pytest.mark.asyncio -async def test_search_without_reranker_keeps_baseline(search_repository): - if not isinstance(search_repository, SQLiteSearchRepository): - pytest.skip("sqlite-vec repository behavior is local SQLite-only.") - - _enable_semantic(search_repository) - await _index_two_auth_notes(search_repository) - search_repository._rerank_provider = None +async def test_search_without_reranker_keeps_baseline(rerank_search_repository): + await _index_two_auth_notes(rerank_search_repository) + rerank_search_repository._rerank_provider = None - results = await search_repository.search( + results = await rerank_search_repository.search( search_text="auth session token", retrieval_mode=SearchRetrievalMode.VECTOR, limit=5, ) # Vector similarity ranks the plain-auth note above the "deep"-tilted one. assert [r.permalink for r in results] == ["specs/alpha", "specs/bravo"] + + +@pytest.mark.asyncio +async def test_fts_title_and_permalink_searches_never_call_reranker( + rerank_search_repository, +): + await _index_two_auth_notes(rerank_search_repository) + rerank_search_repository._rerank_provider = None + searches = { + "text": {"search_text": "auth", "retrieval_mode": SearchRetrievalMode.FTS}, + "title": {"title": "Alpha Auth Guide"}, + "permalink": {"permalink": "specs/alpha"}, + } + + baseline = { + name: await rerank_search_repository.search(**parameters) + for name, parameters in searches.items() + } + exploding_reranker = _ExplodingReranker() + rerank_search_repository._rerank_provider = exploding_reranker + guarded = { + name: await rerank_search_repository.search(**parameters) + for name, parameters in searches.items() + } + + assert all(baseline[name] for name in searches) + assert guarded == baseline + assert exploding_reranker.calls == 0 + + +@pytest.mark.asyncio +async def test_reranker_max_document_chars_config_reaches_provider( + session_maker, + test_project, + app_config, +): + max_chars = 12 + repository = _semantic_search_repository( + session_maker, + test_project.id, + app_config, + reranker_max_document_chars=max_chars, + ) + await _index_two_auth_notes(repository) + reranker = _FakeReranker({"Alpha": 0.1, "Bravo": 0.9}) + repository._rerank_provider = reranker + + await repository.search( + search_text="auth session token", + retrieval_mode=SearchRetrievalMode.VECTOR, + limit=5, + ) + + assert reranker.document_batches == [ + [ + build_rerank_document( + "Alpha Auth Guide", + "auth login session token overview", + max_chars, + ), + build_rerank_document( + "Bravo Auth Guide", + "auth login session token deep dive", + max_chars, + ), + ] + ] + + +@pytest.mark.parametrize( + ("backend", "expected_type"), + [ + (DatabaseBackend.SQLITE, SQLiteSearchRepository), + (DatabaseBackend.POSTGRES, PostgresSearchRepository), + ], +) +def test_create_search_repository_injects_reranker_for_both_backends( + monkeypatch, + backend, + expected_type, +): + reranker = _FakeReranker({}) + embedding_provider = _StubEmbeddingProvider() + vector_index = MagicMock() + config = BasicMemoryConfig( + env="test", + projects={"test-project": "/tmp/test"}, + default_project="test-project", + database_backend=backend, + semantic_search_enabled=True, + ) + monkeypatch.setattr( + search_repository_module, + "create_embedding_provider", + lambda _config: embedding_provider, + ) + monkeypatch.setattr( + search_repository_module, + "create_semantic_vector_index", + lambda **_kwargs: ( + "pgvector" if backend == DatabaseBackend.POSTGRES else "sqlite-vec", + vector_index, + ), + ) + monkeypatch.setattr( + search_repository_module, + "create_rerank_provider", + lambda _config: reranker, + ) + + repository = create_search_repository( + MagicMock(), + project_id=1, + app_config=config, + database_backend=backend, + ) + + assert isinstance(repository, expected_type) + assert repository._rerank_provider is reranker + + +@pytest.mark.parametrize( + ("repository_type", "repository_module", "backend"), + [ + ( + SQLiteSearchRepository, + sqlite_search_repository_module, + DatabaseBackend.SQLITE, + ), + ( + PostgresSearchRepository, + postgres_search_repository_module, + DatabaseBackend.POSTGRES, + ), + ], +) +@pytest.mark.parametrize("semantic_search_enabled", [False, True]) +def test_repository_self_resolves_reranker_only_when_semantic_search_is_enabled( + monkeypatch, + repository_type, + repository_module, + backend, + semantic_search_enabled, +): + reranker = _FakeReranker({}) + resolver = MagicMock(return_value=reranker) + monkeypatch.setattr(repository_module, "create_rerank_provider", resolver) + config = BasicMemoryConfig( + env="test", + projects={"test-project": "/tmp/test"}, + default_project="test-project", + database_backend=backend, + semantic_search_enabled=semantic_search_enabled, + ) + + repository = repository_type( + MagicMock(), + project_id=1, + app_config=config, + embedding_provider=_StubEmbeddingProvider(), + vector_index=MagicMock(), + ) + + if semantic_search_enabled: + resolver.assert_called_once_with(config) + assert repository._rerank_provider is reranker + else: + resolver.assert_not_called() + assert repository._rerank_provider is None