diff --git a/apps/api/.env.example b/apps/api/.env.example index e33e8bc37..33ed8d8ce 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -103,6 +103,7 @@ ARK_API_KEY= # RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy # 3-channel RRF mode. # RETRIEVAL_WORKFLOW_PLANNER_TIMEOUT_SECONDS=10.0 +# RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT=2000 # File handling defaults SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md,.html,.htm 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..6b9583298 --- /dev/null +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -0,0 +1,199 @@ +"""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/apps/worker/.env.example b/apps/worker/.env.example index 7a682e776..404340367 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -97,6 +97,7 @@ ARK_API_KEY= # evidence_text is the primary output and answer_text is always empty. Set # RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy # 3-channel RRF mode. +# RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT=2000 # Required for specific features: billing and analytics BILLING_ENABLED=false diff --git a/packages/shared-python/shared/core/config/__init__.py b/packages/shared-python/shared/core/config/__init__.py index 6889df8a5..5f07b4179 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__ = [ @@ -33,6 +34,7 @@ "JobConfig", "AIConfig", "MineruConfig", + "RetrievalConfig", "AppConfig", "app_config", "settings", 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..5484c2f67 --- /dev/null +++ b/packages/shared-python/shared/core/config/retrieval.py @@ -0,0 +1,17 @@ +"""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..1ddea1beb 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -4,19 +4,35 @@ Each channel queries the full scoped corpus independently and returns ranked rows. Channels are fused via RRF in the orchestrator. """ + 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 ( @@ -94,10 +110,10 @@ def _build_extra_filters( params: dict[str, Any] = {} if allowed_chunk_types is not None: - placeholders = ', '.join(f':_act_{i}' for i in range(len(allowed_chunk_types))) - clauses.append(f'AND LOWER(dc.chunk_type) IN ({placeholders})') + placeholders = ", ".join(f":_act_{i}" for i in range(len(allowed_chunk_types))) + clauses.append(f"AND LOWER(dc.chunk_type) IN ({placeholders})") for i, ct in enumerate(sorted(allowed_chunk_types)): - params[f'_act_{i}'] = ct + params[f"_act_{i}"] = ct if signal_paths: # TODO(intent-step): Current implementation uses OR across @@ -108,16 +124,75 @@ def _build_extra_filters( # vs "path_prefix" (for Intent Step resolved paths). ilike_parts = [] for i, kw in enumerate(signal_paths): - key = f'_sig_{i}' + key = f"_sig_{i}" ilike_parts.append(f"LOWER(COALESCE(ds.section_path, '')) LIKE :{key}") - params[key] = f'%{kw.lower()}%' - combined = ' OR '.join(ilike_parts) - if filter_mode == 'keep': - clauses.append(f'AND ({combined})') + params[key] = f"%{kw.lower()}%" + combined = " OR ".join(ilike_parts) + if filter_mode == "keep": + clauses.append(f"AND ({combined})") else: - clauses.append(f'AND NOT ({combined})') + clauses.append(f"AND NOT ({combined})") + + return "\n ".join(clauses), params + + +def _build_exclude_section_filters( + *, + exclude_sections: list[dict[str, str]], +) -> tuple[str, dict[str, Any]]: + """Exclude an exact section path and its descendants inside the scoped CTE. + + Applied before the FTS candidate LIMIT so excluded sections cannot consume + the bounded candidate budget. Sectionless chunks stay eligible (empty path + does not match), matching ``is_excluded_section``. + """ + clauses: list[str] = [] + params: dict[str, Any] = {} + + for index, item in enumerate(exclude_sections): + if not isinstance(item, dict): + continue + document_id = str(item.get("document_id") or "").strip() + section_path = str(item.get("section_path") or "").strip() + if not document_id or not section_path: + continue + + document_key = f"_exc_section_doc_{index}" + path_key = f"_exc_section_path_{index}" + clauses.append( + f"""AND NOT ( + dc.document_id = :{document_key} + AND ( + COALESCE(ds.section_path, '') = :{path_key} + OR POSITION(:{path_key} || ' / ' IN COALESCE(ds.section_path, '')) = 1 + ) + )""" + ) + params[document_key] = document_id + params[path_key] = section_path + + return "\n ".join(clauses), params + - return '\n '.join(clauses), params +def _join_sql_filters(*filters: str) -> str: + return "\n ".join(filter(None, filters)) + + +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]: @@ -131,7 +206,8 @@ def _filter_excluded_sections( if not exclude_sections: return rows return [ - row for row in rows + row + for row in rows if not is_excluded_section( document_id=row.get("document_id"), section_path=row.get("section_path"), @@ -151,7 +227,7 @@ async def path_channel( exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", ) -> list[dict[str, Any]]: """Path channel: BM25 over pre-tokenized path search text. @@ -184,7 +260,7 @@ async def content_channel( exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", ) -> list[dict[str, Any]]: """Content channel: BM25 over pre-tokenized content search text.""" return await _bm25_channel( @@ -229,25 +305,96 @@ async def _bm25_channel( signal_paths=signal_paths or [], filter_mode=filter_mode, ) + section_sql, section_params = _build_exclude_section_filters( + exclude_sections=exclude_sections, + ) + extra_sql = _join_sql_filters(extra_sql, section_sql) params = _build_base_params( user_id=user_id, namespace=namespace, exclude_document_ids=exclude_document_ids, ) params.update(extra_params) + params.update(section_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) + # Defensive: SQL already owns pre-LIMIT section exclusion. 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( @@ -261,7 +408,7 @@ async def term_channel( exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, signal_paths: list[str] | None = None, - filter_mode: str = 'delete', + filter_mode: str = "delete", ) -> list[dict[str, Any]]: """Term/grep channel: substring matching on term_search_text. @@ -299,12 +446,17 @@ async def term_channel( params["full_query"] = f"%{query_lower}%" where_clause = " OR ".join(ilike_conditions) - sql = _SCOPED_CORPUS_CTE.format(exclude_clause=exclude_clause, extra_filters=extra_sql) + f""" + sql = ( + _SCOPED_CORPUS_CTE.format( + exclude_clause=exclude_clause, extra_filters=extra_sql + ) + + f""" SELECT sc.* FROM scoped_chunks sc WHERE sc.term_search_text IS NOT NULL AND ({where_clause}) """ + ) result = await db.execute(text(sql), params) rows = [_row_to_dict(r) for r in result.all()] diff --git a/packages/shared-python/shared/tests/test_retrieval_search_channels.py b/packages/shared-python/shared/tests/test_retrieval_search_channels.py new file mode 100644 index 000000000..e2e488a26 --- /dev/null +++ b/packages/shared-python/shared/tests/test_retrieval_search_channels.py @@ -0,0 +1,216 @@ +"""Unit coverage for bounded PostgreSQL candidates in lexical channels.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from pytest import MonkeyPatch + +from shared.services.retrieval.search import channels + + +class _FakeRow: + def __init__(self, **values: Any) -> None: + self._mapping = values + + +class _FakeResult: + def __init__(self, rows: list[_FakeRow]) -> None: + self._rows = rows + + def all(self) -> list[_FakeRow]: + return self._rows + + +class _FakeSession: + def __init__(self, *result_sets: list[_FakeRow]) -> None: + self.result_sets = list(result_sets) + self.calls: list[tuple[str, dict[str, Any]]] = [] + + async def execute( + self, + statement: object, + params: dict[str, Any], + ) -> _FakeResult: + self.calls.append((str(statement), dict(params))) + return _FakeResult(self.result_sets.pop(0)) + + +class _FakeLogger: + def __init__(self) -> None: + self.messages: list[str] = [] + + def debug(self, message: str, *args: Any) -> None: + if args: + self.messages.append(message.format(*args)) + else: + self.messages.append(message) + + +def _row( + *, + id_row: int, + content_search_text: str = "alpha beta", + path_search_text: str = "root alpha", +) -> _FakeRow: + return _FakeRow( + id=f"row-{id_row}", + chunk_id=f"chunk-{id_row}", + document_id="doc-1", + section_id=f"section-{id_row}", + chunk_type="text", + content=f"content {id_row}", + content_search_text=content_search_text, + path_search_text=path_search_text, + section_path=f"Root / Section {id_row}", + ) + + +@pytest.mark.asyncio +async def test_content_channel_uses_bounded_or_fts_after_scope_filters( + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setattr( + channels.settings, + "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", + 7, + ) + db = _FakeSession([_row(id_row=1)]) + + rows = await channels.content_channel( + db, # type: ignore[arg-type] + user_id="user-1", + namespace="knowledge", + query="alpha beta", + top_k=3, + exclude_document_ids=["doc-old"], + exclude_sections=[{"document_id": "doc-1", "section_path": "Root / Hidden"}], + allowed_chunk_types={"text"}, + signal_paths=["Root"], + filter_mode="keep", + ) + + assert len(rows) == 1 + assert len(db.calls) == 1 + sql, params = db.calls[0] + 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") + assert sql.index("LOWER(COALESCE(ds.section_path") < sql.index( + "LIMIT :fts_candidate_limit" + ) + assert sql.index("POSITION(:_exc_section_path_0") < sql.index( + "LIMIT :fts_candidate_limit" + ) + assert params["fts_tokens"] == ["alpha", "beta"] + assert params["fts_candidate_limit"] == 7 + assert params["excluded_doc_ids"] == ["doc-old"] + assert params["_exc_section_doc_0"] == "doc-1" + assert params["_exc_section_path_0"] == "Root / Hidden" + + +@pytest.mark.asyncio +async def test_path_channel_uses_path_tsv_and_python_bm25_final_reranker( + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setattr( + channels.settings, + "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", + 4, + ) + candidate_rows = [_row(id_row=index) for index in range(8)] + db = _FakeSession(candidate_rows[:4]) + captured: dict[str, Any] = {} + + def fake_rank( + rows: list[dict[str, Any]], + query_tokens: list[str], + *, + search_field: str, + ) -> list[dict[str, Any]]: + captured["candidate_count"] = len(rows) + captured["query_tokens"] = query_tokens + captured["search_field"] = search_field + return list(reversed(rows)) + + monkeypatch.setattr(channels, "rank_rows_by_bm25", fake_rank) + + rows = await channels.path_channel( + db, # type: ignore[arg-type] + user_id="user-1", + namespace="knowledge", + query="alpha beta", + top_k=2, + exclude_document_ids=[], + exclude_sections=[], + ) + + sql, params = db.calls[0] + assert "sc.path_search_tsv @@" in sql + assert "sc.content_search_tsv @@" not in sql + assert params["fts_candidate_limit"] == 4 + assert captured == { + "candidate_count": 4, + "query_tokens": ["alpha", "beta"], + "search_field": "path_search_text", + } + assert [row["id"] for row in rows] == ["row-3", "row-2"] + + +@pytest.mark.asyncio +async def test_bm25_channel_uses_full_scan_fallback_and_logs_metrics( + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setattr( + channels.settings, + "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", + 5, + ) + db = _FakeSession([], [_row(id_row=1)]) + fake_logger = _FakeLogger() + monkeypatch.setattr(channels, "logger", fake_logger) + + rows = await channels.content_channel( + db, # type: ignore[arg-type] + user_id="user-1", + namespace="knowledge", + query="alpha", + top_k=2, + exclude_document_ids=[], + exclude_sections=[], + ) + + assert len(rows) == 1 + assert len(db.calls) == 2 + fts_sql, fts_params = db.calls[0] + fallback_sql, _fallback_params = db.calls[1] + assert "CAST(:fts_tokens AS text[])" in fts_sql + assert "CAST(:fts_tokens AS text[])" not in fallback_sql + assert "LIMIT :fts_candidate_limit" not in fallback_sql + assert fts_params["fts_candidate_limit"] == 5 + assert len(fake_logger.messages) == 1 + assert "content_search_text" in fake_logger.messages[0] + assert "candidates=1" in fake_logger.messages[0] + assert "limit=5" in fake_logger.messages[0] + assert "fallback=True" in fake_logger.messages[0] + + +@pytest.mark.asyncio +async def test_empty_or_unsafe_query_does_not_load_fallback_candidates() -> None: + db = _FakeSession() + + rows = await channels.content_channel( + db, # type: ignore[arg-type] + user_id="user-1", + namespace="knowledge", + query="!!! ---", + top_k=2, + exclude_document_ids=[], + exclude_sections=[], + ) + + assert rows == [] + assert db.calls == []