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_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 9b8eb1074..9281b0a2b 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -1194,3 +1194,230 @@ async def test_should_exclude_matching_sections_from_the_response( assert len(results) == 1 assert _result_source(results[0])["document_id"] == included_document["document_id"] assert _result_source(results[0])["section_path"] == included_document["section_path"] + + +@pytest.mark.asyncio +async def test_content_channel_fts_should_match_any_query_token( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + async with developer_api_client_factory() as api_client: + alpha_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-or", + source_file_name="alpha.pdf", + section_path="Root / Alpha", + content="alpha evidence only", + ) + beta_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-or", + source_file_name="beta.pdf", + section_path="Root / Beta", + content="beta evidence only", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-or", + source_file_name="filler.pdf", + section_path="Root / Filler", + content="unrelated evidence only", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-content-fts-or", + "query": "alpha beta", + "top_k": 2, + "channels": ["content"], + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + assert response_json["router_used"] == "classic_topk" + assert { + _result_source(result)["document_id"] for result in results + } == {alpha_document["document_id"], beta_document["document_id"]} + + +@pytest.mark.asyncio +async def test_path_channel_fts_should_match_any_query_token( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + async with developer_api_client_factory() as api_client: + first_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-path-fts-or", + source_file_name="first.pdf", + section_path="root / needlepath", + content="generic evidence one", + ) + second_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-path-fts-or", + source_file_name="second.pdf", + section_path="root / alternatepath", + content="generic evidence two", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-path-fts-or", + source_file_name="filler.pdf", + section_path="root / filler", + content="generic evidence three", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-path-fts-or", + "query": "needlepath alternatepath", + "top_k": 2, + "channels": ["path"], + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + assert response_json["router_used"] == "classic_topk" + assert { + _result_source(result)["document_id"] for result in results + } == {first_document["document_id"], second_document["document_id"]} + + +@pytest.mark.asyncio +async def test_content_fts_should_apply_filters_before_candidate_limit( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", "1") + async with developer_api_client_factory() as api_client: + included_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-filter", + source_file_name="included.pdf", + section_path="Root / Allowed", + content="bounded marker included", + ) + excluded_section = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-filter", + source_file_name="excluded.pdf", + section_path="Root / Allowed / Hidden", + content="bounded marker excluded section", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-filter", + source_file_name="image.pdf", + section_path="Root / Allowed", + content="bounded marker excluded type", + chunk_type="image", + file_path="images/marker.png", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-filter", + source_file_name="elsewhere.pdf", + section_path="Root / Elsewhere", + content="bounded marker excluded signal path", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-content-fts-filter", + "query": "bounded marker", + "top_k": 1, + "channels": ["content"], + "chunk_types": ["text"], + "signal_paths": ["allowed"], + "filter_mode": "keep", + "exclude_sections": [ + { + "document_id": excluded_section["document_id"], + "section_path": "Root / Allowed", + } + ], + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + assert response_json["router_used"] == "classic_topk" + assert len(results) == 1 + assert _result_source(results[0])["document_id"] == included_document["document_id"] + + +@pytest.mark.asyncio +async def test_content_fts_should_preserve_sectionless_chunk_for_section_exclusion( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + async with developer_api_client_factory() as api_client: + sectionless_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-sectionless", + source_file_name="sectionless.pdf", + section_path="Root / Published", + content="sectionless marker evidence", + ) + await ContractDatabase.execute( + """ + UPDATE document_chunks + SET section_id = NULL + WHERE document_id = :document_id + AND chunk_id = :chunk_id + """, + { + "document_id": sectionless_document["document_id"], + "chunk_id": sectionless_document["chunk_id"], + }, + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-content-fts-sectionless", + source_file_name="filler.pdf", + section_path="Root / Filler", + content="unrelated filler evidence", + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-content-fts-sectionless", + "query": "sectionless marker", + "top_k": 1, + "channels": ["content"], + "exclude_sections": [ + { + "document_id": sectionless_document["document_id"], + "section_path": "Root / Excluded", + } + ], + "use_agentic": False, + }, + ) + + assert response.status_code == 200 + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + assert len(results) == 1 + assert _result_source(results[0])["document_id"] == sectionless_document[ + "document_id" + ] + assert _result_source(results[0])["section_path"] is None 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/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 8cad8a930..bf87a29ba 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -6,8 +6,11 @@ """ from __future__ import annotations +import json +import time from typing import Any +from loguru import logger from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession @@ -16,6 +19,7 @@ tokenize_query_for_ranker, ) from shared.services.retrieval.search.section_filters import is_excluded_section +from shared.services.retrieval.settings import get_postgres_fts_candidate_limit _SCOPED_CORPUS_CTE = """ @@ -120,6 +124,52 @@ def _build_extra_filters( 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.""" + 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 + + +def _build_fts_or_query(query_tokens: list[str]) -> str: + """Build a safely parameterized websearch query with ANY-token semantics.""" + return " OR ".join( + json.dumps(token, ensure_ascii=False) + for token in query_tokens + if token.strip() + ) + + +def _join_sql_filters(*filters: str) -> str: + return "\n ".join(filter(None, filters)) + + def _row_to_dict(row: Any) -> dict[str, Any]: return dict(row._mapping) @@ -222,6 +272,7 @@ async def _bm25_channel( query_tokens = tokenize_query_for_ranker(query) if not query_tokens: return [] + started_at = time.monotonic() exclude_clause = _build_exclude_clause(exclude_document_ids) extra_sql, extra_params = _build_extra_filters( @@ -229,25 +280,79 @@ 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) - - sql = _SCOPED_CORPUS_CTE.format(exclude_clause=exclude_clause, extra_filters=extra_sql) + 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()] + params.update(section_params) + candidate_limit = get_postgres_fts_candidate_limit() + params["fts_candidate_limit"] = candidate_limit + fts_or_query = _build_fts_or_query(query_tokens) + search_vector_field = { + "content_search_text": "content_search_tsv", + "path_search_text": "path_search_tsv", + }[search_field] + channel_name = search_field.removesuffix("_search_text") + fallback_used = not fts_or_query + + base_sql = _SCOPED_CORPUS_CTE.format( + exclude_clause=exclude_clause, + extra_filters=extra_sql, + ) + rows: list[dict[str, Any]] = [] + if fts_or_query: + params["fts_or_query"] = fts_or_query + fts_sql = base_sql + f""" + SELECT sc.* + FROM scoped_chunks sc + WHERE COALESCE(sc.{search_field}, '') <> '' + AND sc.{search_vector_field} @@ websearch_to_tsquery( + 'simple', :fts_or_query + ) + ORDER BY ts_rank_cd( + sc.{search_vector_field}, + websearch_to_tsquery('simple', :fts_or_query) + ) DESC, sc.id + LIMIT :fts_candidate_limit + """ + result = await db.execute(text(fts_sql), params) + rows = [_row_to_dict(row) for row in result.all()] + + if not rows: + fallback_used = True + fallback_sql = base_sql + f""" + SELECT sc.* + FROM scoped_chunks sc + WHERE COALESCE(sc.{search_field}, '') <> '' + ORDER BY sc.id + LIMIT :fts_candidate_limit + """ + fallback_params = dict(params) + fallback_params.pop("fts_or_query", None) + result = await db.execute(text(fallback_sql), fallback_params) + rows = [_row_to_dict(row) for row in result.all()] + + # Keep the shared predicate as a defensive check while SQL owns pre-LIMIT filtering. rows = _filter_excluded_sections(rows, exclude_sections) - + candidate_count = len(rows) ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) - return ranked_rows[:top_k] + final_rows = ranked_rows[:top_k] + duration_ms = round((time.monotonic() - started_at) * 1000, 2) + logger.info( + "retrieval.bm25_channel " + f"channel={channel_name} scoped_count=unavailable " + f"candidate_count={candidate_count} candidate_limit={candidate_limit} " + f"ranked_count={len(ranked_rows)} duration_ms={duration_ms} " + f"fallback_used={str(fallback_used).lower()}" + ) + + return final_rows async def term_channel( diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py index 533251e56..7525fb62d 100644 --- a/packages/shared-python/shared/services/retrieval/settings.py +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -1,16 +1,34 @@ from __future__ import annotations +import os + CHANNEL_WEIGHT_PATH = 1.0 CHANNEL_WEIGHT_CONTENT = 2.0 CHANNEL_WEIGHT_TERM = 1.5 INTERNAL_RECALL_K_MULTIPLIER = 2 RRF_K = 60 DEFAULT_TOP_K = 10 +DEFAULT_POSTGRES_FTS_CANDIDATE_LIMIT = 2000 VALID_CHUNK_TYPES: set[str] = {"text", "image", "table", "page"} ASSET_CHUNK_TYPES: set[str] = {"image", "table"} +def get_postgres_fts_candidate_limit() -> int: + """Return the positive candidate cap for PostgreSQL FTS prefiltering.""" + raw_limit = os.environ.get( + "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", + str(DEFAULT_POSTGRES_FTS_CANDIDATE_LIMIT), + ) + try: + candidate_limit = int(raw_limit) + except ValueError: + return DEFAULT_POSTGRES_FTS_CANDIDATE_LIMIT + if candidate_limit < 1: + return DEFAULT_POSTGRES_FTS_CANDIDATE_LIMIT + return candidate_limit + + def normalize_chunk_types(chunk_types: list[str] | set[str] | None) -> set[str] | None: """Normalize user-provided chunk_types to a validated set. 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..6a7290efc --- /dev/null +++ b/packages/shared-python/shared/tests/test_retrieval_search_channels.py @@ -0,0 +1,207 @@ +"""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 info(self, message: str) -> None: + 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.setenv("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 "websearch_to_tsquery('simple', :fts_or_query)" 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_or_query"] == '"alpha" OR "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.setenv("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_bounded_fallback_and_logs_metrics( + monkeypatch: MonkeyPatch, +) -> None: + monkeypatch.setenv("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 "@@ websearch_to_tsquery" in fts_sql + assert "@@ websearch_to_tsquery" not in fallback_sql + assert "LIMIT :fts_candidate_limit" in fallback_sql + assert fts_params["fts_candidate_limit"] == 5 + assert fallback_params["fts_candidate_limit"] == 5 + assert len(fake_logger.messages) == 1 + assert "channel=content" in fake_logger.messages[0] + assert "scoped_count=unavailable" in fake_logger.messages[0] + assert "candidate_count=1" in fake_logger.messages[0] + assert "candidate_limit=5" in fake_logger.messages[0] + assert "ranked_count=1" in fake_logger.messages[0] + assert "duration_ms=" in fake_logger.messages[0] + assert "fallback_used=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 == []