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,16 @@
 """
 
 _NOISE_ROWS = 300
+_COMMON_TERM_ROWS = 1000
+_DENSE_ROWS_PER_TERM = 120
+_COVERING_QUERY_TERMS = (
+    "alpha",
+    "bravo",
+    "charlie",
+    "delta",
+    "echo",
+    "foxtrot",
+)
 
 
 @pytest_asyncio.fixture
@@ -197,3 +208,169 @@ 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"
+
+
+@pytest_asyncio.fixture
+async def covering_chunk_session(
+    postgresql_proc: PostgreSQLProcess,
+) -> AsyncGenerator[AsyncSession, None]:
+    """Dense single-term chunks plus one chunk covering every query term.
+
+    ts_rank_cd puts the covering chunk first because it rewards matching more
+    of the query, and BM25 agrees. Spending the whole budget per lexeme loses
+    it, since each lexeme's slice fills with denser single-term rows.
+    """
+    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_cover CASCADE"))
+        await conn.execute(text("CREATE SCHEMA bm25_cover"))
+        await conn.execute(text("SET search_path TO bm25_cover"))
+        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')"))
+        for term in _COVERING_QUERY_TERMS:
+            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 :term || '-' || i, 'd1', 's1', 'text', 'body', 1, i, "
+                    "       repeat(:term || ' ', 5) || 'filler ' || i, 'p ' || i "
+                    "FROM generate_series(1, :dense) AS i"
+                ),
+                {"term": term, "dense": _DENSE_ROWS_PER_TERM},
+            )
+        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 ('cover-all', 'd1', 's1', 'text', 'body', 1, 0, :covering, 'p cover')"
+            ),
+            {"covering": " ".join(_COVERING_QUERY_TERMS)},
+        )
+
+    session_factory = async_sessionmaker(engine, expire_on_commit=False)
+    async with session_factory() as session:
+        await session.execute(text("SET search_path TO bm25_cover"))
+        yield session
+    await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_chunk_covering_every_term_survives_a_saturated_budget(
+    covering_chunk_session: AsyncSession,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        channel_settings, "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", 200, raising=False
+    )
+
+    rows = await content_channel(
+        covering_chunk_session,
+        user_id="u1",
+        namespace="ns1",
+        query=" ".join(_COVERING_QUERY_TERMS),
+        top_k=5,
+        exclude_document_ids=[],
+        exclude_sections=[],
+    )
+
+    # The global ts_rank_cd slice is what keeps this chunk. Dropping it in
+    # favour of a purely per-lexeme budget would hand the top spot to a dense
+    # single-term chunk instead.
+    assert [str(row["chunk_id"]) for row in rows][0] == "cover-all"

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
 
+# Rows each lexeme contributes on top of the global slice. Small on purpose:
+# it exists to keep rare lexemes represented, not to replace the global
+# ordering, which is what ranks chunks covering several query terms.
+_CANDIDATES_PER_LEXEME_FLOOR = 50
+
 _TSV_FIELD_BY_SEARCH_FIELD = {
     "content_search_text": "content_search_tsv",
     "path_search_text": "path_search_tsv",
@@ -339,34 +344,70 @@ 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.
+        #
+        # Two pools are unioned. The global slice is the previous behaviour
+        # unchanged, ordered by ts_rank_cd over the whole query, which is what
+        # surfaces chunks covering several query terms. The per-lexeme floor
+        # adds a few rows for each lexeme on top, so a lexeme matching very
+        # little is still represented. ts_rank_cd scores term density within a
+        # chunk and ignores corpus-wide rarity, while BM25 weights rare terms
+        # heavily, so without the floor a short chunk holding the one rare term
+        # in a query is truncated before BM25 ever sees it. See #278.
+        #
+        # The union is a superset of the global slice, so nothing the previous
+        # ordering kept can be lost here.
         prefilter_sql = (
             corpus_cte
             + f""",
+    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
+    ),
     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
+        FROM fts_lexemes
+    ),
+    global_slice AS (
+        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
+    ),
+    lexeme_floor AS (
+        SELECT per_lexeme.*
+        FROM fts_lexemes fl
+        CROSS JOIN LATERAL (
+            SELECT sc.*
+            FROM scoped_chunks sc
+            WHERE COALESCE(sc.{search_field}, '') <> ''
+                AND sc.{tsv_field} @@ to_tsquery('{_FTS_CONFIG}', quote_literal(fl.lexeme))
+            ORDER BY
+                ts_rank_cd(
+                    sc.{tsv_field},
+                    to_tsquery('{_FTS_CONFIG}', quote_literal(fl.lexeme))
+                ) DESC
+            LIMIT :fts_lexeme_floor
+        ) per_lexeme
     )
-    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 (
+        SELECT * FROM global_slice
+        UNION ALL
+        SELECT * FROM lexeme_floor
+    ) candidates
     """
         )
         prefilter_params = dict(params)
         prefilter_params["fts_tokens"] = fts_tokens
         prefilter_params["fts_candidate_limit"] = candidate_limit
+        prefilter_params["fts_lexeme_floor"] = _CANDIDATES_PER_LEXEME_FLOOR
         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 +426,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