fix: give each query lexeme its own BM25 candidate budget

apps/api/tests/contract/test_bm25_fts_prefilter_contract.py

@@ -12,6 +12,7 @@
 
 import pytest
 import pytest_asyncio
+from shared.core.config import settings as channel_settings
 from shared.services.retrieval.search.channels import content_channel, path_channel
 from shared.testing.contract_runtime import PostgreSQLProcess
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
@@ -53,6 +54,7 @@
 """
 
 _NOISE_ROWS = 300
+_COMMON_TERM_ROWS = 1000
 
 
 @pytest_asyncio.fixture
@@ -197,3 +199,86 @@ async def test_exclusions_still_apply_under_the_prefilter(
         exclude_sections=[],
     )
     assert rows == []
+
+
+@pytest_asyncio.fixture
+async def rare_term_session(
+    postgresql_proc: PostgreSQLProcess,
+) -> AsyncGenerator[AsyncSession, None]:
+    """A corpus larger than the candidate budget where one chunk holds a rare term.
+
+    ts_rank_cd scores term density inside a chunk and ignores corpus-wide
+    rarity, so the rare-term chunk sorts last under a single global ordering
+    even though BM25 ranks it first.
+    """
+    dsn = (
+        f"postgresql+asyncpg://{postgresql_proc.user}@"
+        f"{postgresql_proc.host}:{postgresql_proc.port}/postgres"
+    )
+    engine = create_async_engine(dsn, isolation_level="AUTOCOMMIT")
+    async with engine.begin() as conn:
+        await conn.execute(text("DROP SCHEMA IF EXISTS bm25_rare CASCADE"))
+        await conn.execute(text("CREATE SCHEMA bm25_rare"))
+        await conn.execute(text("SET search_path TO bm25_rare"))
+        for statement in filter(None, (s.strip() for s in _SCHEMA.split(";"))):
+            await conn.execute(text(statement))
+        await conn.execute(text("INSERT INTO job_results VALUES (1, 'job1')"))
+        await conn.execute(
+            text(
+                "INSERT INTO documents VALUES "
+                "('d1', 'u1', 'ns1', 'active', 1, 'sample.pdf')"
+            )
+        )
+        await conn.execute(text("INSERT INTO document_sections VALUES ('s1', '/root')"))
+        await conn.execute(
+            text(
+                "INSERT INTO document_chunks "
+                "(chunk_id, document_id, section_id, chunk_type, content, "
+                " job_result_id, sort_order, content_search_text, path_search_text) "
+                "SELECT 'common-' || i, 'd1', 's1', 'text', 'body', 1, i, "
+                "       'data data data data data filler ' || i, 'p ' || i "
+                "FROM generate_series(1, :common) AS i"
+            ),
+            {"common": _COMMON_TERM_ROWS},
+        )
+        await conn.execute(
+            text(
+                "INSERT INTO document_chunks "
+                "(chunk_id, document_id, section_id, chunk_type, content, "
+                " job_result_id, sort_order, content_search_text, path_search_text) "
+                "VALUES ('rare-zebra', 'd1', 's1', 'text', 'body', 1, 0, "
+                "        'zebra', 'p rare')"
+            )
+        )
+
+    session_factory = async_sessionmaker(engine, expire_on_commit=False)
+    async with session_factory() as session:
+        await session.execute(text("SET search_path TO bm25_rare"))
+        yield session
+    await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_rare_term_chunk_survives_a_saturated_candidate_budget(
+    rare_term_session: AsyncSession,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # Budget well under the number of matching chunks, so the pool saturates.
+    monkeypatch.setattr(
+        channel_settings, "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", 200, raising=False
+    )
+
+    rows = await content_channel(
+        rare_term_session,
+        user_id="u1",
+        namespace="ns1",
+        query="data zebra",
+        top_k=5,
+        exclude_document_ids=[],
+        exclude_sections=[],
+    )
+
+    # BM25 weights the rare term far above the common one, so the chunk holding
+    # it belongs at the top. A single global ts_rank_cd ordering truncates it
+    # before BM25 ever sees it.
+    assert [str(row["chunk_id"]) for row in rows][0] == "rare-zebra"

packages/shared-python/shared/services/retrieval/search/channels.py

@@ -28,6 +28,11 @@
 # Guards against pathological queries producing an enormous tsquery.
 _MAX_FTS_QUERY_TOKENS = 50
 
+# Floor on each lexeme's share of the candidate budget. A query with many
+# lexemes would otherwise divide the budget down to a handful of rows each,
+# which throws away candidates the old global ordering would have kept.
+_MIN_CANDIDATES_PER_LEXEME = 50
+
 _TSV_FIELD_BY_SEARCH_FIELD = {
     "content_search_text": "content_search_tsv",
     "path_search_text": "path_search_tsv",
@@ -339,34 +344,52 @@ async def _bm25_channel(
     used_fallback = True
     if fts_tokens:
         # Postgres lexes the tokens with the same configuration that generated
-        # the tsvector columns, then ORs the resulting lexemes. Building the
-        # tsquery server-side keeps the prefilter aligned with the stored
-        # lexicon and leaves no room for tsquery syntax in user input to
-        # change the query shape. `fts_query.q` is NULL when no token yields a
-        # lexeme, which the caller treats as "no usable prefilter".
+        # the tsvector columns. Building the query server-side keeps the
+        # prefilter aligned with the stored lexicon and leaves no room for
+        # tsquery syntax in user input to change the query shape.
+        #
+        # Each lexeme draws candidates from its own slice of the budget rather
+        # than competing in one global ts_rank_cd ordering. ts_rank_cd scores
+        # term density within a chunk and ignores how rare a term is across the
+        # corpus, while BM25 weights rare terms heavily. A short chunk holding
+        # the one rare term in a query therefore sorts near the bottom of a
+        # global ordering and is truncated first, even though BM25 would rank
+        # it top. Giving every lexeme its own slice keeps those chunks in the
+        # pool. See #278.
         prefilter_sql = (
             corpus_cte
             + f""",
