From de4b17c043a5425e0b251617a8b044848a134dc2 Mon Sep 17 00:00:00 2001 From: nuemaan Date: Mon, 10 Aug 2026 23:16:22 +0530 Subject: [PATCH] perf: prefilter BM25 channel candidates with Postgres FTS path_channel and content_channel loaded every scoped chunk with non-empty search text into Python before ranking, so cost grew with namespace size on every retrieve step. term_channel already pushed its query filter into SQL; these two did not, even though the GIN-indexed content_search_tsv and path_search_tsv columns were already selected by the scoped-corpus CTE. Both channels now prefilter on the tsvector column and hand a bounded candidate pool to the existing Python BM25 ranker, which stays the final per-channel ranker. Ordering by ts_rank_cd only decides which candidates survive the limit. The tsquery is built inside SQL from a text[] of ranker tokens: Postgres lexes each token with the same simple configuration that generates the tsvector columns, then ORs the resulting lexemes. That keeps the prefilter aligned with the stored lexicon, and tsquery operators in user input stay data instead of becoming syntax. OR semantics match how the Python ranker admits a row on any token intersection, so an AND-style query cannot narrow recall. Falls back to the previous unfiltered scan when no token yields a lexeme or the prefilter matches nothing, so recall cannot regress. Candidate pool size is configurable via RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT, default 2000. Closes #195 --- .../test_bm25_fts_prefilter_contract.py | 201 ++++++++++++++++++ .../tests/unit/test_bm25_channel_tsquery.py | 45 ++++ .../shared/core/config/__init__.py | 2 + .../shared-python/shared/core/config/app.py | 2 + .../shared/core/config/retrieval.py | 15 ++ .../services/retrieval/search/channels.py | 99 ++++++++- 6 files changed, 360 insertions(+), 4 deletions(-) create mode 100644 apps/api/tests/contract/test_bm25_fts_prefilter_contract.py create mode 100644 apps/api/tests/unit/test_bm25_channel_tsquery.py create mode 100644 packages/shared-python/shared/core/config/retrieval.py diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py new file mode 100644 index 000000000..78a1f6c8f --- /dev/null +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -0,0 +1,201 @@ +"""Contract tests for the BM25 channel Postgres FTS prefilter. + +These run against a real Postgres so the prefilter is validated against the +same generated tsvector columns and GIN indexes production uses. A pure-Python +fake would not catch a mismatch between the query configuration and the one +the columns were generated with. +""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator + +import pytest +import pytest_asyncio +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 +from sqlalchemy import text + +_SCHEMA = """ +CREATE TABLE documents ( + document_id TEXT PRIMARY KEY, + user_id TEXT, + namespace TEXT, + status TEXT, + current_job_result_id INTEGER, + source_file_name TEXT +); +CREATE TABLE job_results (id INTEGER PRIMARY KEY, job_id TEXT); +CREATE TABLE document_sections (section_id TEXT PRIMARY KEY, section_path TEXT); +CREATE TABLE document_chunks ( + id SERIAL PRIMARY KEY, + chunk_id TEXT, + document_id TEXT, + section_id TEXT, + chunk_type TEXT, + content TEXT, + source_chunk_path TEXT, + file_path TEXT, + chunk_metadata JSONB, + job_result_id INTEGER, + sort_order INTEGER, + content_search_text TEXT, + content_search_tsv TSVECTOR GENERATED ALWAYS AS + (to_tsvector('simple', COALESCE(content_search_text, ''))) STORED, + path_search_text TEXT, + path_search_tsv TSVECTOR GENERATED ALWAYS AS + (to_tsvector('simple', COALESCE(path_search_text, ''))) STORED, + term_search_text TEXT +); +CREATE INDEX idx_chunk_content_search_tsv ON document_chunks USING GIN (content_search_tsv); +CREATE INDEX idx_chunk_path_search_tsv ON document_chunks USING GIN (path_search_tsv); +""" + +_NOISE_ROWS = 300 + + +@pytest_asyncio.fixture +async def seeded_session( + postgresql_proc: PostgreSQLProcess, +) -> AsyncGenerator[AsyncSession, None]: + 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_fts CASCADE")) + await conn.execute(text("CREATE SCHEMA bm25_fts")) + await conn.execute(text("SET search_path TO bm25_fts")) + 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) " + "VALUES " + "('hit-en', 'd1', 's1', 'text', 'body', 1, 1, " + " 'alpha beta gamma', 'invoices alpha'), " + "('hit-cjk', 'd1', 's1', 'text', 'body', 1, 2, " + " '合同 条款 甲方', '合同 目录')" + ) + ) + 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 'noise-' || i, 'd1', 's1', 'text', 'body', 1, i + 10, " + " 'filler unrelated wording ' || i, 'misc path ' || i " + "FROM generate_series(1, :noise) AS i" + ), + {"noise": _NOISE_ROWS}, + ) + + session_factory = async_sessionmaker(engine, expire_on_commit=False) + async with session_factory() as session: + await session.execute(text("SET search_path TO bm25_fts")) + yield session + await engine.dispose() + + +async def _content_hits(session: AsyncSession, query: str) -> list[str]: + rows = await content_channel( + session, + user_id="u1", + namespace="ns1", + query=query, + top_k=50, + exclude_document_ids=[], + exclude_sections=[], + ) + return [str(row["chunk_id"]) for row in rows] + + +@pytest.mark.asyncio +async def test_content_channel_returns_only_query_matching_chunks( + seeded_session: AsyncSession, +) -> None: + # The corpus holds hundreds of unrelated chunks. Before the prefilter every + # one of them was loaded into Python for BM25 scoring. + assert await _content_hits(seeded_session, "alpha") == ["hit-en"] + + +@pytest.mark.asyncio +async def test_content_channel_matches_cjk_tokens( + seeded_session: AsyncSession, +) -> None: + assert await _content_hits(seeded_session, "合同") == ["hit-cjk"] + + +@pytest.mark.asyncio +async def test_content_channel_uses_or_semantics_across_tokens( + seeded_session: AsyncSession, +) -> None: + # A row matching any single query token must survive, matching how the + # Python BM25 ranker admits rows. + hits = await _content_hits(seeded_session, "alpha 合同") + assert sorted(hits) == ["hit-cjk", "hit-en"] + + +@pytest.mark.asyncio +async def test_tsquery_operators_in_query_do_not_change_filter_shape( + seeded_session: AsyncSession, +) -> None: + # Tokens are lexed by Postgres as data. If operators leaked into tsquery + # syntax, "alpha & zzzz" would AND and drop the row. + assert await _content_hits(seeded_session, "alpha & zzzz") == ["hit-en"] + assert await _content_hits(seeded_session, "!alpha") == ["hit-en"] + + +@pytest.mark.asyncio +async def test_query_matching_nothing_returns_no_rows( + seeded_session: AsyncSession, +) -> None: + # The fallback re-runs the unfiltered scan, and BM25 then scores no row + # above zero, so the channel still yields nothing. + assert await _content_hits(seeded_session, "zzzznomatch") == [] + + +@pytest.mark.asyncio +async def test_path_channel_prefilters_on_path_search_tsv( + seeded_session: AsyncSession, +) -> None: + rows = await path_channel( + seeded_session, + user_id="u1", + namespace="ns1", + query="invoices", + top_k=50, + exclude_document_ids=[], + exclude_sections=[], + ) + assert [str(row["chunk_id"]) for row in rows] == ["hit-en"] + + +@pytest.mark.asyncio +async def test_exclusions_still_apply_under_the_prefilter( + seeded_session: AsyncSession, +) -> None: + rows = await content_channel( + seeded_session, + user_id="u1", + namespace="ns1", + query="alpha", + top_k=50, + exclude_document_ids=["d1"], + exclude_sections=[], + ) + assert rows == [] diff --git a/apps/api/tests/unit/test_bm25_channel_tsquery.py b/apps/api/tests/unit/test_bm25_channel_tsquery.py new file mode 100644 index 000000000..aea2d9dfa --- /dev/null +++ b/apps/api/tests/unit/test_bm25_channel_tsquery.py @@ -0,0 +1,45 @@ +"""Unit tests for BM25 channel Postgres FTS prefilter token preparation.""" + +from __future__ import annotations + +from shared.services.retrieval.search.channels import ( + _MAX_FTS_QUERY_TOKENS, + _prepare_fts_tokens, +) + + +def test_returns_empty_for_no_tokens() -> None: + assert _prepare_fts_tokens([]) == [] + + +def test_keeps_token_order() -> None: + assert _prepare_fts_tokens(["alpha", "beta"]) == ["alpha", "beta"] + + +def test_strips_surrounding_whitespace() -> None: + assert _prepare_fts_tokens([" alpha ", "beta"]) == ["alpha", "beta"] + + +def test_drops_blank_tokens() -> None: + assert _prepare_fts_tokens(["", " ", "alpha"]) == ["alpha"] + + +def test_returns_empty_when_every_token_is_blank() -> None: + assert _prepare_fts_tokens(["", " "]) == [] + + +def test_caps_token_count() -> None: + tokens = [f"tok{index}" for index in range(_MAX_FTS_QUERY_TOKENS + 25)] + assert len(_prepare_fts_tokens(tokens)) == _MAX_FTS_QUERY_TOKENS + + +def test_preserves_cjk_tokens() -> None: + assert _prepare_fts_tokens(["合同", "条款"]) == ["合同", "条款"] + + +def test_passes_tsquery_operators_through_untouched() -> None: + # Tokens travel to Postgres as a text[] parameter and are lexed there, so + # operator characters are data rather than syntax. Nothing is escaped or + # dropped here. + raw = ["alpha' & 'zzzz", "!beta", "a|b"] + assert _prepare_fts_tokens(raw) == raw diff --git a/packages/shared-python/shared/core/config/__init__.py b/packages/shared-python/shared/core/config/__init__.py index 6889df8a5..b363997ae 100644 --- a/packages/shared-python/shared/core/config/__init__.py +++ b/packages/shared-python/shared/core/config/__init__.py @@ -19,6 +19,7 @@ from .mineru import MineruConfig from .qstash import QStashConfig from .redis import RedisConfig, RedisConfigManager, RedisPoolManager +from .retrieval import RetrievalConfig from .storage import StorageConfig __all__ = [ @@ -29,6 +30,7 @@ "RedisPoolManager", "CeleryConfig", "QStashConfig", + "RetrievalConfig", "StorageConfig", "JobConfig", "AIConfig", diff --git a/packages/shared-python/shared/core/config/app.py b/packages/shared-python/shared/core/config/app.py index 5ce5f9396..9b03f7889 100644 --- a/packages/shared-python/shared/core/config/app.py +++ b/packages/shared-python/shared/core/config/app.py @@ -13,6 +13,7 @@ from .mineru import MineruConfig from .qstash import QStashConfig from .redis import RedisConfig, RedisConfigManager, RedisPoolManager +from .retrieval import RetrievalConfig from .storage import StorageConfig @@ -27,6 +28,7 @@ class AppConfig( MineruConfig, BillingConfig, JobConfig, + RetrievalConfig, ): """Application configuration — all config components merged.""" diff --git a/packages/shared-python/shared/core/config/retrieval.py b/packages/shared-python/shared/core/config/retrieval.py new file mode 100644 index 000000000..14fc5a20d --- /dev/null +++ b/packages/shared-python/shared/core/config/retrieval.py @@ -0,0 +1,15 @@ +"""Retrieval configuration settings""" + +from pydantic import Field +from pydantic_settings import BaseSettings + + +class RetrievalConfig(BaseSettings): + """Retrieval configuration settings""" + + RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT: int = Field( + default=2000, + ge=1, + description="Maximum rows the Postgres FTS prefilter returns per BM25 channel " + "before Python BM25 reranking. Larger values trade memory for recall.", + ) diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 8cad8a930..a674e3494 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -6,17 +6,32 @@ """ from __future__ import annotations +import time from typing import Any +from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession +from shared.core.config import settings from shared.services.retrieval.search.lexical_ranker import ( rank_rows_by_bm25, tokenize_query_for_ranker, ) from shared.services.retrieval.search.section_filters import is_excluded_section +# The generated tsvector columns are built with to_tsvector('simple', ...), so +# queries must use the same configuration or nothing matches. +_FTS_CONFIG = "simple" + +# Guards against pathological queries producing an enormous tsquery. +_MAX_FTS_QUERY_TOKENS = 50 + +_TSV_FIELD_BY_SEARCH_FIELD = { + "content_search_text": "content_search_tsv", + "path_search_text": "path_search_tsv", +} + _SCOPED_CORPUS_CTE = """ WITH scoped_chunks AS ( @@ -120,6 +135,23 @@ def _build_extra_filters( return '\n '.join(clauses), params +def _prepare_fts_tokens(tokens: list[str]) -> list[str]: + """Return the ranker tokens to hand to the Postgres FTS prefilter. + + Tokens are passed to SQL as a text[] parameter and lexed by Postgres + itself, so nothing here needs to escape tsquery syntax. Only emptiness and + an upper bound are enforced. + """ + prepared: list[str] = [] + for token in tokens: + cleaned = token.strip() + if cleaned: + prepared.append(cleaned) + if len(prepared) >= _MAX_FTS_QUERY_TOKENS: + break + return prepared + + def _row_to_dict(row: Any) -> dict[str, Any]: return dict(row._mapping) @@ -236,18 +268,77 @@ async def _bm25_channel( ) params.update(extra_params) - sql = _SCOPED_CORPUS_CTE.format(exclude_clause=exclude_clause, extra_filters=extra_sql) + f""" + corpus_cte = _SCOPED_CORPUS_CTE.format( + exclude_clause=exclude_clause, + extra_filters=extra_sql, + ) + full_scan_sql = corpus_cte + f""" SELECT sc.* FROM scoped_chunks sc WHERE COALESCE(sc.{search_field}, '') <> '' """ - result = await db.execute(text(sql), params) - rows = [_row_to_dict(r) for r in result.all()] + started_at = time.perf_counter() + tsv_field = _TSV_FIELD_BY_SEARCH_FIELD[search_field] + fts_tokens = _prepare_fts_tokens(query_tokens) + candidate_limit = int(settings.RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT) + + rows: list[dict[str, Any]] = [] + 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". + 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 + ) + 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 + """ + prefilter_params = dict(params) + prefilter_params["fts_tokens"] = fts_tokens + prefilter_params["fts_candidate_limit"] = candidate_limit + result = await db.execute(text(prefilter_sql), prefilter_params) + rows = [_row_to_dict(r) for r in result.all()] + used_fallback = not rows + + # No usable tsquery, or the prefilter matched nothing. Fall back to the + # full scoped scan so recall never regresses against the previous + # behaviour. + if used_fallback: + result = await db.execute(text(full_scan_sql), params) + rows = [_row_to_dict(r) for r in result.all()] + + candidate_count = len(rows) rows = _filter_excluded_sections(rows, exclude_sections) ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) - return ranked_rows[:top_k] + ranked_rows = ranked_rows[:top_k] + + logger.debug( + "bm25_channel field={} candidates={} limit={} ranked={} fallback={} duration_ms={:.1f}", + search_field, + candidate_count, + candidate_limit, + len(ranked_rows), + used_fallback, + (time.perf_counter() - started_at) * 1000, + ) + return ranked_rows async def term_channel(