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(