diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py index 6b9583298..1b3ee59aa 100644 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ b/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" diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 1ddea1beb..b21f77b68 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/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 diff --git a/packages/shared-python/shared/tests/test_retrieval_search_channels.py b/packages/shared-python/shared/tests/test_retrieval_search_channels.py index e2e488a26..ddc31cd73 100644 --- a/packages/shared-python/shared/tests/test_retrieval_search_channels.py +++ b/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