-    fts_query AS (
-        SELECT string_agg(quote_literal(lexeme), ' | ')::tsquery AS q
-        FROM (
-            SELECT DISTINCT
-                unnest(tsvector_to_array(to_tsvector('{_FTS_CONFIG}', token))) AS lexeme
-            FROM unnest(CAST(:fts_tokens AS text[])) AS token
-        ) lexemes
+    fts_lexemes AS (
+        SELECT DISTINCT
+            unnest(tsvector_to_array(to_tsvector('{_FTS_CONFIG}', token))) AS lexeme
+        FROM unnest(CAST(:fts_tokens AS text[])) AS token
+    ),
+    lexeme_budget AS (
+        SELECT
+            fl.lexeme,
+            to_tsquery('{_FTS_CONFIG}', quote_literal(fl.lexeme)) AS q,
+            GREATEST(
+                :fts_candidate_limit / GREATEST(COUNT(*) OVER (), 1),
+                :fts_min_per_lexeme
+            ) AS per_lexeme_limit
+        FROM fts_lexemes fl
     )
-    SELECT sc.*
-    FROM scoped_chunks sc, fts_query fq
-    WHERE COALESCE(sc.{search_field}, '') <> ''
-        AND fq.q IS NOT NULL
-        AND sc.{tsv_field} @@ fq.q
-    ORDER BY ts_rank_cd(sc.{tsv_field}, fq.q) DESC
-    LIMIT :fts_candidate_limit
+    SELECT DISTINCT ON (candidates.id) candidates.*
+    FROM lexeme_budget lb
+    CROSS JOIN LATERAL (
+        SELECT sc.*
+        FROM scoped_chunks sc
+        WHERE COALESCE(sc.{search_field}, '') <> ''
+            AND sc.{tsv_field} @@ lb.q
+        ORDER BY ts_rank_cd(sc.{tsv_field}, lb.q) DESC
+        LIMIT lb.per_lexeme_limit
+    ) candidates
     """
         )
         prefilter_params = dict(params)
         prefilter_params["fts_tokens"] = fts_tokens
         prefilter_params["fts_candidate_limit"] = candidate_limit
+        prefilter_params["fts_min_per_lexeme"] = _MIN_CANDIDATES_PER_LEXEME
         result = await db.execute(text(prefilter_sql), prefilter_params)
         rows = [_row_to_dict(r) for r in result.all()]
         used_fallback = not rows
@@ -385,15 +408,29 @@ async def _bm25_channel(
     ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field)
     ranked_rows = ranked_rows[:top_k]
 
+    duration_ms = (time.perf_counter() - started_at) * 1000
+    saturated = not used_fallback and candidate_count >= candidate_limit
     logger.debug(
-        "bm25_channel field={} candidates={} limit={} ranked={} fallback={} duration_ms={:.1f}",
+        "bm25_channel field={} candidates={} limit={} ranked={} "
+        "fallback={} saturated={} duration_ms={:.1f}",
         search_field,
         candidate_count,
         candidate_limit,
         len(ranked_rows),
         used_fallback,
-        (time.perf_counter() - started_at) * 1000,
+        saturated,
+        duration_ms,
     )
+    if saturated:
+        # The pool filled the budget, so chunks past it never reached BM25.
+        # Distinct from the healthy case, which the debug line alone cannot
+        # convey because both report candidates == limit.
+        logger.warning(
+            "bm25_channel candidate budget saturated field={} limit={}; "
+            "raise RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT if recall looks short",
+            search_field,
+            candidate_limit,
+        )
     return ranked_rows
 
 

packages/shared-python/shared/tests/test_retrieval_search_channels.py

@@ -97,13 +97,19 @@ async def test_content_channel_uses_bounded_or_fts_after_scope_filters(
     assert "sc.content_search_tsv @@" in sql
     assert "CAST(:fts_tokens AS text[])" in sql
     assert "ORDER BY ts_rank_cd" in sql
-    assert "LIMIT :fts_candidate_limit" in sql
-    assert sql.index("LOWER(dc.chunk_type)") < sql.index("LIMIT :fts_candidate_limit")
+    # The budget is derived from the configured limit and then spent per
+    # lexeme, so the bounding clause names the per-lexeme share rather than
+    # the setting directly.
+    assert ":fts_candidate_limit" in sql
+    assert "LIMIT lb.per_lexeme_limit" in sql
+    # Scope filters still land inside the CTE, ahead of any candidate bound,
+    # so excluded rows cannot consume the budget.
+    assert sql.index("LOWER(dc.chunk_type)") < sql.index("LIMIT lb.per_lexeme_limit")
     assert sql.index("LOWER(COALESCE(ds.section_path") < sql.index(
-        "LIMIT :fts_candidate_limit"
+        "LIMIT lb.per_lexeme_limit"
     )
     assert sql.index("POSITION(:_exc_section_path_0") < sql.index(
-        "LIMIT :fts_candidate_limit"
+        "LIMIT lb.per_lexeme_limit"
     )
     assert params["fts_tokens"] == ["alpha", "beta"]
     assert params["fts_candidate_limit"] == 7