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..6fa86a68a 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,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" diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 1ddea1beb..b07529d8b 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 +# 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