From 36e7467c22926378b4f31f9a4725dbfd614e08bc Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Wed, 3 Jun 2026 22:26:45 +0530 Subject: [PATCH 001/115] feat(search): per-assistant + global cross-encoder reranking, unbiased candidates Incremental, A/B-comparable reranking instead of a single global switch. Two-level gate: rerank runs only when the global master switch RERANK_ENABLED (a GPU-backed model server is deployed) AND the per-assistant opt-in Persona.rerank_enabled are both on. Default off everywhere, so existing assistants and the GPU-free local setup are unchanged and need no GPU. - Persona.rerank_enabled column (+ migration f6a7b8c9d0e1, server_default false), threaded through upsert/create_update_persona and the persona API models, with a 'Rerank results (beta)' toggle in the assistant editor. - Single resolver _resolve_skip_rerank() in retrieval_preprocessing is now the one place both chat and Slack decide reranking (Slack passes skip_rerank=None to share it). Legacy ENABLE_RERANKING_* flags kept as a fallback. - RERANK_MODEL_NAME makes the cross-encoder env-selectable (prod can pick a stronger model, e.g. BAAI/bge-reranker-v2-m3); model server warms it when RERANK_ENABLED. - Retrieval split: when reranking is on, skip the two-query source-prioritization flow (it normalizes a narrow source-filtered set independently, inflating those scores and polluting the rerank candidate window) and run a single all-sources query; when off, the legacy prioritized flow is unchanged. Driven by prioritize_sources=query.skip_rerank. - Tests: global x per-assistant resolver matrix; single-vs-two-query split. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../f6a7b8c9d0e1_persona_rerank_enabled.py | 39 ++++++++++ .../slack/handlers/handle_message.py | 6 +- backend/danswer/db/models.py | 7 ++ backend/danswer/db/persona.py | 4 + backend/danswer/document_index/interfaces.py | 6 ++ backend/danswer/document_index/vespa/index.py | 20 ++++- .../search/preprocessing/preprocessing.py | 28 ++++++- .../danswer/search/retrieval/search_runner.py | 6 ++ .../danswer/server/features/persona/models.py | 5 ++ backend/model_server/main.py | 7 +- backend/shared_configs/configs.py | 14 +++- .../unit/danswer/document_index/__init__.py | 0 .../danswer/document_index/vespa/__init__.py | 0 .../vespa/test_query_vespa_prioritization.py | 61 +++++++++++++++ backend/tests/unit/danswer/search/__init__.py | 0 .../danswer/search/preprocessing/__init__.py | 0 .../preprocessing/test_resolve_skip_rerank.py | 75 +++++++++++++++++++ .../app/admin/assistants/AssistantEditor.tsx | 12 ++- web/src/app/admin/assistants/interfaces.ts | 1 + web/src/app/admin/assistants/lib.ts | 4 + 20 files changed, 284 insertions(+), 11 deletions(-) create mode 100644 backend/alembic/versions/f6a7b8c9d0e1_persona_rerank_enabled.py create mode 100644 backend/tests/unit/danswer/document_index/__init__.py create mode 100644 backend/tests/unit/danswer/document_index/vespa/__init__.py create mode 100644 backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py create mode 100644 backend/tests/unit/danswer/search/__init__.py create mode 100644 backend/tests/unit/danswer/search/preprocessing/__init__.py create mode 100644 backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_rerank.py diff --git a/backend/alembic/versions/f6a7b8c9d0e1_persona_rerank_enabled.py b/backend/alembic/versions/f6a7b8c9d0e1_persona_rerank_enabled.py new file mode 100644 index 00000000000..f26ec265572 --- /dev/null +++ b/backend/alembic/versions/f6a7b8c9d0e1_persona_rerank_enabled.py @@ -0,0 +1,39 @@ +"""persona: add rerank_enabled (per-assistant cross-encoder reranking opt-in) + +Per-assistant toggle for cross-encoder reranking. Only takes effect when +reranking is globally available (RERANK_ENABLED + a GPU-backed model server); +default false so existing assistants and the GPU-free local/default setup are +unchanged. Lets reranking be rolled out incrementally / A-B compared per +assistant before becoming the default. See db/models.py::Persona and +search/preprocessing/preprocessing.py. + +Revision ID: f6a7b8c9d0e1 +Revises: e5f6a7b8c9d0 +Create Date: 2026-06-03 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "f6a7b8c9d0e1" +down_revision = "e5f6a7b8c9d0" +branch_labels: None = None +depends_on: None = None + + +def upgrade() -> None: + op.add_column( + "persona", + sa.Column( + "rerank_enabled", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ), + ) + + +def downgrade() -> None: + op.drop_column("persona", "rerank_enabled") diff --git a/backend/danswer/danswerbot/slack/handlers/handle_message.py b/backend/danswer/danswerbot/slack/handlers/handle_message.py index d4b79ef0d2c..b422fd2711d 100644 --- a/backend/danswer/danswerbot/slack/handlers/handle_message.py +++ b/backend/danswer/danswerbot/slack/handlers/handle_message.py @@ -69,7 +69,6 @@ from danswer.search.models import OptionalSearchSetting from danswer.search.models import RetrievalDetails from danswer.utils.logger import setup_logger -from shared_configs.configs import ENABLE_RERANKING_ASYNC_FLOW logger_base = setup_logger() @@ -659,7 +658,10 @@ def _get_answer(new_message_request: DirectQARequest) -> OneShotQAResponse | Non persona_id=persona.id if persona is not None else 0, retrieval_options=retrieval_details, chain_of_thought=not disable_cot, - skip_rerank=not ENABLE_RERANKING_ASYNC_FLOW, + # Leave None so retrieval_preprocessing resolves reranking from + # the global RERANK_ENABLED switch + the assistant's + # rerank_enabled flag — same logic as the chat flow. + skip_rerank=None, ) ) except Exception as e: diff --git a/backend/danswer/db/models.py b/backend/danswer/db/models.py index f99b0da8f24..84751d97a7b 100644 --- a/backend/danswer/db/models.py +++ b/backend/danswer/db/models.py @@ -1014,6 +1014,13 @@ class Persona(Base): recency_bias: Mapped[RecencyBiasSetting] = mapped_column( Enum(RecencyBiasSetting, native_enum=False) ) + # Per-assistant opt-in for cross-encoder reranking (beta). Only takes effect + # when reranking is globally available (RERANK_ENABLED + a GPU-backed model + # server); see search/preprocessing/preprocessing.py. Default off so existing + # assistants and the GPU-free local setup are unchanged until toggled. + rerank_enabled: Mapped[bool] = mapped_column( + Boolean, default=False, server_default="false" + ) # Allows the Persona to specify a different LLM version than is controlled # globablly via env variables. For flexibility, validity is not currently enforced # NOTE: only is applied on the actual response generation - is not used for things like diff --git a/backend/danswer/db/persona.py b/backend/danswer/db/persona.py index e39eb47306b..0a83b78492d 100644 --- a/backend/danswer/db/persona.py +++ b/backend/danswer/db/persona.py @@ -74,6 +74,7 @@ def create_update_persona( llm_relevance_filter=create_persona_request.llm_relevance_filter, llm_filter_extraction=create_persona_request.llm_filter_extraction, recency_bias=create_persona_request.recency_bias, + rerank_enabled=create_persona_request.rerank_enabled, prompt_ids=create_persona_request.prompt_ids, tool_ids=create_persona_request.tool_ids, document_set_ids=create_persona_request.document_set_ids, @@ -342,6 +343,7 @@ def upsert_persona( starter_messages: list[StarterMessage] | None, is_public: bool, db_session: Session, + rerank_enabled: bool = False, prompt_ids: list[int] | None = None, document_set_ids: list[int] | None = None, tool_ids: list[int] | None = None, @@ -393,6 +395,7 @@ def upsert_persona( persona.llm_relevance_filter = llm_relevance_filter persona.llm_filter_extraction = llm_filter_extraction persona.recency_bias = recency_bias + persona.rerank_enabled = rerank_enabled persona.default_persona = default_persona persona.llm_model_provider_override = llm_model_provider_override persona.llm_model_version_override = llm_model_version_override @@ -424,6 +427,7 @@ def upsert_persona( llm_relevance_filter=llm_relevance_filter, llm_filter_extraction=llm_filter_extraction, recency_bias=recency_bias, + rerank_enabled=rerank_enabled, default_persona=default_persona, prompts=prompts or [], document_sets=document_sets or [], diff --git a/backend/danswer/document_index/interfaces.py b/backend/danswer/document_index/interfaces.py index 6adedd45268..e922057cfaa 100644 --- a/backend/danswer/document_index/interfaces.py +++ b/backend/danswer/document_index/interfaces.py @@ -298,10 +298,16 @@ def hybrid_retrieval( num_to_retrieve: int, offset: int = 0, hybrid_alpha: float | None = None, + prioritize_sources: bool = True, ) -> list[InferenceChunk]: """ Run hybrid search and return a list of inference chunks. + prioritize_sources: when True (default) the implementation may apply its + source-prioritization behavior. The reranking path passes False so a + single, comparably-scored result set is returned for the cross-encoder + to reorder. + NOTE: the query passed in here is the unprocessed plain text query. Preprocessing is expected to be handled by this function as it may depend on the index implementation. Things like query expansion, synonym injection, stop word removal, lemmatization, etc. are diff --git a/backend/danswer/document_index/vespa/index.py b/backend/danswer/document_index/vespa/index.py index 8605a9bdd86..6bc36c455bc 100644 --- a/backend/danswer/document_index/vespa/index.py +++ b/backend/danswer/document_index/vespa/index.py @@ -702,7 +702,10 @@ def query_vespa_helper(params): @retry(tries=3, delay=1, backoff=2) -def _query_vespa(query_params: Mapping[str, str | int | float]) -> list[InferenceChunk]: +def _query_vespa( + query_params: Mapping[str, str | int | float], + prioritize_sources: bool = True, +) -> list[InferenceChunk]: if "query" in query_params and not cast(str, query_params["query"]).strip(): raise ValueError("No/empty query received") @@ -715,6 +718,18 @@ def _query_vespa(query_params: Mapping[str, str | int | float]) -> list[Inferenc else {}, ) + if not prioritize_sources: + # Single, all-sources query. Used on the reranking path: every chunk is + # scored on ONE comparable normalize_linear scale and the cross-encoder + # reorders the top chunks afterwards. The two-query prioritized flow + # below normalizes a narrow source-filtered set INDEPENDENTLY, which + # inflates those scores (normalize_linear is relative to each query's + # own candidate set) and would pollute the rerank candidate window. + # Honors the caller's `hits` (num_to_retrieve) rather than hardcoding it. + hits = query_vespa_helper(params) + chunks = [_vespa_hit_to_inference_chunk(hit) for hit in hits] + return sorted(chunks, key=lambda chunk: chunk.score or 0, reverse=True) + # Get prioritized sources from filters, default to web and sfkbarticles if none prioritized_sources = query_params.get("prioritized_sources") or [ "web", @@ -1181,6 +1196,7 @@ def hybrid_retrieval( title_content_ratio: float | None = TITLE_CONTENT_RATIO, distance_cutoff: float | None = SEARCH_DISTANCE_CUTOFF, edit_keyword_query: bool = EDIT_KEYWORD_QUERY, + prioritize_sources: bool = True, ) -> list[InferenceChunk]: vespa_where_clauses = _build_vespa_filters(filters) # Needs to be at least as much as the value set in Vespa schema config @@ -1218,7 +1234,7 @@ def hybrid_retrieval( "prioritized_sources": filters.prioritized_sources, # Use the non-None value } - return _query_vespa(params) + return _query_vespa(params, prioritize_sources=prioritize_sources) def admin_retrieval( self, diff --git a/backend/danswer/search/preprocessing/preprocessing.py b/backend/danswer/search/preprocessing/preprocessing.py index e59ef37c95e..2624b58da1a 100644 --- a/backend/danswer/search/preprocessing/preprocessing.py +++ b/backend/danswer/search/preprocessing/preprocessing.py @@ -5,6 +5,7 @@ from danswer.configs.chat_configs import DISABLE_LLM_FILTER_EXTRACTION from danswer.configs.chat_configs import FAVOR_RECENT_DECAY_MULTIPLIER from danswer.configs.chat_configs import NUM_RETURNED_HITS +from danswer.db.models import Persona from danswer.db.models import User from danswer.llm.interfaces import LLM from danswer.search.enums import QueryFlow @@ -23,11 +24,34 @@ from danswer.utils.threadpool_concurrency import run_functions_in_parallel from danswer.utils.timing import log_function_time from shared_configs.configs import ENABLE_RERANKING_REAL_TIME_FLOW +from shared_configs.configs import RERANK_ENABLED logger = setup_logger() +def _resolve_skip_rerank( + explicit_skip_rerank: bool | None, + persona: Persona | None, +) -> bool: + """Single source of truth for whether to skip cross-encoder reranking, + used by both the chat and Slack flows. + + Reranking runs only when the global master switch (RERANK_ENABLED — i.e. a + GPU-backed model server is deployed) AND the per-assistant opt-in + (Persona.rerank_enabled) are both on. If a caller set skip_rerank explicitly + we honor it. With RERANK_ENABLED off (the local / GPU-free default) + reranking never runs regardless of the per-assistant flag, so no GPU is + required. The legacy ENABLE_RERANKING_REAL_TIME_FLOW env flag is retained + only as a fallback for callers that haven't adopted the per-assistant model. + """ + if explicit_skip_rerank is not None: + return explicit_skip_rerank + persona_opts_in = bool(persona and persona.rerank_enabled) + rerank = (RERANK_ENABLED and persona_opts_in) or ENABLE_RERANKING_REAL_TIME_FLOW + return not rerank + + @log_function_time(print_only=True) def retrieval_preprocessing( search_request: SearchRequest, @@ -168,9 +192,7 @@ def retrieval_preprocessing( ) llm_chunk_filter = False - skip_rerank = search_request.skip_rerank - if skip_rerank is None: - skip_rerank = not ENABLE_RERANKING_REAL_TIME_FLOW + skip_rerank = _resolve_skip_rerank(search_request.skip_rerank, persona) # Decays at 1 / (1 + (multiplier * num years)) if persona and persona.recency_bias == RecencyBiasSetting.NO_DECAY: diff --git a/backend/danswer/search/retrieval/search_runner.py b/backend/danswer/search/retrieval/search_runner.py index 411db5b0f56..845ee73e14d 100644 --- a/backend/danswer/search/retrieval/search_runner.py +++ b/backend/danswer/search/retrieval/search_runner.py @@ -154,6 +154,12 @@ def doc_index_retrieval( num_to_retrieve=query.num_hits, offset=query.offset, hybrid_alpha=hybrid_alpha, + # When reranking is on (skip_rerank=False) we skip source + # prioritization: its two-query flow normalizes a narrow + # source-filtered set independently, inflating those scores and + # polluting the rerank candidate window. When reranking is off, + # keep the existing prioritized behavior unchanged. + prioritize_sources=query.skip_rerank, ) else: diff --git a/backend/danswer/server/features/persona/models.py b/backend/danswer/server/features/persona/models.py index aee39e72af0..ca11c49153a 100644 --- a/backend/danswer/server/features/persona/models.py +++ b/backend/danswer/server/features/persona/models.py @@ -23,6 +23,9 @@ class CreatePersonaRequest(BaseModel): is_public: bool llm_filter_extraction: bool recency_bias: RecencyBiasSetting + # Per-assistant cross-encoder reranking opt-in (beta). Defaults False so + # older clients that omit it keep current behavior. + rerank_enabled: bool = False prompt_ids: list[int] document_set_ids: list[int] # e.g. ID of SearchTool or ImageGenerationTool or @@ -46,6 +49,7 @@ class PersonaSnapshot(BaseModel): num_chunks: float | None llm_relevance_filter: bool llm_filter_extraction: bool + rerank_enabled: bool llm_model_provider_override: str | None llm_model_version_override: str | None starter_messages: list[StarterMessage] | None @@ -82,6 +86,7 @@ def from_model( num_chunks=persona.num_chunks, llm_relevance_filter=persona.llm_relevance_filter, llm_filter_extraction=persona.llm_filter_extraction, + rerank_enabled=persona.rerank_enabled, llm_model_provider_override=persona.llm_model_provider_override, llm_model_version_override=persona.llm_model_version_override, starter_messages=persona.starter_messages, diff --git a/backend/model_server/main.py b/backend/model_server/main.py index 1aaf9567874..c9de03af83e 100644 --- a/backend/model_server/main.py +++ b/backend/model_server/main.py @@ -20,6 +20,7 @@ from shared_configs.configs import MIN_THREADS_ML_MODELS from shared_configs.configs import MODEL_SERVER_ALLOWED_HOST from shared_configs.configs import MODEL_SERVER_PORT +from shared_configs.configs import RERANK_ENABLED os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" @@ -41,7 +42,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator: if not INDEXING_ONLY: warm_up_intent_model() - if ENABLE_RERANKING_REAL_TIME_FLOW or ENABLE_RERANKING_ASYNC_FLOW: + if ( + RERANK_ENABLED + or ENABLE_RERANKING_REAL_TIME_FLOW + or ENABLE_RERANKING_ASYNC_FLOW + ): warm_up_cross_encoders() else: logger.info("This model server should only run document indexing.") diff --git a/backend/shared_configs/configs.py b/backend/shared_configs/configs.py index aeeb9cf2ae9..5e6442aeeb1 100644 --- a/backend/shared_configs/configs.py +++ b/backend/shared_configs/configs.py @@ -21,14 +21,24 @@ DOC_EMBEDDING_CONTEXT_SIZE = 512 # Cross Encoder Settings +# Global master switch for cross-encoder reranking. When true, the model server +# loads/warms the reranker and the app will rerank for assistants that opt in +# (Persona.rerank_enabled). When false (the default, and how the local/GPU-free +# setup runs) reranking is never attempted regardless of per-assistant flags, so +# no GPU is required. Pair with a GPU-backed model server in prod. +RERANK_ENABLED = os.environ.get("RERANK_ENABLED", "").lower() == "true" ENABLE_RERANKING_ASYNC_FLOW = ( os.environ.get("ENABLE_RERANKING_ASYNC_FLOW", "").lower() == "true" ) ENABLE_RERANKING_REAL_TIME_FLOW = ( os.environ.get("ENABLE_RERANKING_REAL_TIME_FLOW", "").lower() == "true" ) -# Only using one cross-encoder for now -CROSS_ENCODER_MODEL_ENSEMBLE = ["mixedbread-ai/mxbai-rerank-xsmall-v1"] +# Only using one cross-encoder for now. Env-overridable so a GPU-backed prod +# deployment can select a stronger reranker (e.g. BAAI/bge-reranker-v2-m3) +# without a code change; local/dev keeps the small default. +CROSS_ENCODER_MODEL_ENSEMBLE = [ + os.environ.get("RERANK_MODEL_NAME") or "mixedbread-ai/mxbai-rerank-xsmall-v1" +] CROSS_EMBED_CONTEXT_SIZE = 512 # This controls the minimum number of pytorch "threads" to allocate to the embedding diff --git a/backend/tests/unit/danswer/document_index/__init__.py b/backend/tests/unit/danswer/document_index/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/document_index/vespa/__init__.py b/backend/tests/unit/danswer/document_index/vespa/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py b/backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py new file mode 100644 index 00000000000..697ef40e531 --- /dev/null +++ b/backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py @@ -0,0 +1,61 @@ +"""Unit tests for _query_vespa's two paths. + +When prioritize_sources is False (the reranking path) it issues a SINGLE +all-sources Vespa query, so scores stay on one comparable normalize_linear +scale for the cross-encoder. When True (reranking off) it keeps the legacy +two-query flow: an all-sources query plus a second source-filtered query whose +hits are concatenated. The second query independently normalizes a narrower +set, which is why we skip it under reranking. +""" +import pytest + +from danswer.document_index.vespa import index as vespa_index + + +@pytest.fixture +def count_vespa_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: + """Replace the network call with a recorder that returns no hits, so the + function exercises its branching without touching Vespa.""" + calls: list[dict] = [] + + def fake_helper(params: dict) -> list: + calls.append(dict(params)) + return [] + + monkeypatch.setattr(vespa_index, "query_vespa_helper", fake_helper) + return calls + + +def test_single_query_when_not_prioritizing(count_vespa_calls: list[dict]) -> None: + out = vespa_index._query_vespa( + {"yql": "select * from sources x where true", "query": "hello"}, + prioritize_sources=False, + ) + assert len(count_vespa_calls) == 1 # one all-sources query, no second query + assert out == [] + # the single query must not have a source_type filter spliced into the YQL + assert "source_type contains" not in count_vespa_calls[0]["yql"] + + +def test_two_queries_when_prioritizing(count_vespa_calls: list[dict]) -> None: + vespa_index._query_vespa( + { + "yql": "select * from sources x where true", + "query": "hello", + "prioritized_sources": ["web"], + }, + prioritize_sources=True, + ) + assert len(count_vespa_calls) == 2 # all-sources + prioritized-sources + # the second (prioritized) query appends the source filter to the YQL + assert "source_type contains" not in count_vespa_calls[0]["yql"] + assert 'source_type contains "web"' in count_vespa_calls[1]["yql"] + + +def test_default_keeps_legacy_two_query_behavior(count_vespa_calls: list[dict]) -> None: + # No explicit flag => prioritize_sources defaults True => legacy behavior, + # so callers other than the reranking path are unaffected. + vespa_index._query_vespa( + {"yql": "select * from sources x where true", "query": "hello"} + ) + assert len(count_vespa_calls) == 2 diff --git a/backend/tests/unit/danswer/search/__init__.py b/backend/tests/unit/danswer/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/search/preprocessing/__init__.py b/backend/tests/unit/danswer/search/preprocessing/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_rerank.py b/backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_rerank.py new file mode 100644 index 00000000000..1a54be7677f --- /dev/null +++ b/backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_rerank.py @@ -0,0 +1,75 @@ +"""Unit tests for _resolve_skip_rerank — the single resolver that decides +whether a query skips cross-encoder reranking. + +The rule: rerank runs only when the global master switch (RERANK_ENABLED, i.e. +a GPU-backed model server is deployed) AND the per-assistant opt-in +(Persona.rerank_enabled) are both on. An explicit skip_rerank is honored as-is. +A legacy ENABLE_RERANKING_REAL_TIME_FLOW=true forces rerank as a fallback. + +The function reads RERANK_ENABLED / ENABLE_RERANKING_REAL_TIME_FLOW as module +globals at call time, so we monkeypatch them on the module. +""" +from types import SimpleNamespace + +import pytest + +from danswer.search.preprocessing import preprocessing as pp + + +def _persona(rerank_enabled: bool) -> SimpleNamespace: + # Stands in for a Persona; _resolve_skip_rerank only reads .rerank_enabled. + return SimpleNamespace(rerank_enabled=rerank_enabled) + + +@pytest.fixture(autouse=True) +def _reset_flags(monkeypatch: pytest.MonkeyPatch) -> None: + # Default both global flags off so each test sets only what it needs. + monkeypatch.setattr(pp, "RERANK_ENABLED", False) + monkeypatch.setattr(pp, "ENABLE_RERANKING_REAL_TIME_FLOW", False) + + +# --- explicit skip_rerank is always honored, regardless of globals/persona --- + + +def test_explicit_true_is_honored(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "RERANK_ENABLED", True) + assert pp._resolve_skip_rerank(True, _persona(True)) is True + + +def test_explicit_false_is_honored(monkeypatch: pytest.MonkeyPatch) -> None: + # Even with everything off, an explicit "don't skip" wins. + assert pp._resolve_skip_rerank(False, _persona(False)) is False + + +# --- the global x per-assistant matrix (explicit None) --- + + +def test_global_off_persona_on_skips(monkeypatch: pytest.MonkeyPatch) -> None: + # Local / GPU-free default: global off => never rerank, even if opted in. + assert pp._resolve_skip_rerank(None, _persona(True)) is True + + +def test_global_on_persona_off_skips(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "RERANK_ENABLED", True) + assert pp._resolve_skip_rerank(None, _persona(False)) is True + + +def test_global_on_no_persona_skips(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "RERANK_ENABLED", True) + assert pp._resolve_skip_rerank(None, None) is True + + +def test_global_on_persona_on_reranks(monkeypatch: pytest.MonkeyPatch) -> None: + # The one combination that actually reranks. + monkeypatch.setattr(pp, "RERANK_ENABLED", True) + assert pp._resolve_skip_rerank(None, _persona(True)) is False + + +# --- legacy fallback flag --- + + +def test_legacy_realtime_flag_forces_rerank(monkeypatch: pytest.MonkeyPatch) -> None: + # Back-compat: the old env flag reranks even without the per-assistant opt-in. + monkeypatch.setattr(pp, "ENABLE_RERANKING_REAL_TIME_FLOW", True) + assert pp._resolve_skip_rerank(None, _persona(False)) is False + assert pp._resolve_skip_rerank(None, None) is False diff --git a/web/src/app/admin/assistants/AssistantEditor.tsx b/web/src/app/admin/assistants/AssistantEditor.tsx index c58cdcdadf9..0856f150170 100644 --- a/web/src/app/admin/assistants/AssistantEditor.tsx +++ b/web/src/app/admin/assistants/AssistantEditor.tsx @@ -178,6 +178,7 @@ export function AssistantEditor({ num_chunks: existingPersona?.num_chunks ?? null, include_citations: existingPersona?.prompts[0]?.include_citations ?? true, llm_relevance_filter: existingPersona?.llm_relevance_filter ?? false, + rerank_enabled: existingPersona?.rerank_enabled ?? false, llm_model_provider_override: existingPersona?.llm_model_provider_override ?? null, llm_model_version_override: @@ -213,6 +214,7 @@ export function AssistantEditor({ num_chunks: Yup.number().nullable(), include_citations: Yup.boolean().required(), llm_relevance_filter: Yup.boolean().required(), + rerank_enabled: Yup.boolean().required(), llm_model_version_override: Yup.string().nullable(), llm_model_provider_override: Yup.string().nullable(), starter_messages: Yup.array().of( @@ -580,6 +582,14 @@ export function AssistantEditor({ } /> + + ); -} \ No newline at end of file +} diff --git a/web/src/app/admin/assistants/interfaces.ts b/web/src/app/admin/assistants/interfaces.ts index 0a06ac4cc82..1ba7ff47e1f 100644 --- a/web/src/app/admin/assistants/interfaces.ts +++ b/web/src/app/admin/assistants/interfaces.ts @@ -32,6 +32,7 @@ export interface Persona { num_chunks?: number; llm_relevance_filter?: boolean; llm_filter_extraction?: boolean; + rerank_enabled?: boolean; llm_model_provider_override?: string; llm_model_version_override?: string; starter_messages: StarterMessage[] | null; diff --git a/web/src/app/admin/assistants/lib.ts b/web/src/app/admin/assistants/lib.ts index 4d42789d810..f74f039a30b 100644 --- a/web/src/app/admin/assistants/lib.ts +++ b/web/src/app/admin/assistants/lib.ts @@ -10,6 +10,7 @@ interface PersonaCreationRequest { include_citations: boolean; is_public: boolean; llm_relevance_filter: boolean | null; + rerank_enabled: boolean; llm_model_provider_override: string | null; llm_model_version_override: string | null; starter_messages: StarterMessage[] | null; @@ -30,6 +31,7 @@ interface PersonaUpdateRequest { include_citations: boolean; is_public: boolean; llm_relevance_filter: boolean | null; + rerank_enabled: boolean; llm_model_provider_override: string | null; llm_model_version_override: string | null; starter_messages: StarterMessage[] | null; @@ -106,6 +108,7 @@ function buildPersonaAPIBody( document_set_ids, num_chunks, llm_relevance_filter, + rerank_enabled, is_public, groups, users, @@ -117,6 +120,7 @@ function buildPersonaAPIBody( description, num_chunks, llm_relevance_filter, + rerank_enabled, llm_filter_extraction: false, is_public, recency_bias: "base_decay", From d41cf2487a541884770adc1d07a24361994bf705 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Wed, 3 Jun 2026 22:26:56 +0530 Subject: [PATCH 002/115] k8s: gpu-inference component (co-located rerank on a GPU node) Optional kustomize component that pins the inference-model-server onto a GPU node pool (nodeSelector agentpool=gpupool, toleration sku=gpu:NoSchedule, nvidia.com/gpu: 1) and serves the cross-encoder reranker (RERANK_MODEL_NAME, default BAAI/bge-reranker-v2-m3) alongside the embedding + intent models. Also sets real cpu/mem requests+limits (base leaves them empty -> eviction-prone). Opt in from the prod overlay's components: and set RERANK_ENABLED=true in env.properties. The existing model-server image already bundles CUDA torch, so no rebuild. Local omits the component and runs GPU-free. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../inference-model-server-gpu.yaml | 48 +++++++++++++++++++ k8s/optional/gpu-inference/kustomization.yaml | 37 ++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 k8s/optional/gpu-inference/inference-model-server-gpu.yaml create mode 100644 k8s/optional/gpu-inference/kustomization.yaml diff --git a/k8s/optional/gpu-inference/inference-model-server-gpu.yaml b/k8s/optional/gpu-inference/inference-model-server-gpu.yaml new file mode 100644 index 00000000000..458557e942b --- /dev/null +++ b/k8s/optional/gpu-inference/inference-model-server-gpu.yaml @@ -0,0 +1,48 @@ +# Strategic-merge patch: move the inference model server onto the GPU node pool +# and have it serve the cross-encoder reranker alongside the embedding + intent +# models (co-located, per the "deploy rerank + other inference models together +# on GPU" decision). +# +# Adjust the nodeSelector label/value and the toleration to match your GPU node +# pool. On AKS a GPU pool is typically labelled `agentpool: ` (or +# `kubernetes.azure.com/agentpool`) and tainted `sku=gpu:NoSchedule`. +apiVersion: apps/v1 +kind: Deployment +metadata: + name: inference-model-server-deployment +spec: + template: + spec: + # Pin to the GPU pool (single node — we accept single-GPU-node risk: if it + # is down, search/embedding degrade, not just reranking). + nodeSelector: + agentpool: gpupool + # GPU node pool taint: sku=gpu:NoSchedule (confirmed). + tolerations: + - key: sku + operator: Equal + value: gpu + effect: NoSchedule + containers: + - name: inference-model-server + env: + # Stronger reranker than the xsmall default; trivially fast on a V100. + # No reindex — the reranker only scores chunk text at query time. + - name: RERANK_MODEL_NAME + value: BAAI/bge-reranker-v2-m3 + # NOTE: RERANK_ENABLED=true must also be set in the overlay's + # env.properties so it reaches BOTH this model server (warms the + # reranker at boot) AND the api-server/background pods (so the app + # actually reranks for opted-in assistants). It is intentionally NOT + # set here so the single source of truth stays the shared configmap. + resources: + # Replaces the empty `resources: {}` in base — also fixes the + # no-requests/limits smell that made this pod eviction-prone. + requests: + cpu: '2' + memory: 8Gi + nvidia.com/gpu: '1' + limits: + cpu: '4' + memory: 12Gi + nvidia.com/gpu: '1' diff --git a/k8s/optional/gpu-inference/kustomization.yaml b/k8s/optional/gpu-inference/kustomization.yaml new file mode 100644 index 00000000000..1268ed5f66c --- /dev/null +++ b/k8s/optional/gpu-inference/kustomization.yaml @@ -0,0 +1,37 @@ +# Kustomize Component — GPU-backed inference + cross-encoder reranking. +# +# The "global flag" half of the incremental reranking rollout. Opting an +# overlay into this component: +# - schedules the inference-model-server onto the GPU node pool (patch in +# inference-model-server-gpu.yaml) and gives it a `nvidia.com/gpu` request, +# - serves the cross-encoder reranker (RERANK_MODEL_NAME, default +# BAAI/bge-reranker-v2-m3) alongside the embedding + intent models. +# +# This patch only handles the model-server pod. You MUST ALSO set, in the +# including overlay's env.properties (so it reaches every pod via env-configmap): +# +# RERANK_ENABLED=true +# +# That is the app-side master switch: with it on, the api-server/background +# pods rerank for assistants whose `rerank_enabled` flag is set (toggled per +# assistant in the admin UI). With it OFF — and by NOT including this component, +# which is how local/dev runs — reranking never runs and NO GPU is required. +# +# Opt in from an overlay: +# # k8s/overlays/prod/kustomization.yaml +# components: +# - ../../optional/gpu-inference +# # + add `RERANK_ENABLED=true` to env.properties +# +# Prereqs on the cluster: a GPU node pool (e.g. Standard_NC6s_v3 / V100) with +# the NVIDIA device plugin so `nvidia.com/gpu` is schedulable. The existing +# danswer-model-server image already bundles CUDA torch — no rebuild needed. +# Adjust the nodeSelector/toleration in the patch to match your pool. +apiVersion: kustomize.config.k8s.io/v1alpha1 +kind: Component + +patches: + - path: inference-model-server-gpu.yaml + target: + kind: Deployment + name: inference-model-server-deployment From c3e2471a2b83b08ff8b0ded3a923880e40251998 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Wed, 3 Jun 2026 22:26:56 +0530 Subject: [PATCH 003/115] docs: branch design doc for reranking/recency/retrieval quality Captures the verified query-path analysis (chat + Slack), the rerank flow and its corrected mental model, the recency decay math + levers (incl. the dead 'auto' auto-detect finding), the source-prioritization normalization bias and its two-path fix, the incremental per-assistant rollout, the GPU sizing/plan, the implementation map, and how to enable in prod. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/search-quality-reranking-and-recency.md | 316 +++++++++++++++++++ 1 file changed, 316 insertions(+) create mode 100644 docs/search-quality-reranking-and-recency.md diff --git a/docs/search-quality-reranking-and-recency.md b/docs/search-quality-reranking-and-recency.md new file mode 100644 index 00000000000..4052917b022 --- /dev/null +++ b/docs/search-quality-reranking-and-recency.md @@ -0,0 +1,316 @@ +# Search Answer Quality: Reranking, Recency, and Retrieval Prioritization + +> Design + investigation notes for branch **`feature/improve-queries`**. +> Goal: improve the **reliability of answers** (chat + Slack), give a path to +> **prioritize recent documents**, and do it **incrementally / A-B-comparably** +> rather than flipping a global switch. All findings below were verified against +> the code on this branch; file:line anchors are approximate (they drift as the +> code changes) but point at the right place. + +--- + +## 1. What problem this branch solves + +The fork retrieves well but **ranks and reranks poorly by default**, which costs +answer quality: + +1. **Cross-encoder reranking is OFF by default** — the LLM receives chunks in raw + bi-encoder hybrid order, so a genuinely-relevant chunk ranked #12 by vector + similarity never reaches the prompt (only ~10 chunks fit). +2. **A fork-specific "prioritized source" hack biases retrieval** — it runs a + second, source-filtered Vespa query and merges it, but because Vespa's + `normalize_linear` scoring is *relative to each query's candidate set*, the + narrow second query's scores are inflated and `web`/`sfkbarticles` get lifted + to the top regardless of true relevance. Default-on for every query. +3. **Recency is only a gentle decay** (and the `auto` setting is effectively + dead — see §4), so there's no real "prefer recent" behavior. +4. There was **no way to roll any of this out gradually** or compare old vs new. + +The branch adds a **two-level rerank gate** (global infra flag + per-assistant +toggle), splits retrieval so reranking gets clean (unbiased) candidates, and +documents the recency levers — without changing behavior for un-opted assistants +or the GPU-free local setup. + +--- + +## 2. The query/answer path (verified) + +Both surfaces converge on the same pipeline: + +``` +Chat: /chat send-message ─┐ + ├─► SearchTool.run ─► SearchPipeline +Slack: handle_message ─────┘ │ + ├─ retrieval_preprocessing (builds SearchQuery) + ├─ VespaIndex.hybrid_retrieval → _query_vespa (retrieve) + ├─ search_postprocessing (rerank + LLM relevance filter) + └─ prune_documents (token budget → ~10 chunks → LLM) +``` + +Key files: +- `danswer/tools/search/search_tool.py` — builds `SearchRequest` (sets `persona`, + leaves `skip_rerank=None`), runs `SearchPipeline`. +- `danswer/search/pipeline.py` — orchestrates the stages. +- `danswer/search/preprocessing/preprocessing.py` — `retrieval_preprocessing` + builds the `SearchQuery`; resolves filters, recency multiplier, and + `skip_rerank` (see §3, §5). +- `danswer/search/retrieval/search_runner.py` — `doc_index_retrieval` calls + `hybrid_retrieval`. +- `danswer/document_index/vespa/index.py` — `hybrid_retrieval` → `_query_vespa` + (the Vespa query); `danswer_chunk.sd` is the rank profile. +- `danswer/search/postprocessing/postprocessing.py` — `semantic_reranking`, + `rerank_chunks`, `filter_chunks`. +- `danswer/danswerbot/slack/handlers/handle_message.py` — Slack entry; now passes + `skip_rerank=None` so it shares the chat resolver. + +### Slack-specific guardrails worth knowing +- No-citations ⇒ **answer suppressed** (primary hallucination guard). +- `@retry(tries=5)` on answer generation; up-to-5× full re-execution on missing + citations (latency/cost + orphaned chat sessions — a known cost). +- Slack **bypasses ACL** for channels with document sets. + +--- + +## 3. Reranking: what it actually does + +**Bi-encoder (default, today):** Vespa scores each chunk as +`alpha·vector_similarity + (1-alpha)·BM25`, then `× document_boost × recency_bias`. +`vector_similarity` is cosine between the query embedding and each chunk's +*pre-computed* embedding — fast, but it can't model fine-grained query↔chunk +interaction. + +**Cross-encoder (reranking):** feeds the query **and** each chunk's *text* +together through a transformer (`mxbai-rerank-xsmall-v1` by default) so attention +runs across both → a far more accurate relevance judgment. Standard +retrieve-broad-then-rerank. + +**The real flow (corrected mental model):** +1. Vespa returns `NUM_RETURNED_HITS = 50` (+ up to 10 from the prioritized query, + deduped) — **not 15**. +2. `rerank_chunks` reranks only the **top 15** (`NUM_RERANKED_RESULTS`, + `chunks_to_rerank[:num_rerank]`); the rest get `score=None`, appended behind. +3. `semantic_reranking` computes the cross-encoder score, then + **`boosted = cross_encoder_score × document_boost × recency_bias`** + (`postprocessing.py:~74`) and re-sorts. So recency/boost are **re-applied** on + top of the cross-encoder score — it's not pure cross-encoder. +4. Token-budget prune → ~10 chunks to the LLM (LLM-relevant ones hoisted first). + +So **Vespa's score still selects *which* 15 are candidates**; the cross-encoder +reorders within that set. This is why a biased candidate set (see §5) matters. + +**Flags (both default `false` → rerank off everywhere):** +- Slack: `skip_rerank = not ENABLE_RERANKING_ASYNC_FLOW` +- Chat: `skip_rerank = not ENABLE_RERANKING_REAL_TIME_FLOW` + +This branch supersedes both with `RERANK_ENABLED` + per-assistant (see §6); the +old flags remain as a fallback. + +--- + +## 4. Recency / freshness + +The decay machinery **already exists and is identical to upstream Onyx** — there +is nothing to "catch up" on; the lever is *tuning* it. + +Vespa rank profile (`danswer_chunk.sd`): +``` +document_age = max(if(isNan(doc_updated_at),7890000, now()-doc_updated_at)/31536000, 0) # years +recency_bias = max(1 / (1 + query(decay_factor) * document_age), 0.75) # floored at 0.75 +# global-phase: (alpha·norm(vector) + (1-alpha)·norm(keyword)) * document_boost * recency_bias +``` +- `decay_factor = DOC_TIME_DECAY(0.5) × recency_bias_multiplier`. +- **Floor 0.75** ⇒ an old doc loses *at most 25%* of its score. It only *decays* + old docs; it never *boosts* fresh ones. + +`recency_bias_multiplier` per persona (`preprocessing.py:~176`): `no_decay`→0, +`base_decay`→0.5, `favor_recent`→1.0, `auto`→LLM-predicted. + +**Gotcha — `auto` is effectively dead.** The LLM time-filter auto-detection +(`enable_auto_detect_filters`) is **never threaded through**: `handle_message` +sets it on `RetrievalDetails`, but `SearchTool` doesn't copy it into +`SearchRequest` and `pipeline.py` doesn't pass it to `retrieval_preprocessing` +(whose param defaults `False`). So personas set to `recency_bias: "auto"` (the +seeded default) fall back to *base* decay and never favor recent. → **use +`favor_recent` for a deterministic recency preference.** + +**Levers to actually prefer recent (by effort):** +| Lever | Effect | Cost | +|---|---|---| +| persona `recency_bias: favor_recent` | doubles decay rate | config only | +| lower the `0.75` floor in `danswer_chunk.sd` | old docs decay further | Vespa schema redeploy | +| raise `DOC_TIME_DECAY` env | sharper 0–2yr decay | env only (floor-capped) | +| add a real `freshness()` boost term | lifts new docs | schema change + **reindex** | + +**Reranking weakens recency further** (a 0.75–1.0 multiplier barely moves a wide +cross-encoder score spread). So if recency matters, tune decay **separately** +from the rerank rollout, and measure independently. + +--- + +## 5. The source-prioritization bias (and the fix) + +`_query_vespa` (`index.py:~705`) historically ran **two** queries and merged them: +``` +Query A: all sources, hits=50 +Query B: + source_type ∈ {web, sfkbarticles}, hits=10 (default-on) +merge = B + A; dedup by (doc_id,chunk_id) keeping MAX score; sort desc +``` +**Why it's a bug:** the rank profile uses `normalize_linear(...)`, which is +min-max **relative to each query's own candidate set**. Query B's narrow set +normalizes its top docs near the ceiling regardless of absolute relevance; dedup +keeps the inflated B-score. ⇒ `web`/`sfkbarticles` are systematically lifted to +the top by a normalization artifact. + +**Interaction with reranking:** rerank only re-scores the **top 15 by this biased +score**, so the bias moves *upstream into candidate selection* — a genuinely +better non-prioritized doc ranked #16 never enters the rerank window. So the hack +**partially undermines** the rerank rollout. + +Secondary bugs in the same function: hardcodes `hits` (ignores +`num_to_retrieve`/persona limit) and ignores `offset` (pagination). + +**The fix (two paths, this branch):** +- **Reranking ON** ⇒ `prioritize_sources=False` ⇒ a **single all-sources query** + (one comparable `normalize_linear` scale, honors the caller's `hits`); the + cross-encoder reorders. +- **Reranking OFF** ⇒ legacy two-query prioritized flow, **byte-for-byte + unchanged**. + +Driven by `prioritize_sources=query.skip_rerank` in `doc_index_retrieval` — so it +rides the same per-assistant + global rerank decision, no separate flag. + +> NOTE: the prioritized-source hack is a deliberate fork divergence. The split +> preserves it whenever reranking is off; if you later remove it entirely, +> confirm the original product intent first (curated web/KB content?). + +--- + +## 6. The incremental design (what was built) + +**Two-level gate — rerank runs iff `RERANK_ENABLED` (global) AND +`persona.rerank_enabled` (per-assistant).** + +- **Global** `RERANK_ENABLED` (env, default false): the master switch. When on, a + GPU-backed model server warms the reranker and the app *may* rerank. Off (local + / default) ⇒ reranking never runs ⇒ **no GPU required**. +- **Per-assistant** `Persona.rerank_enabled` (bool, default false): which + assistants actually rerank. Lets you enable it on one assistant, compare + answers against an un-toggled copy in chat **or** Slack, and flip the default + once convinced. +- **Single resolver** `_resolve_skip_rerank(explicit, persona)` in + `preprocessing.py` is the one place both chat and Slack decide reranking + (`rerank = (RERANK_ENABLED and persona.rerank_enabled) or + ENABLE_RERANKING_REAL_TIME_FLOW`). Slack now passes `skip_rerank=None` so it + shares this logic. An explicit `skip_rerank` is honored as-is. + +Both chat and Slack respect it because `SearchTool` builds `SearchRequest` with +`skip_rerank=None` + `persona=`, and preprocessing reads +`search_request.persona`. + +--- + +## 7. Infrastructure / GPU plan + +Observed on the **darwin** cluster (June 2026): +- **No GPU anywhere** (4 nodes, all `nvidia.com/gpu: `). +- The inference model server (the `INDEXING_ONLY=false` pod) does **query + embedding + intent** today (~7.3 GiB RAM), with **no resource requests/limits** + on its container, on a node already at **87% memory** → eviction-prone. +- With rerank off, the cross-encoder is **not even loaded** (`warm_up_cross_encoders` + is gated). So the reranker is a *net-new* model + per-query compute when enabled. + +Decisions: +- **Self-host on a dedicated GPU node**, `Standard_NC6s_v3` (1× V100 16 GB) — more + than enough (a cross-encoder uses ~1 GB; reranks 15 chunks in <50 ms). +- **Co-locate** embedding + intent + reranker on that GPU node ("deploy rerank + + other inference models together on GPU") — maximizes the GPU, accelerates query + embedding, and evacuates the strained CPU node. **Single-GPU-node risk + accepted** (if it dies, search degrades, not just rerank). +- The existing `danswer-model-server` image **already bundles CUDA torch** — no + rebuild; it auto-uses the GPU once scheduled there. +- **Reranker model:** `BAAI/bge-reranker-v2-m3` (env-selectable via + `RERANK_MODEL_NAME`; local keeps the small default). **Switching rerankers needs + NO reindex** — cross-encoders score chunk *text* at query time; only changing + the *embedding* model forces a reindex. (Avoid late-interaction/ColBERT-style + models, which would need indexing changes.) +- **Upstream's stance:** Onyx made the reranker pluggable (default none; local-dev + = mxbai-xsmall) and **disables local reranking when there's no GPU** (PR #4011) + — i.e. don't self-host cross-encoder reranking on CPU. Hence the GPU node. + +k8s: `k8s/optional/gpu-inference/` component pins the inference deployment to the +GPU pool (`nodeSelector: agentpool=gpupool`, toleration `sku=gpu:NoSchedule`, +`nvidia.com/gpu: 1`, real cpu/mem requests+limits — also fixes the no-limits +smell), and sets `RERANK_MODEL_NAME`. The overlay must also set +`RERANK_ENABLED=true` in `env.properties` (reaches every pod via env-configmap). + +--- + +## 8. Implementation map (files changed on this branch) + +Backend: +- `db/models.py` — `Persona.rerank_enabled` (server_default false). +- `alembic/versions/f6a7b8c9d0e1_persona_rerank_enabled.py` — migration + (down_revision `e5f6a7b8c9d0`). +- `db/persona.py` — thread `rerank_enabled` through `upsert_persona` / + `create_update_persona`. +- `server/features/persona/models.py` — `CreatePersonaRequest` + + `PersonaSnapshot` + `from_model`. +- `shared_configs/configs.py` — `RERANK_ENABLED`, env-selectable + `CROSS_ENCODER_MODEL_ENSEMBLE` via `RERANK_MODEL_NAME`. +- `search/preprocessing/preprocessing.py` — `_resolve_skip_rerank` (single + resolver) + `Persona` import. +- `danswerbot/slack/handlers/handle_message.py` — `skip_rerank=None` (+ dropped + the now-unused `ENABLE_RERANKING_ASYNC_FLOW` import). +- `model_server/main.py` — warm cross-encoder when `RERANK_ENABLED`. +- `document_index/vespa/index.py` — `_query_vespa(prioritize_sources=...)` single + vs two-query split; `hybrid_retrieval(prioritize_sources=...)`. +- `document_index/interfaces.py` — `hybrid_retrieval` abstract signature. +- `search/retrieval/search_runner.py` — pass + `prioritize_sources=query.skip_rerank`. + +Web (`web/src/app/admin/assistants/`): +- `interfaces.ts`, `lib.ts`, `AssistantEditor.tsx` — "Rerank results (beta)" + toggle, mirroring `llm_relevance_filter`. + +Infra: +- `k8s/optional/gpu-inference/` — kustomization + inference patch. + +Tests (`backend/tests/unit/...`): +- `search/preprocessing/test_resolve_skip_rerank.py` — global × per-assistant + matrix + explicit override + legacy fallback (7 cases). +- `document_index/vespa/test_query_vespa_prioritization.py` — single-vs-two-query + split + default-is-legacy (3 cases). + +--- + +## 9. How to enable in prod (when ready) + +1. `alembic upgrade head` (adds `persona.rerank_enabled`) → bounce `dapi` + `dbe`. +2. Add the `Standard_NC6s_v3` GPU node pool (label `agentpool=gpupool`, taint + `sku=gpu:NoSchedule`, NVIDIA device plugin). +3. Prod overlay: add `- ../../optional/gpu-inference` to `components:` **and** set + `RERANK_ENABLED=true` in `env.properties`. Apply. +4. Toggle **"Rerank results"** on one test assistant → A/B compare against an + un-toggled copy in chat + Slack → flip the default once satisfied. + +Local stays GPU-free with zero config: omit the component, leave `RERANK_ENABLED` +unset → reranking never runs. + +--- + +## 10. Open / sequenced follow-ups + +- **Recency tuning is a separate experiment** from reranking — don't bundle. + Start with `favor_recent` on the test assistant (config); lower the `0.75` + floor only if needed (schema redeploy). Measure independently. +- **Confirm the intent** of the prioritized-source hack before ever removing it + outright (it's preserved whenever reranking is off). +- **`enable_auto_detect_filters` is dead** globally (§4) — fixing it would restore + LLM time/source filter extraction *and* the `auto` recency path; tracked + separately. +- **No min relevance-score cutoff** (`SEARCH_DISTANCE_CUTOFF=0` unused) — weak + chunks still fill the context window; candidate for a follow-up. +- **The 5× re-execution on missing citations** (Slack) — latency/cost + orphaned + chat sessions; candidate for a cap. +- Consider a **stronger/larger reranker** or hosted (Cohere) if `bge-reranker-v2-m3` + isn't enough — reindex-free either way. From b6e6ea7119fb759932031a038724048490c8c9e0 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 16:08:35 +0530 Subject: [PATCH 004/115] feat(search): one-shot LLM relevance filter + source diversity + TEI reranker client - LLM relevance filter is now a single listwise call on the MAIN llm (llm_eval_chunks_listwise + LISTWISE_CHUNK_FILTER_PROMPT, fails open), gated independently of reranking by LLM_RELEVANCE_FILTER_ENABLED x per-assistant llm_relevance_filter (resolver _resolve_skip_llm_chunk_filter). No GPU. - Reranker served by Hugging Face TEI on CPU when RERANK_SERVER_URL is set (CrossEncoderEnsembleModel /rerank path); model server skips loading the cross-encoder in that case. Replaces the GPU plan. - _query_vespa simplified to a single all-sources query (removed the two-query source-prioritization union and its score-inflation bias). - Source diversity moved to final selection: ensure_source_diversity in doc_pruning reserves up to SOURCE_DIVERSITY_RESERVED_SLOTS slots for PROTECTED_SOURCES, so KB/web aren't crowded out. Always-on, global, no per-assistant knob. - Chat per-conversation toggles (use_reranking / use_relevance_filter) threaded via SearchTool -> SearchRequest. Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/danswer/chat/process_message.py | 15 ++++- backend/danswer/configs/chat_configs.py | 22 ++++++ backend/danswer/document_index/interfaces.py | 6 -- backend/danswer/document_index/vespa/index.py | 67 +++---------------- backend/danswer/llm/answering/doc_pruning.py | 35 ++++++++++ backend/danswer/prompts/llm_chunk_filter.py | 26 +++++++ backend/danswer/search/pipeline.py | 5 +- .../search/postprocessing/postprocessing.py | 6 +- .../search/preprocessing/preprocessing.py | 43 ++++++++---- .../danswer/search/retrieval/search_runner.py | 6 -- backend/danswer/search/search_nlp_models.py | 29 +++++++- .../secondary_llm_flows/chunk_usefulness.py | 63 +++++++++++++++++ .../danswer/server/query_and_chat/models.py | 7 ++ backend/danswer/tools/search/search_tool.py | 10 +++ backend/model_server/main.py | 7 +- backend/shared_configs/configs.py | 15 +++-- .../inference-model-server-gpu.yaml | 48 ------------- k8s/optional/gpu-inference/kustomization.yaml | 37 ---------- 18 files changed, 266 insertions(+), 181 deletions(-) delete mode 100644 k8s/optional/gpu-inference/inference-model-server-gpu.yaml delete mode 100644 k8s/optional/gpu-inference/kustomization.yaml diff --git a/backend/danswer/chat/process_message.py b/backend/danswer/chat/process_message.py index b94264abd06..9d1e81ef2db 100644 --- a/backend/danswer/chat/process_message.py +++ b/backend/danswer/chat/process_message.py @@ -16,6 +16,7 @@ from danswer.chat.models import StreamingError from danswer.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE from danswer.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH +from danswer.configs.chat_configs import LLM_RELEVANCE_FILTER_ENABLED from danswer.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT from danswer.configs.constants import MessageType from danswer.configs.model_configs import GEN_AI_TEMPERATURE @@ -80,6 +81,7 @@ from danswer.tools.utils import explicit_tool_calling_supported from danswer.utils.logger import setup_logger from danswer.utils.timing import log_generator_function_time +from shared_configs.configs import RERANK_ENABLED logger = setup_logger() @@ -93,11 +95,11 @@ def translate_citations( for db_doc in db_docs: if db_doc.document_id not in doc_id_to_saved_doc_id_map: doc_id_to_saved_doc_id_map[db_doc.document_id] = db_doc.id - #print(f'found doc id: {db_doc.id}') + # print(f'found doc id: {db_doc.id}') citation_to_saved_doc_id_map: dict[int, int] = {} for citation in citations_list: - #print(f'citation id {citation.document_id} for doc num {citation.citation_num}') + # print(f'citation id {citation.document_id} for doc num {citation.citation_num}') if citation.citation_num not in citation_to_saved_doc_id_map: citation_to_saved_doc_id_map[ citation.citation_num @@ -425,6 +427,10 @@ def stream_chat_message_objects( if db_tool_model.in_code_tool_id: tool_cls = get_built_in_tool_by_id(db_tool_model.id, db_session) if tool_cls.__name__ == SearchTool.__name__ and not latest_query_files: + # Chat-page per-conversation toggles (default off, assistant + # settings intentionally ignored in chat). Each is gated by + # its global master switch; we pass explicit skip_* so + # retrieval_preprocessing uses these instead of the persona. search_tool = SearchTool( db_session=db_session, user=user, @@ -438,6 +444,11 @@ def stream_chat_message_objects( chunks_above=new_msg_req.chunks_above, chunks_below=new_msg_req.chunks_below, full_doc=new_msg_req.full_doc, + skip_rerank=not (RERANK_ENABLED and new_msg_req.use_reranking), + skip_llm_chunk_filter=not ( + LLM_RELEVANCE_FILTER_ENABLED + and new_msg_req.use_relevance_filter + ), ) tool_dict[db_tool_model.id] = [search_tool] elif tool_cls.__name__ == ImageGenerationTool.__name__: diff --git a/backend/danswer/configs/chat_configs.py b/backend/danswer/configs/chat_configs.py index e4b85f53cd8..7b3c3ea3af5 100644 --- a/backend/danswer/configs/chat_configs.py +++ b/backend/danswer/configs/chat_configs.py @@ -33,6 +33,28 @@ DISABLE_LLM_CHUNK_FILTER = ( os.environ.get("DISABLE_LLM_CHUNK_FILTER", "").lower() == "true" ) +# Global master switch for the (one-shot, main-LLM) relevance filter, mirroring +# RERANK_ENABLED. When true the app may run the filter for assistants/chats that +# opt in; when false (default) it never runs regardless of per-assistant flags. +# Unlike reranking this is LLM-only — it needs NO GPU — so it can be enabled on +# its own as a cheaper quality tier. DISABLE_LLM_CHUNK_FILTER still hard-kills it. +LLM_RELEVANCE_FILTER_ENABLED = ( + os.environ.get("LLM_RELEVANCE_FILTER_ENABLED", "").lower() == "true" +) +# Source diversity at final doc selection: guarantee that up to +# SOURCE_DIVERSITY_RESERVED_SLOTS of the highest-ranked docs from PROTECTED_SOURCES +# survive into the LLM prompt, so curated KB/web content isn't crowded out by a +# chatty high-relevance source (e.g. Slack). Replaces the old two-query +# source-prioritization hack — always-on, global, operates on the single +# comparably-scored candidate set. Set RESERVED_SLOTS=0 to disable. +PROTECTED_SOURCES = [ + s.strip().lower() + for s in (os.environ.get("PROTECTED_SOURCES") or "web,sfkbarticles").split(",") + if s.strip() +] +SOURCE_DIVERSITY_RESERVED_SLOTS = int( + os.environ.get("SOURCE_DIVERSITY_RESERVED_SLOTS") or 2 +) # Whether the LLM should be used to decide if a search would help given the chat history DISABLE_LLM_CHOOSE_SEARCH = ( os.environ.get("DISABLE_LLM_CHOOSE_SEARCH", "").lower() == "true" diff --git a/backend/danswer/document_index/interfaces.py b/backend/danswer/document_index/interfaces.py index e922057cfaa..6adedd45268 100644 --- a/backend/danswer/document_index/interfaces.py +++ b/backend/danswer/document_index/interfaces.py @@ -298,16 +298,10 @@ def hybrid_retrieval( num_to_retrieve: int, offset: int = 0, hybrid_alpha: float | None = None, - prioritize_sources: bool = True, ) -> list[InferenceChunk]: """ Run hybrid search and return a list of inference chunks. - prioritize_sources: when True (default) the implementation may apply its - source-prioritization behavior. The reranking path passes False so a - single, comparably-scored result set is returned for the cross-encoder - to reorder. - NOTE: the query passed in here is the unprocessed plain text query. Preprocessing is expected to be handled by this function as it may depend on the index implementation. Things like query expansion, synonym injection, stop word removal, lemmatization, etc. are diff --git a/backend/danswer/document_index/vespa/index.py b/backend/danswer/document_index/vespa/index.py index 6bc36c455bc..0af666089c1 100644 --- a/backend/danswer/document_index/vespa/index.py +++ b/backend/danswer/document_index/vespa/index.py @@ -704,7 +704,6 @@ def query_vespa_helper(params): @retry(tries=3, delay=1, backoff=2) def _query_vespa( query_params: Mapping[str, str | int | float], - prioritize_sources: bool = True, ) -> list[InferenceChunk]: if "query" in query_params and not cast(str, query_params["query"]).strip(): raise ValueError("No/empty query received") @@ -718,60 +717,15 @@ def _query_vespa( else {}, ) - if not prioritize_sources: - # Single, all-sources query. Used on the reranking path: every chunk is - # scored on ONE comparable normalize_linear scale and the cross-encoder - # reorders the top chunks afterwards. The two-query prioritized flow - # below normalizes a narrow source-filtered set INDEPENDENTLY, which - # inflates those scores (normalize_linear is relative to each query's - # own candidate set) and would pollute the rerank candidate window. - # Honors the caller's `hits` (num_to_retrieve) rather than hardcoding it. - hits = query_vespa_helper(params) - chunks = [_vespa_hit_to_inference_chunk(hit) for hit in hits] - return sorted(chunks, key=lambda chunk: chunk.score or 0, reverse=True) - - # Get prioritized sources from filters, default to web and sfkbarticles if none - prioritized_sources = query_params.get("prioritized_sources") or [ - "web", - "sfkbarticles", - ] - # All records - params["hits"] = 50 - filtered_hits_all = query_vespa_helper(params) - - # Records from prioritized sources - params["hits"] = 10 - source_conditions = " or ".join( - f'source_type contains "{source}"' for source in prioritized_sources - ) - params["yql"] = params["yql"] + f" and ({source_conditions})" - filtered_hits_prioritized = query_vespa_helper(params) - - filtered_hits_final = filtered_hits_prioritized + filtered_hits_all - - inference_chunks = [ - _vespa_hit_to_inference_chunk(hit) for hit in filtered_hits_final - ] - # inplace sorting based on score - # inference_chunks.sort(key=lambda x: x.score, reverse=True) - - unique_chunks: dict[tuple[str, int], InferenceChunk] = {} - for chunk in inference_chunks: - key = (chunk.document_id, chunk.chunk_id) - if key not in unique_chunks: - unique_chunks[key] = chunk - continue - - stored_chunk_score = unique_chunks[key].score or 0 - this_chunk_score = chunk.score or 0 - if stored_chunk_score < this_chunk_score: - unique_chunks[key] = chunk - - inference_chunks = sorted( - unique_chunks.values(), key=lambda x: x.score or 0, reverse=True - ) - # Good Debugging Spot - return inference_chunks + # Single, all-sources query: every chunk is scored on ONE comparable + # normalize_linear scale (honors the caller's `hits`). Source diversity — + # making sure curated KB/web docs aren't crowded out by a chatty source — + # is handled later, at final doc selection (see llm/answering/doc_pruning.py + # ::ensure_source_diversity), rather than by a second, independently- + # normalized query (which inflated those scores). + hits = query_vespa_helper(params) + inference_chunks = [_vespa_hit_to_inference_chunk(hit) for hit in hits] + return sorted(inference_chunks, key=lambda chunk: chunk.score or 0, reverse=True) @retry(tries=3, delay=1, backoff=2) @@ -1196,7 +1150,6 @@ def hybrid_retrieval( title_content_ratio: float | None = TITLE_CONTENT_RATIO, distance_cutoff: float | None = SEARCH_DISTANCE_CUTOFF, edit_keyword_query: bool = EDIT_KEYWORD_QUERY, - prioritize_sources: bool = True, ) -> list[InferenceChunk]: vespa_where_clauses = _build_vespa_filters(filters) # Needs to be at least as much as the value set in Vespa schema config @@ -1234,7 +1187,7 @@ def hybrid_retrieval( "prioritized_sources": filters.prioritized_sources, # Use the non-None value } - return _query_vespa(params, prioritize_sources=prioritize_sources) + return _query_vespa(params) def admin_retrieval( self, diff --git a/backend/danswer/llm/answering/doc_pruning.py b/backend/danswer/llm/answering/doc_pruning.py index 5a43ab3c6fd..a26be4894ac 100644 --- a/backend/danswer/llm/answering/doc_pruning.py +++ b/backend/danswer/llm/answering/doc_pruning.py @@ -5,6 +5,8 @@ from danswer.chat.models import ( LlmDoc, ) +from danswer.configs.chat_configs import PROTECTED_SOURCES +from danswer.configs.chat_configs import SOURCE_DIVERSITY_RESERVED_SLOTS from danswer.configs.constants import IGNORE_FOR_QA from danswer.configs.model_configs import DOC_EMBEDDING_CONTEXT_SIZE from danswer.llm.answering.models import DocumentPruningConfig @@ -84,6 +86,37 @@ def reorder_docs( return reordered_docs +def ensure_source_diversity(docs: list[T]) -> list[T]: + """Guarantee that up to SOURCE_DIVERSITY_RESERVED_SLOTS of the highest-ranked + docs from PROTECTED_SOURCES survive final selection, so curated KB/web + content isn't crowded out of the prompt by a chatty high-relevance source + (e.g. Slack). Promotes those protected docs to the front (keeping their + relative order); everything else keeps its order. No-op when disabled + (reserved <= 0), when there are no protected sources, or when none are + present in `docs`. + """ + if SOURCE_DIVERSITY_RESERVED_SLOTS <= 0 or not PROTECTED_SOURCES: + return docs + + protected = set(PROTECTED_SOURCES) + promote_indices: list[int] = [] + for ind, doc in enumerate(docs): + source = doc.source_type + source_str = (source.value if hasattr(source, "value") else str(source)).lower() + if source_str in protected: + promote_indices.append(ind) + if len(promote_indices) >= SOURCE_DIVERSITY_RESERVED_SLOTS: + break + + if not promote_indices: + return docs + + promote_set = set(promote_indices) + promoted = [docs[i] for i in promote_indices] + rest = [doc for i, doc in enumerate(docs) if i not in promote_set] + return promoted + rest + + def _remove_docs_to_ignore(docs: list[LlmDoc]) -> list[LlmDoc]: return [doc for doc in docs if not doc.metadata.get(IGNORE_FOR_QA)] @@ -103,6 +136,8 @@ def _apply_pruning( docs = reorder_docs(docs=docs, doc_relevance_list=doc_relevance_list) # remove docs that are explicitly marked as not for QA docs = _remove_docs_to_ignore(docs=docs) + # guarantee curated KB/web docs aren't crowded out before the token-budget cut + docs = ensure_source_diversity(docs) tokens_per_doc: list[int] = [] final_doc_ind = None diff --git a/backend/danswer/prompts/llm_chunk_filter.py b/backend/danswer/prompts/llm_chunk_filter.py index 623ae587703..8cec50b1985 100644 --- a/backend/danswer/prompts/llm_chunk_filter.py +++ b/backend/danswer/prompts/llm_chunk_filter.py @@ -25,6 +25,32 @@ """.strip() +# Listwise variant: judge ALL candidate sections in ONE call (cheaper + +# lower-latency than one call per chunk, and lets the model compare them). +# Run on the MAIN llm. {sections} is a numbered list; the model returns a JSON +# array of the USEFUL section numbers. +LISTWISE_CHUNK_FILTER_PROMPT = """ +You are given {count} numbered reference sections and a user query. For EACH +section, decide whether it is USEFUL for answering the query. It is NOT enough +to be related — the section must contain information USEFUL for answering. If a +section contains ANY useful information that counts; it need not fully answer +the query. + +Reference Sections: +{sections} + +User Query: +``` +{user_query} +``` + +Respond with EXACTLY AND ONLY a JSON array of the numbers of the useful +sections, in any order, e.g. [1, 3, 4]. If none are useful, respond with []. +""".strip() + + # Use the following for easy viewing of prompts if __name__ == "__main__": print(CHUNK_FILTER_PROMPT) + print("\n\n") + print(LISTWISE_CHUNK_FILTER_PROMPT) diff --git a/backend/danswer/search/pipeline.py b/backend/danswer/search/pipeline.py index 98b1a87161d..d73de5dd54e 100644 --- a/backend/danswer/search/pipeline.py +++ b/backend/danswer/search/pipeline.py @@ -300,7 +300,10 @@ def reranked_chunks(self) -> list[InferenceChunk]: self._postprocessing_generator = search_postprocessing( search_query=self.search_query, retrieved_chunks=self.retrieved_chunks, - llm=self.fast_llm, # use fast_llm for relevance, since it is a relatively easier task + # Use the MAIN llm (not fast_llm) for the relevance filter: it now + # judges all chunks in one listwise call, and the main model is more + # reliable at that structured multi-item judgment. + llm=self.llm, rerank_metrics_callback=self.rerank_metrics_callback, ) self._reranked_chunks = cast( diff --git a/backend/danswer/search/postprocessing/postprocessing.py b/backend/danswer/search/postprocessing/postprocessing.py index 3b36bcff3a9..0aeb3dec7d2 100644 --- a/backend/danswer/search/postprocessing/postprocessing.py +++ b/backend/danswer/search/postprocessing/postprocessing.py @@ -17,7 +17,7 @@ from danswer.search.models import SearchQuery from danswer.search.models import SearchType from danswer.search.search_nlp_models import CrossEncoderEnsembleModel -from danswer.secondary_llm_flows.chunk_usefulness import llm_batch_eval_chunks +from danswer.secondary_llm_flows.chunk_usefulness import llm_eval_chunks_listwise from danswer.utils.logger import setup_logger from danswer.utils.threadpool_concurrency import FunctionCall from danswer.utils.threadpool_concurrency import run_functions_in_parallel @@ -141,7 +141,9 @@ def filter_chunks( Returns a list of the unique chunk IDs that were marked as relevant""" chunks_to_filter = chunks_to_filter[: query.max_llm_filter_chunks] - llm_chunk_selection = llm_batch_eval_chunks( + # One listwise call over all candidates (on the main LLM) rather than one + # call per chunk — cheaper, lower latency, and lets the model compare them. + llm_chunk_selection = llm_eval_chunks_listwise( query=query.query, chunk_contents=[chunk.content for chunk in chunks_to_filter], llm=llm, diff --git a/backend/danswer/search/preprocessing/preprocessing.py b/backend/danswer/search/preprocessing/preprocessing.py index 2624b58da1a..7dfef923015 100644 --- a/backend/danswer/search/preprocessing/preprocessing.py +++ b/backend/danswer/search/preprocessing/preprocessing.py @@ -4,6 +4,7 @@ from danswer.configs.chat_configs import DISABLE_LLM_CHUNK_FILTER from danswer.configs.chat_configs import DISABLE_LLM_FILTER_EXTRACTION from danswer.configs.chat_configs import FAVOR_RECENT_DECAY_MULTIPLIER +from danswer.configs.chat_configs import LLM_RELEVANCE_FILTER_ENABLED from danswer.configs.chat_configs import NUM_RETURNED_HITS from danswer.db.models import Persona from danswer.db.models import User @@ -52,6 +53,30 @@ def _resolve_skip_rerank( return not rerank +def _resolve_skip_llm_chunk_filter( + explicit_skip: bool | None, + persona: Persona | None, + disable_llm_chunk_filter: bool, +) -> bool: + """Single source of truth for whether to skip the LLM relevance filter. + + Independent of reranking (it's LLM-only, needs no GPU). The filter runs only + when the global master switch LLM_RELEVANCE_FILTER_ENABLED AND the + per-assistant opt-in (Persona.llm_relevance_filter) are both on. The global + DISABLE_LLM_CHUNK_FILTER kill-switch always wins. If a caller set skip + explicitly (e.g. the chat flow, which has already applied the global gate + + its per-conversation toggle), honor it — but the kill-switch still applies. + """ + if disable_llm_chunk_filter: + return True + if explicit_skip is not None: + return explicit_skip + use = LLM_RELEVANCE_FILTER_ENABLED and bool( + persona and persona.llm_relevance_filter + ) + return not use + + @log_function_time(print_only=True) def retrieval_preprocessing( search_request: SearchRequest, @@ -179,19 +204,9 @@ def retrieval_preprocessing( prioritized_sources=preset_filters.prioritized_sources, # Use prioritized_sources from filters ) - llm_chunk_filter = False - if search_request.skip_llm_chunk_filter is not None: - llm_chunk_filter = not search_request.skip_llm_chunk_filter - elif persona: - llm_chunk_filter = persona.llm_relevance_filter - - if disable_llm_chunk_filter: - if llm_chunk_filter: - logger.info( - "LLM chunk filtering would have run but has been globally disabled" - ) - llm_chunk_filter = False - + skip_llm_chunk_filter = _resolve_skip_llm_chunk_filter( + search_request.skip_llm_chunk_filter, persona, disable_llm_chunk_filter + ) skip_rerank = _resolve_skip_rerank(search_request.skip_rerank, persona) # Decays at 1 / (1 + (multiplier * num years)) @@ -216,7 +231,7 @@ def retrieval_preprocessing( num_hits=limit if limit is not None else NUM_RETURNED_HITS, offset=offset or 0, skip_rerank=skip_rerank, - skip_llm_chunk_filter=not llm_chunk_filter, + skip_llm_chunk_filter=skip_llm_chunk_filter, chunks_above=search_request.chunks_above, chunks_below=search_request.chunks_below, full_doc=search_request.full_doc, diff --git a/backend/danswer/search/retrieval/search_runner.py b/backend/danswer/search/retrieval/search_runner.py index 845ee73e14d..411db5b0f56 100644 --- a/backend/danswer/search/retrieval/search_runner.py +++ b/backend/danswer/search/retrieval/search_runner.py @@ -154,12 +154,6 @@ def doc_index_retrieval( num_to_retrieve=query.num_hits, offset=query.offset, hybrid_alpha=hybrid_alpha, - # When reranking is on (skip_rerank=False) we skip source - # prioritization: its two-query flow normalizes a narrow - # source-filtered set independently, inflating those scores and - # polluting the rerank candidate window. When reranking is off, - # keep the existing prioritized behavior unchanged. - prioritize_sources=query.skip_rerank, ) else: diff --git a/backend/danswer/search/search_nlp_models.py b/backend/danswer/search/search_nlp_models.py index 761d9aa791f..13597c146c5 100644 --- a/backend/danswer/search/search_nlp_models.py +++ b/backend/danswer/search/search_nlp_models.py @@ -13,6 +13,7 @@ from danswer.utils.logger import setup_logger from shared_configs.configs import MODEL_SERVER_HOST from shared_configs.configs import MODEL_SERVER_PORT +from shared_configs.configs import RERANK_SERVER_URL from shared_configs.model_server_models import EmbedRequest from shared_configs.model_server_models import EmbedResponse from shared_configs.model_server_models import IntentRequest @@ -128,20 +129,44 @@ def __init__( self, model_server_host: str = MODEL_SERVER_HOST, model_server_port: int = MODEL_SERVER_PORT, + rerank_server_url: str = RERANK_SERVER_URL, ) -> None: + # When a TEI rerank server is configured, talk to it directly (its + # /rerank API). Otherwise fall back to our model server's + # /cross-encoder-scores (sentence-transformers) path. + self.tei_rerank_endpoint = ( + f"{rerank_server_url}/rerank" if rerank_server_url else None + ) model_server_url = build_model_server_url(model_server_host, model_server_port) self.rerank_server_endpoint = model_server_url + "/encoder/cross-encoder-scores" def predict(self, query: str, passages: list[str]) -> list[list[float]]: - rerank_request = RerankRequest(query=query, documents=passages) + if self.tei_rerank_endpoint: + return [self._predict_tei(query, passages)] + rerank_request = RerankRequest(query=query, documents=passages) response = requests.post( self.rerank_server_endpoint, json=rerank_request.dict() ) response.raise_for_status() - return RerankResponse(**response.json()).scores + def _predict_tei(self, query: str, passages: list[str]) -> list[float]: + """Call a Hugging Face TEI /rerank server and return scores in the SAME + order as `passages`. TEI returns [{index, score}, ...] sorted by score, + so we scatter them back to the input order. raw_scores=true keeps the + cross-encoder logits (the downstream normalization expects a logit-like + scale, not a 0-1 probability).""" + response = requests.post( + self.tei_rerank_endpoint, + json={"query": query, "texts": passages, "raw_scores": True}, + ) + response.raise_for_status() + scores = [0.0] * len(passages) + for item in response.json(): + scores[item["index"]] = item["score"] + return scores + class IntentModel: def __init__( diff --git a/backend/danswer/secondary_llm_flows/chunk_usefulness.py b/backend/danswer/secondary_llm_flows/chunk_usefulness.py index 8148a37138c..6dff17beeab 100644 --- a/backend/danswer/secondary_llm_flows/chunk_usefulness.py +++ b/backend/danswer/secondary_llm_flows/chunk_usefulness.py @@ -1,9 +1,12 @@ +import json +import re from collections.abc import Callable from danswer.llm.interfaces import LLM from danswer.llm.utils import dict_based_prompt_to_langchain_prompt from danswer.llm.utils import message_to_string from danswer.prompts.llm_chunk_filter import CHUNK_FILTER_PROMPT +from danswer.prompts.llm_chunk_filter import LISTWISE_CHUNK_FILTER_PROMPT from danswer.prompts.llm_chunk_filter import NONUSEFUL_PAT from danswer.utils.logger import setup_logger from danswer.utils.threadpool_concurrency import run_functions_tuples_in_parallel @@ -11,6 +14,66 @@ logger = setup_logger() +def _parse_useful_indices(model_output: str, count: int) -> set[int] | None: + """Parse the listwise filter's reply into a set of 1-based useful indices. + + Returns None when no JSON array can be found (a parse failure → caller + should fail OPEN and keep all chunks). An explicitly empty array `[]` is a + valid "none useful" answer and returns an empty set (not None). + """ + match = re.search(r"\[[\s\d,]*\]", model_output) + if match is None: + return None + try: + parsed = json.loads(match.group(0)) + except json.JSONDecodeError: + return None + return { + int(n) for n in parsed if isinstance(n, (int, float)) and 1 <= int(n) <= count + } + + +def llm_eval_chunks_listwise( + query: str, chunk_contents: list[str], llm: LLM +) -> list[bool]: + """Judge all chunks in a SINGLE LLM call (vs one call per chunk). + + Returns a parallel list of booleans. Fails OPEN: on any error or an + unparseable reply, every chunk is kept (True) — same philosophy as the + per-chunk path ("better to trust the (re)ranking if the LLM fails"). + """ + if not chunk_contents: + return [] + + sections = "\n\n".join( + f"Section {i + 1}:\n```\n{content}\n```" + for i, content in enumerate(chunk_contents) + ) + messages = [ + { + "role": "user", + "content": LISTWISE_CHUNK_FILTER_PROMPT.format( + count=len(chunk_contents), sections=sections, user_query=query + ), + } + ] + filled_prompt = dict_based_prompt_to_langchain_prompt(messages) + try: + model_output = message_to_string(llm.invoke(filled_prompt)) + except Exception: + logger.exception("Listwise relevance filter call failed — keeping all chunks") + return [True] * len(chunk_contents) + + useful = _parse_useful_indices(model_output, len(chunk_contents)) + if useful is None: + logger.warning( + "Could not parse listwise relevance filter output — keeping all chunks" + ) + return [True] * len(chunk_contents) + + return [(i + 1) in useful for i in range(len(chunk_contents))] + + def llm_eval_chunk(query: str, chunk_content: str, llm: LLM) -> bool: def _get_usefulness_messages() -> list[dict[str, str]]: messages = [ diff --git a/backend/danswer/server/query_and_chat/models.py b/backend/danswer/server/query_and_chat/models.py index ea1ce1ff680..91233197b6f 100644 --- a/backend/danswer/server/query_and_chat/models.py +++ b/backend/danswer/server/query_and_chat/models.py @@ -113,6 +113,13 @@ class CreateChatMessageRequest(ChunkContext): # used for seeded chats to kick off the generation of an AI answer use_existing_user_message: bool = False + # Per-conversation toggles for the search-quality features, default OFF and + # independent of the assistant's own settings (the chat page exposes these + # so a user can opt in for just this conversation). Each still requires its + # global master switch (RERANK_ENABLED / LLM_RELEVANCE_FILTER_ENABLED). + use_reranking: bool = False + use_relevance_filter: bool = False + @root_validator def check_search_doc_ids_or_retrieval_options(cls: BaseModel, values: dict) -> dict: search_doc_ids, retrieval_options = values.get("search_doc_ids"), values.get( diff --git a/backend/danswer/tools/search/search_tool.py b/backend/danswer/tools/search/search_tool.py index 75770e69f62..fe41724c984 100644 --- a/backend/danswer/tools/search/search_tool.py +++ b/backend/danswer/tools/search/search_tool.py @@ -78,6 +78,12 @@ def __init__( chunks_below: int = 0, full_doc: bool = False, bypass_acl: bool = False, + # Per-request reranking / relevance-filter overrides. None => let + # retrieval_preprocessing decide from the global flags + the assistant's + # settings (Slack / default path). The chat flow passes explicit values + # derived from the per-conversation toggles + global flags. + skip_rerank: bool | None = None, + skip_llm_chunk_filter: bool | None = None, ) -> None: self.user = user self.persona = persona @@ -93,6 +99,8 @@ def __init__( self.chunks_below = chunks_below self.full_doc = full_doc self.bypass_acl = bypass_acl + self.skip_rerank = skip_rerank + self.skip_llm_chunk_filter = skip_llm_chunk_filter self.db_session = db_session def name(self) -> str: @@ -211,6 +219,8 @@ def run(self, **kwargs: str) -> Generator[ToolResponse, None, None]: chunks_above=self.chunks_above, chunks_below=self.chunks_below, full_doc=self.full_doc, + skip_rerank=self.skip_rerank, + skip_llm_chunk_filter=self.skip_llm_chunk_filter, ), user=self.user, llm=self.llm, diff --git a/backend/model_server/main.py b/backend/model_server/main.py index c9de03af83e..c1fa2a5c653 100644 --- a/backend/model_server/main.py +++ b/backend/model_server/main.py @@ -21,6 +21,7 @@ from shared_configs.configs import MODEL_SERVER_ALLOWED_HOST from shared_configs.configs import MODEL_SERVER_PORT from shared_configs.configs import RERANK_ENABLED +from shared_configs.configs import RERANK_SERVER_URL os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" @@ -42,11 +43,15 @@ async def lifespan(app: FastAPI) -> AsyncGenerator: if not INDEXING_ONLY: warm_up_intent_model() + # Only load the cross-encoder here when reranking is enabled AND it's + # NOT served by an external TEI server (RERANK_SERVER_URL). With TEI, + # the reranker lives in that container, so this server stays embedding + + # intent only. if ( RERANK_ENABLED or ENABLE_RERANKING_REAL_TIME_FLOW or ENABLE_RERANKING_ASYNC_FLOW - ): + ) and not RERANK_SERVER_URL: warm_up_cross_encoders() else: logger.info("This model server should only run document indexing.") diff --git a/backend/shared_configs/configs.py b/backend/shared_configs/configs.py index 5e6442aeeb1..36b0aad8cb7 100644 --- a/backend/shared_configs/configs.py +++ b/backend/shared_configs/configs.py @@ -21,12 +21,17 @@ DOC_EMBEDDING_CONTEXT_SIZE = 512 # Cross Encoder Settings -# Global master switch for cross-encoder reranking. When true, the model server -# loads/warms the reranker and the app will rerank for assistants that opt in -# (Persona.rerank_enabled). When false (the default, and how the local/GPU-free -# setup runs) reranking is never attempted regardless of per-assistant flags, so -# no GPU is required. Pair with a GPU-backed model server in prod. +# Global master switch for cross-encoder reranking. When true, the reranker is +# available and the app will rerank for assistants that opt in +# (Persona.rerank_enabled). When false (default) reranking is never attempted +# regardless of per-assistant flags. RERANK_ENABLED = os.environ.get("RERANK_ENABLED", "").lower() == "true" +# If set, reranking is served by a Hugging Face Text-Embeddings-Inference (TEI) +# container at this base URL (its /rerank endpoint) — a CPU-optimized, +# full-precision way to host the cross-encoder WITHOUT a GPU. When set, our own +# model server does NOT load the cross-encoder (TEI owns it). Empty => use the +# legacy in-model-server sentence-transformers path. +RERANK_SERVER_URL = (os.environ.get("RERANK_SERVER_URL") or "").rstrip("/") ENABLE_RERANKING_ASYNC_FLOW = ( os.environ.get("ENABLE_RERANKING_ASYNC_FLOW", "").lower() == "true" ) diff --git a/k8s/optional/gpu-inference/inference-model-server-gpu.yaml b/k8s/optional/gpu-inference/inference-model-server-gpu.yaml deleted file mode 100644 index 458557e942b..00000000000 --- a/k8s/optional/gpu-inference/inference-model-server-gpu.yaml +++ /dev/null @@ -1,48 +0,0 @@ -# Strategic-merge patch: move the inference model server onto the GPU node pool -# and have it serve the cross-encoder reranker alongside the embedding + intent -# models (co-located, per the "deploy rerank + other inference models together -# on GPU" decision). -# -# Adjust the nodeSelector label/value and the toleration to match your GPU node -# pool. On AKS a GPU pool is typically labelled `agentpool: ` (or -# `kubernetes.azure.com/agentpool`) and tainted `sku=gpu:NoSchedule`. -apiVersion: apps/v1 -kind: Deployment -metadata: - name: inference-model-server-deployment -spec: - template: - spec: - # Pin to the GPU pool (single node — we accept single-GPU-node risk: if it - # is down, search/embedding degrade, not just reranking). - nodeSelector: - agentpool: gpupool - # GPU node pool taint: sku=gpu:NoSchedule (confirmed). - tolerations: - - key: sku - operator: Equal - value: gpu - effect: NoSchedule - containers: - - name: inference-model-server - env: - # Stronger reranker than the xsmall default; trivially fast on a V100. - # No reindex — the reranker only scores chunk text at query time. - - name: RERANK_MODEL_NAME - value: BAAI/bge-reranker-v2-m3 - # NOTE: RERANK_ENABLED=true must also be set in the overlay's - # env.properties so it reaches BOTH this model server (warms the - # reranker at boot) AND the api-server/background pods (so the app - # actually reranks for opted-in assistants). It is intentionally NOT - # set here so the single source of truth stays the shared configmap. - resources: - # Replaces the empty `resources: {}` in base — also fixes the - # no-requests/limits smell that made this pod eviction-prone. - requests: - cpu: '2' - memory: 8Gi - nvidia.com/gpu: '1' - limits: - cpu: '4' - memory: 12Gi - nvidia.com/gpu: '1' diff --git a/k8s/optional/gpu-inference/kustomization.yaml b/k8s/optional/gpu-inference/kustomization.yaml deleted file mode 100644 index 1268ed5f66c..00000000000 --- a/k8s/optional/gpu-inference/kustomization.yaml +++ /dev/null @@ -1,37 +0,0 @@ -# Kustomize Component — GPU-backed inference + cross-encoder reranking. -# -# The "global flag" half of the incremental reranking rollout. Opting an -# overlay into this component: -# - schedules the inference-model-server onto the GPU node pool (patch in -# inference-model-server-gpu.yaml) and gives it a `nvidia.com/gpu` request, -# - serves the cross-encoder reranker (RERANK_MODEL_NAME, default -# BAAI/bge-reranker-v2-m3) alongside the embedding + intent models. -# -# This patch only handles the model-server pod. You MUST ALSO set, in the -# including overlay's env.properties (so it reaches every pod via env-configmap): -# -# RERANK_ENABLED=true -# -# That is the app-side master switch: with it on, the api-server/background -# pods rerank for assistants whose `rerank_enabled` flag is set (toggled per -# assistant in the admin UI). With it OFF — and by NOT including this component, -# which is how local/dev runs — reranking never runs and NO GPU is required. -# -# Opt in from an overlay: -# # k8s/overlays/prod/kustomization.yaml -# components: -# - ../../optional/gpu-inference -# # + add `RERANK_ENABLED=true` to env.properties -# -# Prereqs on the cluster: a GPU node pool (e.g. Standard_NC6s_v3 / V100) with -# the NVIDIA device plugin so `nvidia.com/gpu` is schedulable. The existing -# danswer-model-server image already bundles CUDA torch — no rebuild needed. -# Adjust the nodeSelector/toleration in the patch to match your pool. -apiVersion: kustomize.config.k8s.io/v1alpha1 -kind: Component - -patches: - - path: inference-model-server-gpu.yaml - target: - kind: Deployment - name: inference-model-server-deployment From d3f4207448d2ef7ba227401b353425b7644b92f0 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 16:08:48 +0530 Subject: [PATCH 005/115] feat(chat): per-conversation Rerank / Relevance toggles in the chat input bar Two per-conversation switches (default off) wired through sendMessage -> CreateChatMessageRequest (use_reranking / use_relevance_filter), independent of the assistant's own settings. Co-Authored-By: Claude Opus 4.8 (1M context) --- web/src/app/chat/ChatPage.tsx | 12 ++++++++ web/src/app/chat/input/ChatInputBar.tsx | 38 +++++++++++++++++++++++++ web/src/app/chat/lib.tsx | 9 ++++++ 3 files changed, 59 insertions(+) diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx index dd0d34cb341..9fbfeadd41a 100644 --- a/web/src/app/chat/ChatPage.tsx +++ b/web/src/app/chat/ChatPage.tsx @@ -363,6 +363,12 @@ export function ChatPage({ ); const [isStreaming, setIsStreaming] = useState(false); + // Per-conversation search-quality toggles (default OFF, independent of the + // assistant's own settings). Each is also gated server-side by its global + // master switch. Reset per page load — "default off" is intentional. + const [useReranking, setUseReranking] = useState(false); + const [useRelevanceFilter, setUseRelevanceFilter] = useState(false); + // uploaded files const [currentMessageFiles, setCurrentMessageFiles] = useState< FileDescriptor[] @@ -798,6 +804,8 @@ export function ChatPage({ systemPromptOverride: searchParams.get(SEARCH_PARAM_NAMES.SYSTEM_PROMPT) || undefined, useExistingUserMessage: isSeededChat, + useReranking: useReranking, + useRelevanceFilter: useRelevanceFilter, }); const updateFn = (messages: Message[]) => { const replacementsMap = finalMessage @@ -1595,6 +1603,10 @@ export function ChatPage({ } filterManager={filterManager} llmOverrideManager={llmOverrideManager} + useReranking={useReranking} + setUseReranking={setUseReranking} + useRelevanceFilter={useRelevanceFilter} + setUseRelevanceFilter={setUseRelevanceFilter} selectedAssistant={livePersona} files={currentMessageFiles} setFiles={setCurrentMessageFiles} diff --git a/web/src/app/chat/input/ChatInputBar.tsx b/web/src/app/chat/input/ChatInputBar.tsx index 3e5763bcbec..460400a94dd 100644 --- a/web/src/app/chat/input/ChatInputBar.tsx +++ b/web/src/app/chat/input/ChatInputBar.tsx @@ -40,6 +40,10 @@ export function ChatInputBar({ retrievalDisabled, filterManager, llmOverrideManager, + useReranking, + setUseReranking, + useRelevanceFilter, + setUseRelevanceFilter, onSetSelectedAssistant, selectedAssistant, files, @@ -59,6 +63,10 @@ export function ChatInputBar({ retrievalDisabled: boolean; filterManager: FilterManager; llmOverrideManager: LlmOverrideManager; + useReranking: boolean; + setUseReranking: (value: boolean) => void; + useRelevanceFilter: boolean; + setUseRelevanceFilter: (value: boolean) => void; selectedAssistant: Persona; alternativeAssistant: Persona | null; files: FileDescriptor[]; @@ -412,6 +420,36 @@ export function ChatInputBar({ /> )} + {/* Per-conversation search-quality toggles (default off). Only + take effect if enabled globally by an admin. */} + {!retrievalDisabled && ( + + )} + + {!retrievalDisabled && ( + + )} + 0; @@ -161,6 +168,8 @@ export async function* sendMessage({ } : null, use_existing_user_message: useExistingUserMessage, + use_reranking: useReranking ?? false, + use_relevance_filter: useRelevanceFilter ?? false, }), }); if (!sendMessageResponse.ok) { From f77f411733bbf5848f00a8107cacd55283902cfc Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 16:08:48 +0530 Subject: [PATCH 006/115] k8s: serve the reranker via CPU TEI (drop the GPU plan) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - optional/tei-rerank: Hugging Face TEI on CPU (bge-reranker-v2-m3, fp32), Deployment + Service, health probes, model-cache volume. - prod overlay includes it + sets RERANK_ENABLED / RERANK_SERVER_URL / LLM_RELEVANCE_FILTER_ENABLED. Local sets RERANK_ENABLED (no RERANK_SERVER_URL) so the model server loads the reranker in-process — no GPU anywhere. - Removed the optional/gpu-inference component. Co-Authored-By: Claude Opus 4.8 (1M context) --- k8s/optional/tei-rerank/kustomization.yaml | 26 ++++++ k8s/optional/tei-rerank/tei-rerank.yaml | 94 ++++++++++++++++++++++ k8s/overlays/local/env.properties | 7 ++ k8s/overlays/prod/env.properties | 6 ++ k8s/overlays/prod/kustomization.yaml | 3 + 5 files changed, 136 insertions(+) create mode 100644 k8s/optional/tei-rerank/kustomization.yaml create mode 100644 k8s/optional/tei-rerank/tei-rerank.yaml diff --git a/k8s/optional/tei-rerank/kustomization.yaml b/k8s/optional/tei-rerank/kustomization.yaml new file mode 100644 index 00000000000..3e9c9109bc6 --- /dev/null +++ b/k8s/optional/tei-rerank/kustomization.yaml @@ -0,0 +1,26 @@ +# Kustomize Component — CPU-hosted cross-encoder reranker via Hugging Face TEI. +# +# Adds a `tei-rerank` Deployment + Service that serves BAAI/bge-reranker-v2-m3 +# on CPU at full precision (no GPU, no accuracy loss vs GPU — just a bit more +# latency, ~100-250ms for ~20 chunks). +# +# Opting an overlay into this component is the "global enable" half of +# reranking. You MUST ALSO set, in the including overlay's env.properties (so +# they reach the api-server / background pods via env-configmap): +# +# RERANK_ENABLED=true +# RERANK_SERVER_URL=http://tei-rerank-service:80 +# +# With those set, the app reranks for assistants whose `rerank_enabled` flag is +# on (and for chat conversations whose Rerank toggle is on). Our own model +# server then does NOT load the cross-encoder — TEI owns it. +# +# Opt in from an overlay: +# components: +# - ../../optional/tei-rerank +# # + RERANK_ENABLED=true and RERANK_SERVER_URL=... in env.properties +apiVersion: kustomize.config.k8s.io/v1alpha1 +kind: Component + +resources: + - tei-rerank.yaml diff --git a/k8s/optional/tei-rerank/tei-rerank.yaml b/k8s/optional/tei-rerank/tei-rerank.yaml new file mode 100644 index 00000000000..dd9ab02fbc2 --- /dev/null +++ b/k8s/optional/tei-rerank/tei-rerank.yaml @@ -0,0 +1,94 @@ +# Hugging Face Text-Embeddings-Inference (TEI) serving the cross-encoder +# reranker on CPU at full precision (FP32). CPU-optimized (Rust + native token +# batching), so no GPU is needed. The app reaches it via RERANK_SERVER_URL +# (set in the overlay's env.properties) → its /rerank endpoint. +# +# Scaling: stateless — bump `replicas` (or add an HPA on CPU) for more +# rerank throughput. Rule of thumb ~4 vCPU + 4-8Gi per replica, ~1 replica +# per ~5 sustained rerank-QPS. +apiVersion: apps/v1 +kind: Deployment +metadata: + name: tei-rerank-deployment +spec: + replicas: 1 + selector: + matchLabels: + app: tei-rerank + template: + metadata: + labels: + app: tei-rerank + spec: + containers: + - name: tei-rerank + # PIN a real CPU tag (never :latest — see the Vespa lesson). Verify + # the latest available cpu-* tag at ghcr.io/huggingface/text-embeddings-inference. + image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.5 + args: + - "--model-id" + - "BAAI/bge-reranker-v2-m3" + - "--dtype" + - "float32" # full precision: no accuracy loss vs GPU + - "--port" + - "80" + # Bound memory/load so a burst queues instead of OOMing. + - "--max-concurrent-requests" + - "64" + - "--max-batch-tokens" + - "16384" + - "--max-client-batch-size" + - "32" + ports: + - containerPort: 80 + protocol: TCP + env: + # TEI caches the model under /data; back it with a volume so a + # restart doesn't re-download. emptyDir here for simplicity — use a + # PVC in prod to make restarts instant and avoid egress. + - name: HUGGINGFACE_HUB_CACHE + value: /data + resources: + requests: + cpu: "4" + memory: 4Gi + limits: + cpu: "4" # request==limit cpu → predictable rerank latency + memory: 8Gi + readinessProbe: + httpGet: + path: /health + port: 80 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /health + port: 80 + periodSeconds: 30 + # Model download on first boot can take minutes — don't kill it early. + startupProbe: + httpGet: + path: /health + port: 80 + periodSeconds: 10 + failureThreshold: 60 + volumeMounts: + - mountPath: /data + name: tei-model-cache + volumes: + - name: tei-model-cache + emptyDir: {} +--- +apiVersion: v1 +kind: Service +metadata: + name: tei-rerank-service +spec: + ports: + - name: tei-rerank-port + port: 80 + protocol: TCP + targetPort: 80 + selector: + app: tei-rerank + type: ClusterIP diff --git a/k8s/overlays/local/env.properties b/k8s/overlays/local/env.properties index b75f1d918ed..ee84eb01f3c 100644 --- a/k8s/overlays/local/env.properties +++ b/k8s/overlays/local/env.properties @@ -37,6 +37,13 @@ ASYM_QUERY_PREFIX= ASYM_PASSAGE_PREFIX= ENABLE_RERANKING_REAL_TIME_FLOW= ENABLE_RERANKING_ASYNC_FLOW= +# Cross-encoder reranking. Local loads the reranker IN-PROCESS in the model +# server (no TEI container; RERANK_SERVER_URL intentionally unset). Assistants/ +# chats still opt in per their own flag. (Prod serves it via the tei-rerank +# component instead.) +RERANK_ENABLED=true +# LLM relevance filter (LLM-based). Independent of reranking. +LLM_RELEVANCE_FILTER_ENABLED=true # --- LLM --- GEN_AI_MODEL_PROVIDER=custom diff --git a/k8s/overlays/prod/env.properties b/k8s/overlays/prod/env.properties index 331889a5472..10c3fa2e776 100644 --- a/k8s/overlays/prod/env.properties +++ b/k8s/overlays/prod/env.properties @@ -75,6 +75,12 @@ ASYM_QUERY_PREFIX= ASYM_PASSAGE_PREFIX= ENABLE_RERANKING_REAL_TIME_FLOW= ENABLE_RERANKING_ASYNC_FLOW= +# Cross-encoder reranking, served by the tei-rerank component (CPU, no GPU). +# Global enable; assistants/chats still opt in per their own flag. +RERANK_ENABLED=true +RERANK_SERVER_URL=http://tei-rerank-service:80 +# LLM relevance filter (LLM-based, no GPU). Independent of reranking. +LLM_RELEVANCE_FILTER_ENABLED=true # --- LLM --- GEN_AI_MODEL_PROVIDER=custom diff --git a/k8s/overlays/prod/kustomization.yaml b/k8s/overlays/prod/kustomization.yaml index 7e492172b1c..15bae752af0 100644 --- a/k8s/overlays/prod/kustomization.yaml +++ b/k8s/overlays/prod/kustomization.yaml @@ -16,6 +16,9 @@ resources: components: - ../../optional/background-scaling + # CPU-hosted cross-encoder reranker (no GPU). Also set RERANK_ENABLED=true and + # RERANK_SERVER_URL in env.properties for the app to use it. + - ../../optional/tei-rerank namespace: darwin From 15f52818e7bf0d0d5b50328c346112c7211b59d7 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 16:08:48 +0530 Subject: [PATCH 007/115] test+docs: rerank/relevance/diversity tests + design-doc and dev-setup updates - Unit: relevance-filter gating matrix, listwise parser, source-diversity promotion/caps/disable (replaces the removed two-query test). - Integration: TEI rerank transport (mocked), real CPU cross-encoder reordering (MiniLM), filter_chunks with a stub LLM. - Docs updated to the final design (TEI-on-CPU, source diversity at selection, two assistant knobs); CONTRIBUTING documents local in-process reranking. Co-Authored-By: Claude Opus 4.8 (1M context) --- CONTRIBUTING.md | 11 + .../test_filter_chunks_listwise.py | 48 +++ .../integration/test_reranking_cpu_model.py | 67 ++++ .../integration/test_tei_rerank_client.py | 64 ++++ .../vespa/test_query_vespa_prioritization.py | 61 ---- backend/tests/unit/danswer/llm/__init__.py | 0 .../unit/danswer/llm/answering/__init__.py | 0 .../llm/answering/test_source_diversity.py | 68 ++++ .../test_resolve_skip_llm_chunk_filter.py | 70 ++++ .../danswer/secondary_llm_flows/__init__.py | 0 .../test_listwise_chunk_filter.py | 45 +++ docs/how-darwin-answers-questions.md | 310 ++++++++++++++++++ docs/search-quality-reranking-and-recency.md | 272 +++++++++------ 13 files changed, 855 insertions(+), 161 deletions(-) create mode 100644 backend/tests/integration/test_filter_chunks_listwise.py create mode 100644 backend/tests/integration/test_reranking_cpu_model.py create mode 100644 backend/tests/integration/test_tei_rerank_client.py delete mode 100644 backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py create mode 100644 backend/tests/unit/danswer/llm/__init__.py create mode 100644 backend/tests/unit/danswer/llm/answering/__init__.py create mode 100644 backend/tests/unit/danswer/llm/answering/test_source_diversity.py create mode 100644 backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_llm_chunk_filter.py create mode 100644 backend/tests/unit/danswer/secondary_llm_flows/__init__.py create mode 100644 backend/tests/unit/danswer/secondary_llm_flows/test_listwise_chunk_filter.py create mode 100644 docs/how-darwin-answers-questions.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1e70f61b346..ebd2311e2a2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -327,6 +327,17 @@ export INDEXING_MODEL_SERVER_HOST=localhost export INDEXING_MODEL_SERVER_PORT=9000 export REDIS_HOST=cache # matches the compose service name +# Cross-encoder reranking, available locally. The model server (`dmo`) loads the +# reranker IN-PROCESS (sentence-transformers, CPU) — no extra container. Uses the +# small default model (mxbai-rerank-xsmall-v1); set RERANK_MODEL_NAME to try a +# bigger one. Reranking still only runs for assistants / chats that opt in. +# (Prod serves the reranker via a TEI container instead — see k8s/optional/tei-rerank.) +export RERANK_ENABLED=true +export LLM_RELEVANCE_FILTER_ENABLED=true # LLM relevance filter; independent of rerank +# Advanced: to mirror prod and offload the reranker to a local TEI container +# instead of in-process, run TEI yourself and set: +# export RERANK_SERVER_URL=http://localhost:8086 + # --------------------------------------------------------------------------- # LLM (Generative AI) — UiPath LLM Gateway via OAuth client credentials # Replace with your own gateway / model-provider settings if different. diff --git a/backend/tests/integration/test_filter_chunks_listwise.py b/backend/tests/integration/test_filter_chunks_listwise.py new file mode 100644 index 00000000000..87e1b3d5d19 --- /dev/null +++ b/backend/tests/integration/test_filter_chunks_listwise.py @@ -0,0 +1,48 @@ +"""Integration test: the one-shot LLM relevance filter end-to-end (filter_chunks). + +Drives the real filter_chunks → llm_eval_chunks_listwise → _parse_useful_indices +path with a stub LLM (we can't run a real chat model locally), verifying the +listwise selection maps to the right chunk ids and that it fails OPEN. +""" +from types import SimpleNamespace + +from danswer.search.postprocessing.postprocessing import filter_chunks + + +class _StubLLM: + """Minimal LLM whose .invoke returns a message with the given content + (message_to_string only needs `.content` to be a str).""" + + def __init__(self, reply: str) -> None: + self._reply = reply + + def invoke(self, *_: object, **__: object) -> SimpleNamespace: + return SimpleNamespace(content=self._reply) + + +def _chunks(n: int) -> list[SimpleNamespace]: + return [ + SimpleNamespace(content=f"chunk number {i}", unique_id=f"u{i}") + for i in range(n) + ] + + +def _query() -> SimpleNamespace: + return SimpleNamespace(query="some question", max_llm_filter_chunks=15) + + +def test_listwise_selection_maps_to_chunk_ids() -> None: + # LLM says sections 1 and 3 are useful (1-based) → chunks u0 and u2. + out = filter_chunks(_query(), _chunks(3), _StubLLM("[1, 3]")) + assert out == ["u0", "u2"] + + +def test_none_useful() -> None: + out = filter_chunks(_query(), _chunks(3), _StubLLM("[]")) + assert out == [] + + +def test_fail_open_on_unparseable_reply() -> None: + # No JSON array → keep everything (fail open). + out = filter_chunks(_query(), _chunks(3), _StubLLM("I'm not sure.")) + assert out == ["u0", "u1", "u2"] diff --git a/backend/tests/integration/test_reranking_cpu_model.py b/backend/tests/integration/test_reranking_cpu_model.py new file mode 100644 index 00000000000..4489cbc7cae --- /dev/null +++ b/backend/tests/integration/test_reranking_cpu_model.py @@ -0,0 +1,67 @@ +"""Integration test: cross-encoder reranking with a REAL model on CPU. + +Loads a small cross-encoder (ms-marco-MiniLM-L-6-v2 — tiny, CPU-friendly) and +runs it through the real rerank_chunks / semantic_reranking ordering logic to +prove a relevant chunk gets reordered to the top even when retrieval put an +irrelevant one first. Uses the small MiniLM model (not the prod +bge-reranker-v2-m3) because the *ordering logic* is what's under test, and it +keeps the download light. Skips cleanly if the model can't be fetched (offline). +""" +from types import SimpleNamespace + +import pytest + +pytest.importorskip("sentence_transformers") + + +def _load_cross_encoder(): # type: ignore[no-untyped-def] + from sentence_transformers import CrossEncoder + + try: + return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2") + except Exception as exc: # network / offline / disk + pytest.skip(f"cross-encoder model unavailable: {exc}") + + +def _chunk(content: str, score: float) -> SimpleNamespace: + # SimpleNamespace stands in for InferenceChunk — semantic_reranking only + # does attribute access (content, boost, recency_bias, score). + return SimpleNamespace( + content=content, + boost=0, + recency_bias=1.0, + score=score, + document_id="d", + chunk_id=0, + source_links=None, + ) + + +def test_reranker_reorders_relevant_chunk_to_top(monkeypatch) -> None: # type: ignore[no-untyped-def] + model = _load_cross_encoder() + + import danswer.search.postprocessing.postprocessing as pp + + class _RealEnsemble: + def __init__(self, *_: object, **__: object) -> None: + pass + + def predict(self, query: str, passages: list[str]) -> list[list[float]]: + scores = model.predict([(query, p) for p in passages]) + return [[float(s) for s in scores]] + + monkeypatch.setattr(pp, "CrossEncoderEnsembleModel", _RealEnsemble) + + # Deliberately WRONG initial order: the off-topic chunk has the top score. + chunks = [ + _chunk("Bananas are a good source of potassium.", score=0.9), + _chunk("The capital of France is Paris.", score=0.5), + _chunk("Python is a programming language.", score=0.4), + ] + query = SimpleNamespace(query="What is the capital of France?", num_rerank=15) + + ranked = pp.rerank_chunks(query=query, chunks_to_rerank=chunks) + + # After reranking, the France/Paris chunk should be on top, not the banana. + assert "Paris" in ranked[0].content + assert ranked[0].content != "Bananas are a good source of potassium." diff --git a/backend/tests/integration/test_tei_rerank_client.py b/backend/tests/integration/test_tei_rerank_client.py new file mode 100644 index 00000000000..178c5f0eea6 --- /dev/null +++ b/backend/tests/integration/test_tei_rerank_client.py @@ -0,0 +1,64 @@ +"""Integration test for the TEI rerank transport in CrossEncoderEnsembleModel. + +When RERANK_SERVER_URL is configured, the client talks to a Hugging Face TEI +server's /rerank endpoint. TEI returns [{index, score}, ...] sorted by score, +so the client must scatter the scores back into the INPUT passage order and +wrap them as list[list[float]] (the shape semantic_reranking expects). +""" + +import danswer.search.search_nlp_models as nlp + + +class _FakeResponse: + def __init__(self, payload: object) -> None: + self._payload = payload + + def raise_for_status(self) -> None: + pass + + def json(self) -> object: + return self._payload + + +def test_tei_path_scatters_scores_back_to_passage_order(monkeypatch) -> None: # type: ignore[no-untyped-def] + captured: dict = {} + + # TEI replies sorted by score (best first), referencing the input by index. + def fake_post(url: str, json: dict | None = None, **_: object) -> _FakeResponse: + captured["url"] = url + captured["json"] = json + return _FakeResponse( + [ + {"index": 2, "score": 9.0}, + {"index": 0, "score": 1.0}, + {"index": 1, "score": -3.0}, + ] + ) + + monkeypatch.setattr(nlp.requests, "post", fake_post) + + model = nlp.CrossEncoderEnsembleModel(rerank_server_url="http://tei-rerank:80") + out = model.predict("the query", ["a", "b", "c"]) + + # one ensemble entry, scores back in passage order [a, b, c] + assert out == [[1.0, -3.0, 9.0]] + assert captured["url"] == "http://tei-rerank:80/rerank" + assert captured["json"]["texts"] == ["a", "b", "c"] + assert captured["json"]["raw_scores"] is True + + +def test_legacy_path_used_when_no_tei_url(monkeypatch) -> None: # type: ignore[no-untyped-def] + captured: dict = {} + + def fake_post(url: str, json: dict | None = None, **_: object) -> _FakeResponse: + captured["url"] = url + return _FakeResponse({"scores": [[0.1, 0.2]]}) + + monkeypatch.setattr(nlp.requests, "post", fake_post) + + model = nlp.CrossEncoderEnsembleModel(rerank_server_url="") + out = model.predict("q", ["a", "b"]) + + assert out == [[0.1, 0.2]] + # legacy model-server endpoint, NOT /rerank + assert captured["url"].endswith("/encoder/cross-encoder-scores") diff --git a/backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py b/backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py deleted file mode 100644 index 697ef40e531..00000000000 --- a/backend/tests/unit/danswer/document_index/vespa/test_query_vespa_prioritization.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Unit tests for _query_vespa's two paths. - -When prioritize_sources is False (the reranking path) it issues a SINGLE -all-sources Vespa query, so scores stay on one comparable normalize_linear -scale for the cross-encoder. When True (reranking off) it keeps the legacy -two-query flow: an all-sources query plus a second source-filtered query whose -hits are concatenated. The second query independently normalizes a narrower -set, which is why we skip it under reranking. -""" -import pytest - -from danswer.document_index.vespa import index as vespa_index - - -@pytest.fixture -def count_vespa_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict]: - """Replace the network call with a recorder that returns no hits, so the - function exercises its branching without touching Vespa.""" - calls: list[dict] = [] - - def fake_helper(params: dict) -> list: - calls.append(dict(params)) - return [] - - monkeypatch.setattr(vespa_index, "query_vespa_helper", fake_helper) - return calls - - -def test_single_query_when_not_prioritizing(count_vespa_calls: list[dict]) -> None: - out = vespa_index._query_vespa( - {"yql": "select * from sources x where true", "query": "hello"}, - prioritize_sources=False, - ) - assert len(count_vespa_calls) == 1 # one all-sources query, no second query - assert out == [] - # the single query must not have a source_type filter spliced into the YQL - assert "source_type contains" not in count_vespa_calls[0]["yql"] - - -def test_two_queries_when_prioritizing(count_vespa_calls: list[dict]) -> None: - vespa_index._query_vespa( - { - "yql": "select * from sources x where true", - "query": "hello", - "prioritized_sources": ["web"], - }, - prioritize_sources=True, - ) - assert len(count_vespa_calls) == 2 # all-sources + prioritized-sources - # the second (prioritized) query appends the source filter to the YQL - assert "source_type contains" not in count_vespa_calls[0]["yql"] - assert 'source_type contains "web"' in count_vespa_calls[1]["yql"] - - -def test_default_keeps_legacy_two_query_behavior(count_vespa_calls: list[dict]) -> None: - # No explicit flag => prioritize_sources defaults True => legacy behavior, - # so callers other than the reranking path are unaffected. - vespa_index._query_vespa( - {"yql": "select * from sources x where true", "query": "hello"} - ) - assert len(count_vespa_calls) == 2 diff --git a/backend/tests/unit/danswer/llm/__init__.py b/backend/tests/unit/danswer/llm/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/llm/answering/__init__.py b/backend/tests/unit/danswer/llm/answering/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/llm/answering/test_source_diversity.py b/backend/tests/unit/danswer/llm/answering/test_source_diversity.py new file mode 100644 index 00000000000..71ecd5267af --- /dev/null +++ b/backend/tests/unit/danswer/llm/answering/test_source_diversity.py @@ -0,0 +1,68 @@ +"""Unit tests for ensure_source_diversity — guarantees curated KB/web docs +aren't crowded out of the final prompt by a chatty high-relevance source. + +Promotes up to SOURCE_DIVERSITY_RESERVED_SLOTS of the highest-ranked +protected-source docs to the front, preserving the rest of the order. This is +the replacement for the old two-query source-prioritization hack. +""" +from types import SimpleNamespace + +import pytest + +from danswer.llm.answering import doc_pruning as dp + + +def _doc(name: str, source: str) -> SimpleNamespace: + return SimpleNamespace(semantic_identifier=name, source_type=source) + + +def _ids(docs: list[SimpleNamespace]) -> list[str]: + return [d.semantic_identifier for d in docs] + + +@pytest.fixture(autouse=True) +def _cfg(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dp, "PROTECTED_SOURCES", ["web", "sfkbarticles"]) + monkeypatch.setattr(dp, "SOURCE_DIVERSITY_RESERVED_SLOTS", 2) + + +def test_promotes_top_protected_docs_to_front() -> None: + # Slack dominates the top; KB/web are ranked lower and would be cut. + docs = [ + _doc("s1", "slack"), + _doc("s2", "slack"), + _doc("kb1", "sfkbarticles"), + _doc("s3", "slack"), + _doc("w1", "web"), + ] + out = dp.ensure_source_diversity(docs) + # top 2 protected hoisted (in their original relative order); rest preserved + assert _ids(out) == ["kb1", "w1", "s1", "s2", "s3"] + + +def test_caps_at_reserved_slots() -> None: + # 3 protected present, but only the top 2 are promoted. + docs = [ + _doc("s1", "slack"), + _doc("kb1", "web"), + _doc("kb2", "web"), + _doc("kb3", "sfkbarticles"), + ] + out = dp.ensure_source_diversity(docs) + assert _ids(out) == ["kb1", "kb2", "s1", "kb3"] + + +def test_noop_when_no_protected_docs() -> None: + docs = [_doc("s1", "slack"), _doc("s2", "slack")] + assert dp.ensure_source_diversity(docs) is docs + + +def test_disabled_when_reserved_zero(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dp, "SOURCE_DIVERSITY_RESERVED_SLOTS", 0) + docs = [_doc("kb1", "web"), _doc("s1", "slack")] + assert dp.ensure_source_diversity(docs) is docs + + +def test_already_at_front_unchanged() -> None: + docs = [_doc("kb1", "web"), _doc("kb2", "web"), _doc("s1", "slack")] + assert _ids(dp.ensure_source_diversity(docs)) == ["kb1", "kb2", "s1"] diff --git a/backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_llm_chunk_filter.py b/backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_llm_chunk_filter.py new file mode 100644 index 00000000000..07eb2a2ed45 --- /dev/null +++ b/backend/tests/unit/danswer/search/preprocessing/test_resolve_skip_llm_chunk_filter.py @@ -0,0 +1,70 @@ +"""Unit tests for _resolve_skip_llm_chunk_filter — decides whether to skip the +LLM relevance filter. + +Rule: the filter runs only when the global master switch +LLM_RELEVANCE_FILTER_ENABLED AND the per-assistant opt-in +(Persona.llm_relevance_filter) are both on. The global DISABLE_LLM_CHUNK_FILTER +kill-switch always wins. An explicit skip (from the chat flow) is honored unless +the kill-switch is set. Independent of reranking (LLM-only, no GPU). +""" +from types import SimpleNamespace + +import pytest + +from danswer.search.preprocessing import preprocessing as pp + + +def _persona(llm_relevance_filter: bool) -> SimpleNamespace: + return SimpleNamespace(llm_relevance_filter=llm_relevance_filter) + + +@pytest.fixture(autouse=True) +def _reset(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "LLM_RELEVANCE_FILTER_ENABLED", False) + + +def _resolve(explicit, persona, disable=False): # type: ignore[no-untyped-def] + return pp._resolve_skip_llm_chunk_filter(explicit, persona, disable) + + +# --- kill-switch wins --- +def test_kill_switch_forces_skip(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "LLM_RELEVANCE_FILTER_ENABLED", True) + assert _resolve(False, _persona(True), disable=True) is True + + +# --- explicit override (chat) honored when not killed --- +def test_explicit_false_runs_filter() -> None: + assert _resolve(False, _persona(False)) is False + + +def test_explicit_true_skips() -> None: + assert _resolve(True, _persona(True)) is True + + +# --- global x per-assistant matrix (explicit None) --- +def test_global_off_persona_on_skips() -> None: + assert _resolve(None, _persona(True)) is True + + +def test_global_on_persona_off_skips(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "LLM_RELEVANCE_FILTER_ENABLED", True) + assert _resolve(None, _persona(False)) is True + + +def test_global_on_no_persona_skips(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "LLM_RELEVANCE_FILTER_ENABLED", True) + assert _resolve(None, None) is True + + +def test_global_on_persona_on_runs(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(pp, "LLM_RELEVANCE_FILTER_ENABLED", True) + assert _resolve(None, _persona(True)) is False + + +# --- independence from reranking: relevance filter on with rerank irrelevant --- +def test_independent_of_rerank(monkeypatch: pytest.MonkeyPatch) -> None: + # Relevance filter enabled while reranking is globally OFF (no GPU path). + monkeypatch.setattr(pp, "LLM_RELEVANCE_FILTER_ENABLED", True) + monkeypatch.setattr(pp, "RERANK_ENABLED", False) + assert _resolve(None, _persona(True)) is False # filter still runs diff --git a/backend/tests/unit/danswer/secondary_llm_flows/__init__.py b/backend/tests/unit/danswer/secondary_llm_flows/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backend/tests/unit/danswer/secondary_llm_flows/test_listwise_chunk_filter.py b/backend/tests/unit/danswer/secondary_llm_flows/test_listwise_chunk_filter.py new file mode 100644 index 00000000000..4f2deca80a3 --- /dev/null +++ b/backend/tests/unit/danswer/secondary_llm_flows/test_listwise_chunk_filter.py @@ -0,0 +1,45 @@ +"""Unit tests for the listwise (one-shot) relevance filter parsing. + +The filter asks the LLM, in a single call, to return a JSON array of the useful +section numbers. Parsing must: extract the array even amid prose, treat an empty +array as a valid "none useful", clamp out-of-range numbers, and FAIL OPEN (keep +all chunks) when no array can be parsed. +""" +from danswer.secondary_llm_flows.chunk_usefulness import _parse_useful_indices +from danswer.secondary_llm_flows.chunk_usefulness import llm_eval_chunks_listwise + + +def test_parses_array() -> None: + assert _parse_useful_indices("[1, 3, 4]", count=5) == {1, 3, 4} + + +def test_parses_array_amid_prose() -> None: + assert _parse_useful_indices("Sure! The useful ones are [2, 5].", count=5) == {2, 5} + + +def test_empty_array_means_none_useful() -> None: + # Explicit [] is a valid answer (not a parse failure) → empty set. + assert _parse_useful_indices("[]", count=5) == set() + + +def test_out_of_range_clamped() -> None: + assert _parse_useful_indices("[0, 2, 9]", count=3) == {2} + + +def test_no_array_is_parse_failure() -> None: + # None signals the caller to fail OPEN. + assert _parse_useful_indices("I cannot decide.", count=3) is None + + +def test_empty_input_returns_empty() -> None: + assert llm_eval_chunks_listwise("q", [], llm=None) == [] # type: ignore[arg-type] + + +def test_fail_open_on_llm_error() -> None: + class _BoomLLM: + def invoke(self, *_args: object, **_kwargs: object) -> object: + raise RuntimeError("model down") + + # Any exception → keep all chunks (all True). + out = llm_eval_chunks_listwise("q", ["a", "b", "c"], llm=_BoomLLM()) # type: ignore[arg-type] + assert out == [True, True, True] diff --git a/docs/how-darwin-answers-questions.md b/docs/how-darwin-answers-questions.md new file mode 100644 index 00000000000..a54c125effa --- /dev/null +++ b/docs/how-darwin-answers-questions.md @@ -0,0 +1,310 @@ +# How Darwin Answers a Question + +A visual tour of how Darwin turns a question (in **Slack** or the **web chat**) +into a **grounded, cited answer** — and the knobs that change that behavior. +Written for a mixed audience: enough pictures for a stakeholder conversation, +enough specifics that engineers trust it. + +> **TL;DR** — Darwin is **not** a "stuff some chunks into a prompt" toy RAG. +> Every question runs through hybrid retrieval → recency-aware ranking → an +> optional neural reranker → an LLM relevance filter → answer-grounding +> guardrails — all configurable **per assistant**. That depth is where answer +> quality and trust come from. +> +> Engineering deep-dive (file:line, exact ranking math, rollout plan): +> [`search-quality-reranking-and-recency.md`](./search-quality-reranking-and-recency.md). + +--- + +## 1. The moving parts + +```mermaid +flowchart LR + SL["💬 Slack"]:::surface + WEB["🖥️ Web chat"]:::surface + + subgraph CORE["🧩 Answering core"] + direction TB + API["⚙️ API server
orchestration"]:::core + PIPE["🔎 Search pipeline"]:::core + API --> PIPE + end + + MS["🧠 Inference model server
embeddings · intent · reranker"]:::model + VES[("📚 Vespa
vector + keyword index")]:::store + LLM["✨ LLM provider"]:::llm + PG[("🗄️ Postgres")]:::store + REDIS[("⚡ Redis
cache · rate-limit")]:::store + + SL --> API + WEB --> API + PIPE -->|"embed query"| MS + PIPE -->|"hybrid search"| VES + PIPE -->|"generate"| LLM + API -.-> PG + API -.-> REDIS + + subgraph ING["📥 Ingestion · runs continuously"] + direction LR + CONN["🔌 Connectors
Slack · Jira · Web · Salesforce …"]:::surface + DASK["🧵 Dask workers
chunk + embed"]:::core + CONN --> DASK + end + DASK -->|"write chunks + vectors"| VES + + classDef surface fill:#dbeafe,stroke:#2563eb,color:#1e3a8a,stroke-width:1px + classDef core fill:#ede9fe,stroke:#7c3aed,color:#4c1d95,stroke-width:1px + classDef model fill:#fef3c7,stroke:#d97706,color:#7c2d12,stroke-width:1px + classDef store fill:#f1f5f9,stroke:#64748b,color:#334155,stroke-width:1px + classDef llm fill:#fae8ff,stroke:#c026d3,color:#701a75,stroke-width:1px + style CORE fill:#faf5ff,stroke:#a78bfa,color:#4c1d95 + style ING fill:#f0fdf4,stroke:#86efac,color:#14532d +``` + +Two independent halves: +- **📥 Ingestion (always on):** connectors pull documents → Dask workers chunk & + embed them → everything lands in **Vespa**. (Smart enough to *skip* + re-indexing unchanged content.) +- **🧩 Answering (per question):** the path in §2. + +--- + +## 2. The journey of a question + +The same pipeline serves **both** Slack and web chat — they differ only in entry +point and presentation, not in how retrieval/ranking work. + +```mermaid +sequenceDiagram + autonumber + actor U as 👤 User + participant API as ⚙️ API + participant MS as 🧠 Models + participant V as 📚 Vespa + participant R as 🎯 Reranker + participant L as ✨ LLM + + U->>API: question + chosen assistant + Note over API: preprocess — filters,
recency, rerank decision + + rect rgb(220, 252, 231) + Note over API,V: ① RETRIEVE + API->>MS: embed the question + API->>V: hybrid search
(meaning + keywords, recency-weighted) + V-->>API: ≈ 50 candidate chunks + end + + rect rgb(254, 243, 199) + Note over API,R: ② RERANK · only if enabled for this assistant + API->>R: re-score top 15
(question + chunk text together) + R-->>API: reordered by true relevance + end + + rect rgb(250, 232, 255) + Note over API,L: ③ GENERATE + API->>API: relevance filter → keep best ≈ 10 + API->>L: prompt grounded in those chunks + L-->>API: streamed answer + citations + end + + rect rgb(254, 226, 226) + Note over API: 🛡️ guardrail — no citations ⇒ don't answer + end + API-->>U: ✅ grounded answer with sources +``` + +--- + +## 3. The retrieval funnel — where quality comes from + +A question doesn't get "the top chunks." It gets **progressively narrowed** by +increasingly precise (and expensive) stages — broad recall first, sharp +precision last: + +``` + hundreds of thousands of indexed chunks + ████████████████████████████████████████████████ KNOWLEDGE BASE + │ ① hybrid retrieval (meaning + keywords + recency) + ██████████████████████████ ≈ 50 candidates + │ ② cross-encoder rerank (question + chunk together) + ██████████████ top 15, re-scored + │ ③ LLM relevance filter + token budget + ████████ ≈ 10 best chunks + │ ④ grounded generation + citation check + ▼ + ✅ 1 trustworthy answer +``` + +```mermaid +flowchart TB + CORP["📚 Entire knowledge base
(100k+ chunks)"]:::broad + CORP --> S1["① Hybrid retrieval · Vespa
vector + keyword, recency-weighted"]:::retrieve + S1 --> C50(["≈ 50 candidates"]):::count + C50 --> S2["② Cross-encoder reranker
reads question + chunk together"]:::rerank + S2 --> C15(["top 15 re-scored"]):::count + C15 --> S3["③ LLM relevance filter
+ token budget"]:::filter + S3 --> C10(["≈ 10 best chunks"]):::count + C10 --> S4["④ Grounded LLM generation"]:::llm + S4 --> ANS(["✅ 1 cited answer"]):::answer + + classDef broad fill:#f1f5f9,stroke:#64748b,color:#334155 + classDef retrieve fill:#dcfce7,stroke:#16a34a,color:#14532d + classDef rerank fill:#fef3c7,stroke:#d97706,color:#7c2d12 + classDef filter fill:#ffedd5,stroke:#ea580c,color:#7c2d12 + classDef llm fill:#fae8ff,stroke:#c026d3,color:#701a75 + classDef count fill:#ffffff,stroke:#94a3b8,color:#0f172a,stroke-dasharray:3 3 + classDef answer fill:#bbf7d0,stroke:#15803d,color:#14532d,stroke-width:2px +``` + +Why each stage matters: + +| Stage | What it does | Why it's not trivial | +|---|---|---| +| **① Hybrid retrieval** | Finds candidates by **meaning** (vector) *and* **exact terms** (keyword/BM25), fused, then weighted by recency | Pure vector misses exact IDs/error codes; pure keyword misses paraphrases. Fusion catches both. | +| **② Reranking** | A neural **cross-encoder** reads the question and each chunk *together* and re-scores | Retrieval embeds the chunk *before* it sees your question; the reranker judges actual relevance and fixes "right doc, ranked too low." | +| **③ Relevance filter** | An LLM pass drops off-topic chunks before answering | Stops near-misses from diluting the prompt → fewer confident-but-wrong answers. | +| **④ Grounded generation** | Answer built only from selected chunks, with citations | If it can't cite, Darwin **stays silent** rather than hallucinate. | + +--- + +## 4. Hybrid search, in one picture + +Every candidate's score blends two signals, then is nudged by freshness and human +feedback: + +```mermaid +flowchart LR + SEM["🧭 Semantic similarity
meaning match (vectors)"]:::sem + KW["🔤 Keyword match
exact terms (BM25)"]:::kw + SEM -->|"× α"| MIX(("➕ blend")):::mix + KW -->|"× (1 − α)"| MIX + MIX --> BOOST["👍 feedback boost
(promote / bury docs)"]:::boost + BOOST --> REC["🕒 recency factor
(newer scores higher)"]:::rec + REC --> SCORE(["⭐ final relevance score"]):::score + + classDef sem fill:#dbeafe,stroke:#2563eb,color:#1e3a8a + classDef kw fill:#dcfce7,stroke:#16a34a,color:#14532d + classDef mix fill:#ede9fe,stroke:#7c3aed,color:#4c1d95 + classDef boost fill:#fef3c7,stroke:#d97706,color:#7c2d12 + classDef rec fill:#cffafe,stroke:#0891b2,color:#155e75 + classDef score fill:#bbf7d0,stroke:#15803d,color:#14532d,stroke-width:2px +``` + +> `score = ( α · semantic + (1 − α) · keyword ) × feedback_boost × recency` +> — **α** is the dial between meaning and exact terms (default leans semantic). + +This is the *search engine* under the chatbot — a real ranking system, not a +nearest-neighbor lookup. + +--- + +## 5. Recency — preferring fresh knowledge + +Two equally-relevant docs shouldn't tie when one is from last week and one is +three years old. Darwin multiplies each score by a **recency factor** that decays +with document age: + +> `recency_factor = max( 1 / (1 + decay × age_in_years), 0.75 )` + +``` +score multiplier by document age (1.00 = full score) + + new 6 mo 1 yr 2 yr+ +default ██████ █████▍ █████ ████▌ 1.00 → 0.89 → 0.80 → 0.75 (floor) +favor- ██████ █████ ████▌ ████▌ 1.00 → 0.80 → 0.75 → 0.75 +recent +``` + +- A **soft nudge**, not a cliff — a clearly-better old doc can still win. +- The **0.75 floor** caps the penalty at 25% (tunable). Decay strength is set + **per assistant** (`favor_recent` is more aggressive) or globally. +- The lever for "prefer recent" without discarding authoritative older content. + +--- + +## 6. The knobs — same engine, different behavior + +Darwin's behavior is **configured, not hardcoded** — globally and **per +assistant** (each assistant has its own knowledge scope, prompt, and ranking +behavior). + +| Knob | Where | Default | When on | +|---|---|---|---| +| **Neural reranking** | global switch **and** per-assistant toggle | raw hybrid order | top candidates reordered by a cross-encoder (sharper relevance) | +| **Recency preference** | per assistant | mild decay | stronger tilt toward recent docs | +| **LLM relevance filter** | per assistant | — | off-topic chunks dropped pre-answer | +| **Citations required** | per Slack channel | — | no citations ⇒ **no answer** (won't bluff) | +| **Source diversity** | automatic (global) | on | guarantees curated KB/web docs aren't crowded out of the prompt by a chatty source | +| **Guardrails** | built in | — | ACL filtering · rate limiting · retry/backoff | + +```mermaid +flowchart LR + QQ(["❓ Same question"]):::q + QQ --> A1["🅰️ Assistant A
rerank OFF · broad scope"]:::dim + QQ --> A2["🅱️ Assistant B
rerank ON · favor-recent · curated"]:::bright + A1 --> R1(["fast · recall-oriented"]):::dimout + A2 --> R2(["sharper · fresher · more precise"]):::brightout + + classDef q fill:#e0e7ff,stroke:#4f46e5,color:#312e81,stroke-width:2px + classDef dim fill:#f1f5f9,stroke:#94a3b8,color:#475569 + classDef bright fill:#fef3c7,stroke:#d97706,color:#7c2d12,stroke-width:2px + classDef dimout fill:#f8fafc,stroke:#cbd5e1,color:#64748b + classDef brightout fill:#dcfce7,stroke:#16a34a,color:#14532d,stroke-width:2px +``` + +This is what makes a **controlled rollout** possible: enable a capability on +*one* assistant, compare answers side-by-side, then make it the default — no +big-bang switch. + +--- + +## 7. Why this isn't a toy RAG + +A weekend RAG demo is: embed docs → nearest-neighbor → stuff prompt. Darwin adds +the parts that decide whether answers are **trustworthy at scale**: + +```mermaid +flowchart TB + subgraph TOY["🧪 Toy RAG"] + direction TB + T1["vector nearest-neighbor"]:::t --> T2["stuff into prompt"]:::t --> T3["hope it's right"]:::t + end + subgraph DARWIN["🦾 Darwin"] + direction TB + D1["hybrid retrieval · meaning + keywords"]:::d + D2["two-stage ranking · recall → precision"]:::d + D3["recency-aware scoring"]:::d + D4["LLM relevance filtering"]:::d + D5["grounded + cited · suppress if unsure"]:::d + D6["per-assistant config · ACL · rate-limit · retention"]:::d + D1 --> D2 --> D3 --> D4 --> D5 --> D6 + end + classDef t fill:#fee2e2,stroke:#dc2626,color:#7f1d1d + classDef d fill:#dcfce7,stroke:#16a34a,color:#14532d + style TOY fill:#fef2f2,stroke:#fca5a5,color:#7f1d1d + style DARWIN fill:#f0fdf4,stroke:#86efac,color:#14532d +``` + +- **Two-signal hybrid retrieval** (meaning *and* keywords), not just vectors. +- **Two-stage ranking** — cheap recall (bi-encoder) then expensive precision + (cross-encoder) — the standard of serious search systems. +- **Recency-aware ranking** so stale content doesn't masquerade as current. +- **LLM relevance filtering** before generation. +- **Grounded, cited answers with hallucination suppression** — would rather say + nothing than make something up. +- **Per-assistant configurability** — different teams, knowledge scopes, prompts, + and ranking behavior from one platform. +- **Enterprise plumbing** — permissions/ACL, rate limiting, retention, a broad + connector ecosystem, and a horizontally-scaled indexing pipeline. +- **Operational depth** — separate embedding/indexing/reranking model servers, + Redis caching, CPU-optimized reranking (TEI), incremental & measurable rollouts. + +Each layer is a deliberate quality or trust decision. That's the "meat": the gap +between a chatbot that *sounds* right and one you can put in front of the +business. + +--- + +*Companion deep-dive:* +[`search-quality-reranking-and-recency.md`](./search-quality-reranking-and-recency.md) +*(architecture, exact ranking math, file references, rollout plan).* diff --git a/docs/search-quality-reranking-and-recency.md b/docs/search-quality-reranking-and-recency.md index 4052917b022..c3ae2935fe8 100644 --- a/docs/search-quality-reranking-and-recency.md +++ b/docs/search-quality-reranking-and-recency.md @@ -85,8 +85,7 @@ runs across both → a far more accurate relevance judgment. Standard retrieve-broad-then-rerank. **The real flow (corrected mental model):** -1. Vespa returns `NUM_RETURNED_HITS = 50` (+ up to 10 from the prioritized query, - deduped) — **not 15**. +1. Vespa returns `NUM_RETURNED_HITS = 50` (a single all-sources query) — **not 15**. 2. `rerank_chunks` reranks only the **top 15** (`NUM_RERANKED_RESULTS`, `chunks_to_rerank[:num_rerank]`); the rest get `score=None`, appended behind. 3. `semantic_reranking` computes the cross-encoder score, then @@ -147,41 +146,38 @@ from the rerank rollout, and measure independently. --- -## 5. The source-prioritization bias (and the fix) - -`_query_vespa` (`index.py:~705`) historically ran **two** queries and merged them: -``` -Query A: all sources, hits=50 -Query B: + source_type ∈ {web, sfkbarticles}, hits=10 (default-on) -merge = B + A; dedup by (doc_id,chunk_id) keeping MAX score; sort desc -``` -**Why it's a bug:** the rank profile uses `normalize_linear(...)`, which is -min-max **relative to each query's own candidate set**. Query B's narrow set -normalizes its top docs near the ceiling regardless of absolute relevance; dedup -keeps the inflated B-score. ⇒ `web`/`sfkbarticles` are systematically lifted to -the top by a normalization artifact. - -**Interaction with reranking:** rerank only re-scores the **top 15 by this biased -score**, so the bias moves *upstream into candidate selection* — a genuinely -better non-prioritized doc ranked #16 never enters the rerank window. So the hack -**partially undermines** the rerank rollout. - -Secondary bugs in the same function: hardcodes `hits` (ignores -`num_to_retrieve`/persona limit) and ignores `offset` (pagination). - -**The fix (two paths, this branch):** -- **Reranking ON** ⇒ `prioritize_sources=False` ⇒ a **single all-sources query** - (one comparable `normalize_linear` scale, honors the caller's `hits`); the - cross-encoder reorders. -- **Reranking OFF** ⇒ legacy two-query prioritized flow, **byte-for-byte - unchanged**. - -Driven by `prioritize_sources=query.skip_rerank` in `doc_index_retrieval` — so it -rides the same per-assistant + global rerank decision, no separate flag. - -> NOTE: the prioritized-source hack is a deliberate fork divergence. The split -> preserves it whenever reranking is off; if you later remove it entirely, -> confirm the original product intent first (curated web/KB content?). +## 5. Source diversity: keeping KB/web from getting lost + +**The requirement:** if a chatty source (e.g. Slack) produces the +highest-relevance chunks, it shouldn't crowd authoritative **KB/web doc** +content out of the ~10 chunks that reach the LLM. + +**What the fork used to do (removed on this branch):** `_query_vespa` ran **two** +queries and merged them — an all-sources query plus a second, source-filtered +query — to force `web`/`sfkbarticles` in. That was the wrong mechanism: the rank +profile uses `normalize_linear(...)`, which is min-max **relative to each query's +own candidate set**, so the narrow second query's top docs normalized near the +ceiling **regardless of absolute relevance**. Result: prioritized sources were +*over-promoted* by a normalization artifact, and because rerank/relevance only +see the **top-N candidate window**, the inflated ordering polluted what those +stages got to evaluate. + +**What we do instead:** `_query_vespa` now issues a **single, comparably-scored +query**, and source diversity is enforced at **final doc selection** — +`ensure_source_diversity` in `llm/answering/doc_pruning.py` (called from +`_apply_pruning`, after the relevance reorder, before the token-budget cut): + +- It promotes up to **`SOURCE_DIVERSITY_RESERVED_SLOTS`** (default **2**) of the + highest-ranked **`PROTECTED_SOURCES`** (default `web,sfkbarticles`) docs to the + front, preserving the rest of the order. Actual promotion is + `min(reserved, #protected docs present)`. +- It's **always-on and globally configured** (env), operates on the **single + comparably-scored** candidate set (no inflation, plays correctly with rerank), + and is **not a per-assistant decision** — so it doesn't add an assistant knob. +- Disable with `SOURCE_DIVERSITY_RESERVED_SLOTS=0`. + +So the *goal* of the old prioritized-source hack is preserved (KB/web aren't +lost), without the score-inflation bug and without a per-assistant toggle. --- @@ -190,9 +186,9 @@ rides the same per-assistant + global rerank decision, no separate flag. **Two-level gate — rerank runs iff `RERANK_ENABLED` (global) AND `persona.rerank_enabled` (per-assistant).** -- **Global** `RERANK_ENABLED` (env, default false): the master switch. When on, a - GPU-backed model server warms the reranker and the app *may* rerank. Off (local - / default) ⇒ reranking never runs ⇒ **no GPU required**. +- **Global** `RERANK_ENABLED` (env, default false): the master switch. When on, + the reranker is available (served by **TEI on CPU** — see §7 — or a GPU) and + the app *may* rerank. Off (local / default) ⇒ reranking never runs. - **Per-assistant** `Persona.rerank_enabled` (bool, default false): which assistants actually rerank. Lets you enable it on one assistant, compare answers against an un-toggled copy in chat **or** Slack, and flip the default @@ -209,39 +205,84 @@ Both chat and Slack respect it because `SearchTool` builds `SearchRequest` with --- -## 7. Infrastructure / GPU plan - -Observed on the **darwin** cluster (June 2026): -- **No GPU anywhere** (4 nodes, all `nvidia.com/gpu: `). -- The inference model server (the `INDEXING_ONLY=false` pod) does **query - embedding + intent** today (~7.3 GiB RAM), with **no resource requests/limits** - on its container, on a node already at **87% memory** → eviction-prone. -- With rerank off, the cross-encoder is **not even loaded** (`warm_up_cross_encoders` - is gated). So the reranker is a *net-new* model + per-query compute when enabled. - -Decisions: -- **Self-host on a dedicated GPU node**, `Standard_NC6s_v3` (1× V100 16 GB) — more - than enough (a cross-encoder uses ~1 GB; reranks 15 chunks in <50 ms). -- **Co-locate** embedding + intent + reranker on that GPU node ("deploy rerank + - other inference models together on GPU") — maximizes the GPU, accelerates query - embedding, and evacuates the strained CPU node. **Single-GPU-node risk - accepted** (if it dies, search degrades, not just rerank). -- The existing `danswer-model-server` image **already bundles CUDA torch** — no - rebuild; it auto-uses the GPU once scheduled there. -- **Reranker model:** `BAAI/bge-reranker-v2-m3` (env-selectable via - `RERANK_MODEL_NAME`; local keeps the small default). **Switching rerankers needs - NO reindex** — cross-encoders score chunk *text* at query time; only changing - the *embedding* model forces a reindex. (Avoid late-interaction/ColBERT-style - models, which would need indexing changes.) -- **Upstream's stance:** Onyx made the reranker pluggable (default none; local-dev - = mxbai-xsmall) and **disables local reranking when there's no GPU** (PR #4011) - — i.e. don't self-host cross-encoder reranking on CPU. Hence the GPU node. - -k8s: `k8s/optional/gpu-inference/` component pins the inference deployment to the -GPU pool (`nodeSelector: agentpool=gpupool`, toleration `sku=gpu:NoSchedule`, -`nvidia.com/gpu: 1`, real cpu/mem requests+limits — also fixes the no-limits -smell), and sets `RERANK_MODEL_NAME`. The overlay must also set -`RERANK_ENABLED=true` in `env.properties` (reaches every pod via env-configmap). +## 6b. The two assistant knobs + valid combinations + +There are **two per-assistant search-quality knobs**, each gated the same way +(global master switch × per-assistant flag, with a per-conversation chat toggle +that ignores the assistant). Source diversity (§5) is **not** a knob — it's +automatic and globally configured. + +| Knob | Global flag | Per-assistant | Chat toggle (default) | Needs GPU? | +|---|---|---|---|---| +| **Reranking** (cross-encoder) | `RERANK_ENABLED` | `Persona.rerank_enabled` | off | No — TEI on CPU (§7) | +| **LLM relevance filter** (one-shot, **main LLM**) | `LLM_RELEVANCE_FILTER_ENABLED` | `Persona.llm_relevance_filter` | off | **No** (LLM-only) | + +- **LLM relevance filter** is a **single listwise call on the main LLM** + (`llm_eval_chunks_listwise`, fails open on parse/error), not 15 fast-LLM + calls. It needs **no GPU**, so it's a cheaper quality tier on its own. +- **Source diversity** (KB/web protected from being crowded out) is handled + automatically at doc selection — global env (`PROTECTED_SOURCES`, + `SOURCE_DIVERSITY_RESERVED_SLOTS`), no per-assistant or chat decision (§5). +- **Resolution precedence** for the two knobs: chat per-conversation toggle (if + the request set it) → assistant flag → default. Slack/one-shot always use the + assistant flag. + +**Reranking × relevance filter — all four combinations work** (source diversity +applies underneath all of them): + +| `rerank_enabled` | `llm_relevance_filter` | Behavior | Infra | +|---|---|---|---| +| off | off | raw hybrid order (today's default) | none | +| **on** | off | cross-encoder reordering | **TEI-CPU** (or GPU) | +| off | **on** | LLM relevance filter only (1 LLM call) | **none / no GPU** | +| **on** | **on** | rerank **then** relevance-filter | **TEI-CPU** (or GPU) | + +The "relevance-filter-only, no GPU" row is the cheap middle tier; the "rerank-on" +rows need the TEI reranker but still **no GPU**. + +--- + +## 7. Reranker serving: TEI on CPU (no GPU) + +**Decision: serve the reranker on CPU via Hugging Face TEI — no GPU.** + +The darwin cluster has **no GPU** (4 nodes, all `nvidia.com/gpu: `). The +*naive* path (our model server's `sentence_transformers.CrossEncoder` on CPU) is +seconds-slow — which is why upstream Onyx disables local rerank without a GPU +(PR #4011). But that's an artifact of the **unoptimized runtime**, not the model: +with an optimized CPU runtime (**TEI** — Rust, native token batching), the same +`BAAI/bge-reranker-v2-m3` (568M) reranks ~20 chunks in **~100–250 ms on CPU at +full precision** — acceptable for the retrieval phase, and **no accuracy loss** +vs a GPU (CPU vs GPU doesn't change the math; only INT8 *quantization* would, and +we don't use it). + +How it's wired: +- A **`tei-rerank`** deployment (`ghcr.io/huggingface/text-embeddings-inference:cpu-*`) + serves `bge-reranker-v2-m3` at `--dtype float32`, exposing `/rerank`. +- The app's `CrossEncoderEnsembleModel` calls TEI's `/rerank` when + **`RERANK_SERVER_URL`** is set (scattering TEI's score-sorted reply back to + passage order); otherwise it uses the legacy model-server path. When TEI is in + use, our own model server does **not** load the cross-encoder. +- **No reindex** to adopt or switch rerankers — cross-encoders score chunk *text* + at query time; only changing the *embedding* model forces a reindex. (Avoid + late-interaction/ColBERT-style models, which would need indexing changes.) + +Sizing / scaling (per replica): **~4 vCPU, request 4 GiB / limit 8 GiB** (weights +~2.3 GB FP32 + batch headroom). TEI is stateless → scale with replicas or an HPA +on CPU; rule of thumb ~1 replica per ~5 sustained rerank-QPS. Set CPU +request==limit for predictable latency. Optional INT8 later trades a small +accuracy hit for ~2× speed / ~0.6 GB. + +k8s: **`k8s/optional/tei-rerank/`** component (Deployment + Service, CPU +resources, `/health` probes, model-cache volume). Included by **both** the prod +and local overlays. The overlay's `env.properties` sets `RERANK_ENABLED=true`, +`RERANK_SERVER_URL=http://tei-rerank-service:80`, and +`LLM_RELEVANCE_FILTER_ENABLED=true`. (The earlier GPU `gpu-inference` component +was removed in favor of this.) + +> If a GPU is ever desired for lowest latency, the same model runs on a **CUDA** +> node (e.g. `NV6ads_A10_v5` — fractional NVIDIA A10; *not* `NV8as_v4`, whose GPU +> is AMD and unusable by TEI/PyTorch). Not needed for current scale. --- @@ -262,39 +303,69 @@ Backend: - `danswerbot/slack/handlers/handle_message.py` — `skip_rerank=None` (+ dropped the now-unused `ENABLE_RERANKING_ASYNC_FLOW` import). - `model_server/main.py` — warm cross-encoder when `RERANK_ENABLED`. -- `document_index/vespa/index.py` — `_query_vespa(prioritize_sources=...)` single - vs two-query split; `hybrid_retrieval(prioritize_sources=...)`. -- `document_index/interfaces.py` — `hybrid_retrieval` abstract signature. -- `search/retrieval/search_runner.py` — pass - `prioritize_sources=query.skip_rerank`. - -Web (`web/src/app/admin/assistants/`): -- `interfaces.ts`, `lib.ts`, `AssistantEditor.tsx` — "Rerank results (beta)" - toggle, mirroring `llm_relevance_filter`. +- `document_index/vespa/index.py` — `_query_vespa` simplified to a **single + all-sources query** (removed the two-query union). + +LLM relevance filter (independent gate, one-shot, main LLM): +- `configs/chat_configs.py` — `LLM_RELEVANCE_FILTER_ENABLED`. +- `preprocessing.py` — `_resolve_skip_llm_chunk_filter` resolver. +- `prompts/llm_chunk_filter.py` + `secondary_llm_flows/chunk_usefulness.py` — + `LISTWISE_CHUNK_FILTER_PROMPT` + `llm_eval_chunks_listwise` (+ `_parse_useful_indices`). +- `search/pipeline.py` — relevance filter now uses the **main** llm (not fast). +- `search/postprocessing/postprocessing.py` — `filter_chunks` → listwise call. + +Source diversity (automatic, global — replaces the old two-query prioritization): +- `configs/chat_configs.py` — `PROTECTED_SOURCES`, `SOURCE_DIVERSITY_RESERVED_SLOTS`. +- `llm/answering/doc_pruning.py` — `ensure_source_diversity`, called in + `_apply_pruning` after the relevance reorder. + +Chat per-conversation toggles + TEI serving: +- `server/query_and_chat/models.py` — `use_reranking` / `use_relevance_filter` + on `CreateChatMessageRequest`. +- `tools/search/search_tool.py` + `chat/process_message.py` — thread the + per-conversation skips into the `SearchRequest`. +- `shared_configs/configs.py` — `RERANK_SERVER_URL`; `search_nlp_models.py` + `CrossEncoderEnsembleModel` TEI `/rerank` path; `model_server/main.py` skips + loading the cross-encoder when TEI serves it. + +Web: +- `admin/assistants/{interfaces,lib,AssistantEditor}.tsx` — "Rerank results" + checkbox (relevance filter reuses the existing "Apply LLM Relevance Filter"). +- `chat/{lib.tsx,ChatPage.tsx,input/ChatInputBar.tsx}` — two per-conversation + chat toggles (Rerank, Relevance). Infra: -- `k8s/optional/gpu-inference/` — kustomization + inference patch. - -Tests (`backend/tests/unit/...`): -- `search/preprocessing/test_resolve_skip_rerank.py` — global × per-assistant - matrix + explicit override + legacy fallback (7 cases). -- `document_index/vespa/test_query_vespa_prioritization.py` — single-vs-two-query - split + default-is-legacy (3 cases). +- `k8s/optional/tei-rerank/` — CPU TEI reranker component; included by the **prod** + overlay (sets `RERANK_ENABLED` / `RERANK_SERVER_URL` / + `LLM_RELEVANCE_FILTER_ENABLED`). **Local** loads the reranker in-process (no + TEI), via `RERANK_ENABLED` with `RERANK_SERVER_URL` unset. (Replaced the + removed `gpu-inference`.) + +Tests: +- `tests/unit/.../test_resolve_skip_rerank.py`, `test_resolve_skip_llm_chunk_filter.py` + — the two gating matrices. +- `tests/unit/.../test_listwise_chunk_filter.py` — listwise parser. +- `tests/unit/.../test_source_diversity.py` — diversity promotion / caps / disable. +- `tests/integration/` — TEI rerank transport (mocked), **real CPU cross-encoder + reordering** (MiniLM), and `filter_chunks` with a stub LLM. --- ## 9. How to enable in prod (when ready) 1. `alembic upgrade head` (adds `persona.rerank_enabled`) → bounce `dapi` + `dbe`. -2. Add the `Standard_NC6s_v3` GPU node pool (label `agentpool=gpupool`, taint - `sku=gpu:NoSchedule`, NVIDIA device plugin). -3. Prod overlay: add `- ../../optional/gpu-inference` to `components:` **and** set - `RERANK_ENABLED=true` in `env.properties`. Apply. -4. Toggle **"Rerank results"** on one test assistant → A/B compare against an - un-toggled copy in chat + Slack → flip the default once satisfied. - -Local stays GPU-free with zero config: omit the component, leave `RERANK_ENABLED` -unset → reranking never runs. +2. The prod overlay already includes `../../optional/tei-rerank` and sets + `RERANK_ENABLED` / `RERANK_SERVER_URL` / `LLM_RELEVANCE_FILTER_ENABLED` in + `env.properties` — `kubectl apply -k k8s/overlays/prod`. TEI downloads the + model on first boot (back `/data` with a PVC to avoid re-download); **no GPU, + no image rebuild.** (Verify the pinned `cpu-*` TEI image tag first.) +3. Per assistant (admin editor): toggle **Rerank results** and/or **Apply LLM + Relevance Filter**. Or use the **chat-page toggles** to A/B per conversation. + Compare answers, then flip the defaults once satisfied. Source diversity is + automatic (tune via `PROTECTED_SOURCES` / `SOURCE_DIVERSITY_RESERVED_SLOTS`). + +No GPU required at any point. Local mirrors prod (the `local` overlay includes +the same TEI component), so reranking can be exercised locally. --- @@ -303,8 +374,9 @@ unset → reranking never runs. - **Recency tuning is a separate experiment** from reranking — don't bundle. Start with `favor_recent` on the test assistant (config); lower the `0.75` floor only if needed (schema redeploy). Measure independently. -- **Confirm the intent** of the prioritized-source hack before ever removing it - outright (it's preserved whenever reranking is off). +- **Graceful rerank fallback:** if the TEI reranker errors, search currently + degrades rather than falling back to bi-encoder order — worth adding before + enabling rerank by default (see §3 / reliability). - **`enable_auto_detect_filters` is dead** globally (§4) — fixing it would restore LLM time/source filter extraction *and* the `auto` recency path; tracked separately. From 68568e01988f9e3ae7e69b674ef47017f0d43b06 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 17:30:33 +0530 Subject: [PATCH 008/115] k8s: bake bge-reranker-v2-m3 into our own ONNX TEI image MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The TEI CPU image is ONNX-Runtime-based and bge-reranker-v2-m3 ships no ONNX weights, so the stock image 404s on onnx/model.onnx. Our own image exports the model to ONNX at build time (HF Optimum; pinned optimum 1.23.3 + torch 2.2.2 + numpy<2 for the >2GB external-data path) and bakes it at /model — no runtime download, no re-download on restart, no HF dependency. Deployment uses the ACR image with --model-id /model (dropped the emptyDir that shadowed /data and the unneeded istio annotation). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/search-quality-reranking-and-recency.md | 15 ++++--- k8s/optional/tei-rerank/Dockerfile | 45 ++++++++++++++++++++ k8s/optional/tei-rerank/tei-rerank.yaml | 27 ++++-------- 3 files changed, 64 insertions(+), 23 deletions(-) create mode 100644 k8s/optional/tei-rerank/Dockerfile diff --git a/docs/search-quality-reranking-and-recency.md b/docs/search-quality-reranking-and-recency.md index c3ae2935fe8..479f11bf063 100644 --- a/docs/search-quality-reranking-and-recency.md +++ b/docs/search-quality-reranking-and-recency.md @@ -257,8 +257,12 @@ vs a GPU (CPU vs GPU doesn't change the math; only INT8 *quantization* would, an we don't use it). How it's wired: -- A **`tei-rerank`** deployment (`ghcr.io/huggingface/text-embeddings-inference:cpu-*`) - serves `bge-reranker-v2-m3` at `--dtype float32`, exposing `/rerank`. +- A **`tei-rerank`** deployment runs **our own image** (`tei-reranker:bge-v2-m3-*`, + built from `k8s/optional/tei-rerank/Dockerfile`) — TEI's CPU base with + `bge-reranker-v2-m3` **exported to ONNX and baked in** at `/model`. The CPU + TEI runtime is ONNX-Runtime-based and the model ships no ONNX weights, so we + export them at build time (HF Optimum) — this also means **no runtime + download, no re-download on restart, no HuggingFace runtime dependency**. - The app's `CrossEncoderEnsembleModel` calls TEI's `/rerank` when **`RERANK_SERVER_URL`** is set (scattering TEI's score-sorted reply back to passage order); otherwise it uses the legacy model-server path. When TEI is in @@ -356,9 +360,10 @@ Tests: 1. `alembic upgrade head` (adds `persona.rerank_enabled`) → bounce `dapi` + `dbe`. 2. The prod overlay already includes `../../optional/tei-rerank` and sets `RERANK_ENABLED` / `RERANK_SERVER_URL` / `LLM_RELEVANCE_FILTER_ENABLED` in - `env.properties` — `kubectl apply -k k8s/overlays/prod`. TEI downloads the - model on first boot (back `/data` with a PVC to avoid re-download); **no GPU, - no image rebuild.** (Verify the pinned `cpu-*` TEI image tag first.) + `env.properties` — `kubectl apply -k k8s/overlays/prod`. The reranker image + has the ONNX model baked in (build/push it from + `k8s/optional/tei-rerank/Dockerfile`), so it starts fast with **no runtime + download and no GPU**. 3. Per assistant (admin editor): toggle **Rerank results** and/or **Apply LLM Relevance Filter**. Or use the **chat-page toggles** to A/B per conversation. Compare answers, then flip the defaults once satisfied. Source diversity is diff --git a/k8s/optional/tei-rerank/Dockerfile b/k8s/optional/tei-rerank/Dockerfile new file mode 100644 index 00000000000..2817559fd66 --- /dev/null +++ b/k8s/optional/tei-rerank/Dockerfile @@ -0,0 +1,45 @@ +# Our own TEI reranker image with the model exported to ONNX and BAKED IN. +# +# Why ONNX: the TEI CPU image (`cpu-*`) runs on ONNX Runtime and expects an +# `onnx/model.onnx` in the model dir. BAAI/bge-reranker-v2-m3 ships only +# PyTorch/safetensors weights (no onnx/), so the stock image 404s trying to +# download `onnx/model.onnx`. We export it to ONNX with HF Optimum at BUILD +# time and lay it out the way TEI expects, so at runtime TEI loads it locally +# (HF_HUB_OFFLINE=1) — no download, no re-download on restart, no HF dependency. +# +# Build + push (linux/amd64): +# docker build -f k8s/optional/tei-rerank/Dockerfile \ +# -t sfbrdevhelmweacr.azurecr.io/danswer/tei-reranker:bge-v2-m3-2 \ +# --platform linux/amd64 --load k8s/optional/tei-rerank +# docker push sfbrdevhelmweacr.azurecr.io/danswer/tei-reranker:bge-v2-m3-2 + +ARG TEI_BASE=ghcr.io/huggingface/text-embeddings-inference:cpu-1.5 +ARG RERANK_MODEL=BAAI/bge-reranker-v2-m3 + +# Stage 1: export the cross-encoder reranker (sequence-classification, 1 logit) +# to ONNX and arrange it the way TEI's CPU runtime wants — config + tokenizer at +# the model root, weights under onnx/ (model.onnx [+ model.onnx_data for >2GB]). +FROM python:3.11-slim AS export +ARG RERANK_MODEL +# Pin a coherent export toolchain: +# - optimum 1.23.3: `optimum-cli export onnx` is built into [exporters] +# (≥1.24 split it into a separate optimum-onnx package). +# - torch 2.2.2: the stable LEGACY torch.onnx exporter, which names the >2GB +# external-data file the way optimum 1.23.3 expects. Newer torch defaults to +# the dynamo/onnxscript exporter, whose external-data naming breaks optimum +# 1.23.3's cleanup (FileNotFoundError on model.onnx.data). bge-reranker-v2-m3 +# is ~2.3GB fp32 so it always uses external data. +# numpy<2: torch 2.2.2 was built against numpy 1.x; numpy 2.x breaks its +# tensor<->numpy bridge ("Numpy is not available"), which the export VALIDATION +# step uses (the export itself succeeds, but validation then fails the build). +RUN pip install --no-cache-dir "optimum[exporters]==1.23.3" "torch==2.2.2" "numpy<2" onnxruntime onnx +RUN optimum-cli export onnx --model "${RERANK_MODEL}" --task text-classification /model \ + && mkdir -p /model/onnx \ + && mv /model/model.onnx /model/onnx/model.onnx \ + && (mv /model/model.onnx_data /model/onnx/model.onnx_data 2>/dev/null || true) + +# Stage 2: TEI runtime with the exported model baked in. Served from a local +# path so nothing is fetched at runtime. +FROM ${TEI_BASE} +COPY --from=export /model /model +ENV HF_HUB_OFFLINE=1 diff --git a/k8s/optional/tei-rerank/tei-rerank.yaml b/k8s/optional/tei-rerank/tei-rerank.yaml index dd9ab02fbc2..f3da4ef6fbe 100644 --- a/k8s/optional/tei-rerank/tei-rerank.yaml +++ b/k8s/optional/tei-rerank/tei-rerank.yaml @@ -22,12 +22,14 @@ spec: spec: containers: - name: tei-rerank - # PIN a real CPU tag (never :latest — see the Vespa lesson). Verify - # the latest available cpu-* tag at ghcr.io/huggingface/text-embeddings-inference. - image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.5 + # Our own image with BAAI/bge-reranker-v2-m3 baked in (see Dockerfile). + # No runtime HuggingFace download → fast, reliable, offline; sidesteps + # the stock-image download bug and the re-download-on-restart problem. + image: sfbrdevhelmweacr.azurecr.io/danswer/tei-reranker:bge-v2-m3-2 args: + # Local path — the ONNX model is baked into the image at /model. - "--model-id" - - "BAAI/bge-reranker-v2-m3" + - "/model" - "--dtype" - "float32" # full precision: no accuracy loss vs GPU - "--port" @@ -42,12 +44,6 @@ spec: ports: - containerPort: 80 protocol: TCP - env: - # TEI caches the model under /data; back it with a volume so a - # restart doesn't re-download. emptyDir here for simplicity — use a - # PVC in prod to make restarts instant and avoid egress. - - name: HUGGINGFACE_HUB_CACHE - value: /data resources: requests: cpu: "4" @@ -65,19 +61,14 @@ spec: path: /health port: 80 periodSeconds: 30 - # Model download on first boot can take minutes — don't kill it early. + # Model is baked in (loaded from the image's /data), so startup is + # just model load — but keep headroom for CPU model init. startupProbe: httpGet: path: /health port: 80 periodSeconds: 10 - failureThreshold: 60 - volumeMounts: - - mountPath: /data - name: tei-model-cache - volumes: - - name: tei-model-cache - emptyDir: {} + failureThreshold: 30 --- apiVersion: v1 kind: Service From fc3c6541e96092bfe9e5a83793b4e34d45037360 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 17:30:33 +0530 Subject: [PATCH 009/115] k8s(prod): bump backend vha-148, web vha-78 (rerank/relevance/diversity) Built + deployed from feature/improve-queries; api-server self-migrated persona.rerank_enabled on rollout. Validated live: TEI reranker healthy and scoring correctly. Co-Authored-By: Claude Opus 4.8 (1M context) --- k8s/overlays/prod/kustomization.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/k8s/overlays/prod/kustomization.yaml b/k8s/overlays/prod/kustomization.yaml index 15bae752af0..06767df7e85 100644 --- a/k8s/overlays/prod/kustomization.yaml +++ b/k8s/overlays/prod/kustomization.yaml @@ -28,10 +28,10 @@ namespace: darwin images: - name: danswer-backend newName: sfbrdevhelmweacr.azurecr.io/danswer/danswer-backend - newTag: vha-147 + newTag: vha-148 - name: danswer-web-server newName: sfbrdevhelmweacr.azurecr.io/danswer/danswer-web-server - newTag: vha-77 + newTag: vha-78 - name: danswer-model-server newName: danswer/danswer-model-server newTag: v0.3.94 From af74ee7cc0a9eb488c52c3db521ed610b279858e Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Sun, 7 Jun 2026 19:02:18 +0530 Subject: [PATCH 010/115] k8s: serve the reranker on GPU via the upstream TEI image CPU serving was validated as non-viable (24-98s for 15 chunks on prod; the long blocking inference also starved /health and got the pod liveness-killed). Move reranking to a T4 GPU node pool (Standard_NC4as_T4_v3, tainted gpu=true:NoSchedule) where the same bge-reranker-v2-m3 reranks in ~0.6-3.2s (45ms for short inputs). - Use the UPSTREAM TEI GPU image directly (ghcr turing-1.5) instead of building our own: TEI's GPU backend is Candle+safetensors, which the model ships, so the ONNX-export/custom-image dance (a CPU-runtime-only constraint) is gone. Deleted the custom Dockerfile. - Take the pod out of the istio mesh (sidecar.istio.io/inject: false): leaf service, PERMISSIVE mTLS so the api-server still reaches it, and an injected sidecar isn't up during the init phase (so the prefetch init container couldn't reach the network). - Prefetch the model into the PVC with the Python HF client in an init container: TEI's Rust hf-hub client fails on HF's redirect with 'relative URL without a base'; the Python client handles it. TEI loads offline from /data/model. Cluster-side (not in repo): added a gpu=true:NoSchedule toleration to the nvidia-device-plugin-daemonset so it schedules on the tainted GPU node and advertises nvidia.com/gpu. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/search-quality-reranking-and-recency.md | 95 +++++++++-------- k8s/optional/tei-rerank/Dockerfile | 45 -------- k8s/optional/tei-rerank/tei-rerank.yaml | 106 +++++++++++++++---- 3 files changed, 135 insertions(+), 111 deletions(-) delete mode 100644 k8s/optional/tei-rerank/Dockerfile diff --git a/docs/search-quality-reranking-and-recency.md b/docs/search-quality-reranking-and-recency.md index 479f11bf063..48a9c95eb9f 100644 --- a/docs/search-quality-reranking-and-recency.md +++ b/docs/search-quality-reranking-and-recency.md @@ -233,36 +233,43 @@ applies underneath all of them): | `rerank_enabled` | `llm_relevance_filter` | Behavior | Infra | |---|---|---|---| | off | off | raw hybrid order (today's default) | none | -| **on** | off | cross-encoder reordering | **TEI-CPU** (or GPU) | +| **on** | off | cross-encoder reordering | **TEI on GPU** | | off | **on** | LLM relevance filter only (1 LLM call) | **none / no GPU** | -| **on** | **on** | rerank **then** relevance-filter | **TEI-CPU** (or GPU) | +| **on** | **on** | rerank **then** relevance-filter | **TEI on GPU** | -The "relevance-filter-only, no GPU" row is the cheap middle tier; the "rerank-on" -rows need the TEI reranker but still **no GPU**. +The "relevance-filter-only, no GPU" row is the cheap middle tier (just one extra +LLM call); the "rerank-on" rows need the **GPU** TEI reranker (CPU was measured at +24–98 s/query — see §7). --- -## 7. Reranker serving: TEI on CPU (no GPU) +## 7. Reranker serving: TEI on GPU (NVIDIA T4) -**Decision: serve the reranker on CPU via Hugging Face TEI — no GPU.** +**Decision: serve the reranker on GPU via the upstream Hugging Face TEI image.** -The darwin cluster has **no GPU** (4 nodes, all `nvidia.com/gpu: `). The -*naive* path (our model server's `sentence_transformers.CrossEncoder` on CPU) is -seconds-slow — which is why upstream Onyx disables local rerank without a GPU -(PR #4011). But that's an artifact of the **unoptimized runtime**, not the model: -with an optimized CPU runtime (**TEI** — Rust, native token batching), the same -`BAAI/bge-reranker-v2-m3` (568M) reranks ~20 chunks in **~100–250 ms on CPU at -full precision** — acceptable for the retrieval phase, and **no accuracy loss** -vs a GPU (CPU vs GPU doesn't change the math; only INT8 *quantization* would, and -we don't use it). +We first tried **CPU** (the cluster had no GPU). It was functionally correct but +**not interactive-viable**: measured live on prod (`bge-reranker-v2-m3`, 568M, +fp32, 4 vCPU) reranking 15 chunks took **24 s** (passages ≤512 tok) to **98 s** +(longer passages) — ~1.6 s/passage — and the long blocking inference starved the +`/health` endpoint, so the liveness probe killed the pod (503s under load). +`--auto-truncate` caps the tail but not the ~24 s floor; replicas add concurrency, +not single-query speed; fp16 doesn't help on CPU (x86 has no native fp16 matmul — +ORT upcasts to fp32); int8 is only ~2–4×. So reranking moved to **GPU**, where the +same model reranks 15 chunks in **~20–80 ms**. + +Why the **upstream image, no custom build**: TEI's **GPU** backend is Candle + +**safetensors**, which `bge-reranker-v2-m3` ships — so the model loads directly. +The custom-image/ONNX-export dance was *only* needed by TEI's **CPU** runtime +(ONNX-Runtime-based; the model has no ONNX weights). On GPU that constraint is +gone, so we point straight at `ghcr.io/huggingface/text-embeddings-inference:turing-1.5` +(`turing` == T4, compute 7.5; use `86-1.5` for A10, `1.5` for A100/H100). How it's wired: -- A **`tei-rerank`** deployment runs **our own image** (`tei-reranker:bge-v2-m3-*`, - built from `k8s/optional/tei-rerank/Dockerfile`) — TEI's CPU base with - `bge-reranker-v2-m3` **exported to ONNX and baked in** at `/model`. The CPU - TEI runtime is ONNX-Runtime-based and the model ships no ONNX weights, so we - export them at build time (HF Optimum) — this also means **no runtime - download, no re-download on restart, no HuggingFace runtime dependency**. +- A **`tei-rerank`** Deployment (upstream GPU image, `--dtype float16`) on the + tainted GPU node pool — it tolerates `gpu=true:NoSchedule` and requests + `nvidia.com/gpu: 1` (only GPU nodes advertise it, which also pins scheduling). + Model weights download once into a PVC-backed HF cache (`/data`); restarts reuse + it (no re-download), so no custom image is needed to avoid re-downloads. - The app's `CrossEncoderEnsembleModel` calls TEI's `/rerank` when **`RERANK_SERVER_URL`** is set (scattering TEI's score-sorted reply back to passage order); otherwise it uses the legacy model-server path. When TEI is in @@ -271,18 +278,18 @@ How it's wired: at query time; only changing the *embedding* model forces a reindex. (Avoid late-interaction/ColBERT-style models, which would need indexing changes.) -Sizing / scaling (per replica): **~4 vCPU, request 4 GiB / limit 8 GiB** (weights -~2.3 GB FP32 + batch headroom). TEI is stateless → scale with replicas or an HPA -on CPU; rule of thumb ~1 replica per ~5 sustained rerank-QPS. Set CPU -request==limit for predictable latency. Optional INT8 later trades a small -accuracy hit for ~2× speed / ~0.6 GB. +Node pool: **`Standard_NC4as_T4_v3`** (1× T4 16 GB, 4 vCPU, ~$480/mo) tainted +`gpu=true:NoSchedule`. The reranker needs only ~1.5–2 GB VRAM, so the T4 is ample +and GPU compute is never the bottleneck. TEI is stateless → scale with replicas / +an HPA for throughput. -k8s: **`k8s/optional/tei-rerank/`** component (Deployment + Service, CPU -resources, `/health` probes, model-cache volume). Included by **both** the prod -and local overlays. The overlay's `env.properties` sets `RERANK_ENABLED=true`, +k8s: **`k8s/optional/tei-rerank/`** component (PVC + Deployment + Service, GPU +request + taint toleration, `/health` probes). Included by the **prod** overlay +(local dev loads the cross-encoder in-process — no TEI container). The overlay's +`env.properties` sets `RERANK_ENABLED=true`, `RERANK_SERVER_URL=http://tei-rerank-service:80`, and -`LLM_RELEVANCE_FILTER_ENABLED=true`. (The earlier GPU `gpu-inference` component -was removed in favor of this.) +`LLM_RELEVANCE_FILTER_ENABLED=true`. (The earlier CPU `tei-rerank` image and its +ONNX-export Dockerfile were removed in favor of this.) > If a GPU is ever desired for lowest latency, the same model runs on a **CUDA** > node (e.g. `NV6ads_A10_v5` — fractional NVIDIA A10; *not* `NV8as_v4`, whose GPU @@ -339,11 +346,11 @@ Web: chat toggles (Rerank, Relevance). Infra: -- `k8s/optional/tei-rerank/` — CPU TEI reranker component; included by the **prod** - overlay (sets `RERANK_ENABLED` / `RERANK_SERVER_URL` / - `LLM_RELEVANCE_FILTER_ENABLED`). **Local** loads the reranker in-process (no - TEI), via `RERANK_ENABLED` with `RERANK_SERVER_URL` unset. (Replaced the - removed `gpu-inference`.) +- `k8s/optional/tei-rerank/` — GPU TEI reranker component (upstream TEI image on a + T4 node pool); included by the **prod** overlay (sets `RERANK_ENABLED` / + `RERANK_SERVER_URL` / `LLM_RELEVANCE_FILTER_ENABLED`). **Local** loads the + reranker in-process (no TEI, no GPU), via `RERANK_ENABLED` with + `RERANK_SERVER_URL` unset. Tests: - `tests/unit/.../test_resolve_skip_rerank.py`, `test_resolve_skip_llm_chunk_filter.py` @@ -358,19 +365,19 @@ Tests: ## 9. How to enable in prod (when ready) 1. `alembic upgrade head` (adds `persona.rerank_enabled`) → bounce `dapi` + `dbe`. -2. The prod overlay already includes `../../optional/tei-rerank` and sets - `RERANK_ENABLED` / `RERANK_SERVER_URL` / `LLM_RELEVANCE_FILTER_ENABLED` in - `env.properties` — `kubectl apply -k k8s/overlays/prod`. The reranker image - has the ONNX model baked in (build/push it from - `k8s/optional/tei-rerank/Dockerfile`), so it starts fast with **no runtime - download and no GPU**. +2. Add a GPU node pool tainted `gpu=true:NoSchedule` (prod uses + `Standard_NC4as_T4_v3`). The prod overlay already includes + `../../optional/tei-rerank` and sets `RERANK_ENABLED` / `RERANK_SERVER_URL` / + `LLM_RELEVANCE_FILTER_ENABLED` in `env.properties` — `kubectl apply -k + k8s/overlays/prod`. TEI pulls the upstream GPU image and downloads the model + once into its PVC-backed cache. 3. Per assistant (admin editor): toggle **Rerank results** and/or **Apply LLM Relevance Filter**. Or use the **chat-page toggles** to A/B per conversation. Compare answers, then flip the defaults once satisfied. Source diversity is automatic (tune via `PROTECTED_SOURCES` / `SOURCE_DIVERSITY_RESERVED_SLOTS`). -No GPU required at any point. Local mirrors prod (the `local` overlay includes -the same TEI component), so reranking can be exercised locally. +Reranking needs the GPU node; the **relevance filter alone needs no GPU**. Local +exercises reranking in-process (no GPU) for dev. --- diff --git a/k8s/optional/tei-rerank/Dockerfile b/k8s/optional/tei-rerank/Dockerfile deleted file mode 100644 index 2817559fd66..00000000000 --- a/k8s/optional/tei-rerank/Dockerfile +++ /dev/null @@ -1,45 +0,0 @@ -# Our own TEI reranker image with the model exported to ONNX and BAKED IN. -# -# Why ONNX: the TEI CPU image (`cpu-*`) runs on ONNX Runtime and expects an -# `onnx/model.onnx` in the model dir. BAAI/bge-reranker-v2-m3 ships only -# PyTorch/safetensors weights (no onnx/), so the stock image 404s trying to -# download `onnx/model.onnx`. We export it to ONNX with HF Optimum at BUILD -# time and lay it out the way TEI expects, so at runtime TEI loads it locally -# (HF_HUB_OFFLINE=1) — no download, no re-download on restart, no HF dependency. -# -# Build + push (linux/amd64): -# docker build -f k8s/optional/tei-rerank/Dockerfile \ -# -t sfbrdevhelmweacr.azurecr.io/danswer/tei-reranker:bge-v2-m3-2 \ -# --platform linux/amd64 --load k8s/optional/tei-rerank -# docker push sfbrdevhelmweacr.azurecr.io/danswer/tei-reranker:bge-v2-m3-2 - -ARG TEI_BASE=ghcr.io/huggingface/text-embeddings-inference:cpu-1.5 -ARG RERANK_MODEL=BAAI/bge-reranker-v2-m3 - -# Stage 1: export the cross-encoder reranker (sequence-classification, 1 logit) -# to ONNX and arrange it the way TEI's CPU runtime wants — config + tokenizer at -# the model root, weights under onnx/ (model.onnx [+ model.onnx_data for >2GB]). -FROM python:3.11-slim AS export -ARG RERANK_MODEL -# Pin a coherent export toolchain: -# - optimum 1.23.3: `optimum-cli export onnx` is built into [exporters] -# (≥1.24 split it into a separate optimum-onnx package). -# - torch 2.2.2: the stable LEGACY torch.onnx exporter, which names the >2GB -# external-data file the way optimum 1.23.3 expects. Newer torch defaults to -# the dynamo/onnxscript exporter, whose external-data naming breaks optimum -# 1.23.3's cleanup (FileNotFoundError on model.onnx.data). bge-reranker-v2-m3 -# is ~2.3GB fp32 so it always uses external data. -# numpy<2: torch 2.2.2 was built against numpy 1.x; numpy 2.x breaks its -# tensor<->numpy bridge ("Numpy is not available"), which the export VALIDATION -# step uses (the export itself succeeds, but validation then fails the build). -RUN pip install --no-cache-dir "optimum[exporters]==1.23.3" "torch==2.2.2" "numpy<2" onnxruntime onnx -RUN optimum-cli export onnx --model "${RERANK_MODEL}" --task text-classification /model \ - && mkdir -p /model/onnx \ - && mv /model/model.onnx /model/onnx/model.onnx \ - && (mv /model/model.onnx_data /model/onnx/model.onnx_data 2>/dev/null || true) - -# Stage 2: TEI runtime with the exported model baked in. Served from a local -# path so nothing is fetched at runtime. -FROM ${TEI_BASE} -COPY --from=export /model /model -ENV HF_HUB_OFFLINE=1 diff --git a/k8s/optional/tei-rerank/tei-rerank.yaml b/k8s/optional/tei-rerank/tei-rerank.yaml index f3da4ef6fbe..cc68144ea71 100644 --- a/k8s/optional/tei-rerank/tei-rerank.yaml +++ b/k8s/optional/tei-rerank/tei-rerank.yaml @@ -1,17 +1,44 @@ -# Hugging Face Text-Embeddings-Inference (TEI) serving the cross-encoder -# reranker on CPU at full precision (FP32). CPU-optimized (Rust + native token -# batching), so no GPU is needed. The app reaches it via RERANK_SERVER_URL -# (set in the overlay's env.properties) → its /rerank endpoint. +# HuggingFace Text-Embeddings-Inference (TEI) serving the cross-encoder reranker +# on GPU (NVIDIA T4). The app reaches it via RERANK_SERVER_URL → /rerank. # -# Scaling: stateless — bump `replicas` (or add an HPA on CPU) for more -# rerank throughput. Rule of thumb ~4 vCPU + 4-8Gi per replica, ~1 replica -# per ~5 sustained rerank-QPS. +# Why the UPSTREAM image (no custom build): TEI's GPU backend is Candle + +# safetensors, and BAAI/bge-reranker-v2-m3 ships safetensors — so the model +# loads directly, no ONNX export needed. (That dance was only required by TEI's +# *CPU* runtime, which is ONNX-Runtime-based and the model has no ONNX weights.) +# CPU serving was also far too slow (~1.6s/passage, 24–98s for 15 chunks), so +# reranking lives on GPU (~20–80ms). +# +# Why NO istio sidecar: TEI is a leaf service that needs no mesh features, mTLS +# is PERMISSIVE (the api-server still reaches it), and — critically — an istio +# sidecar isn't running during the init phase, so the model-prefetch init +# container below couldn't reach the network with the sidecar injected. +# +# Why prefetch the model in an init container (vs letting TEI download it): TEI's +# Rust hf-hub client fails on Hugging Face's redirect with "relative URL without +# a base". The Python huggingface_hub client resolves it fine, so we prefetch the +# weights into the PVC and run TEI offline against the local path. +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: tei-rerank-model-cache +spec: + accessModes: ["ReadWriteOnce"] + # No storageClassName → cluster default SC (AKS managed-csi, + # volumeBindingMode=WaitForFirstConsumer, so the disk lands in the GPU node's + # zone — no cross-zone attach conflict). + resources: + requests: + storage: 10Gi +--- apiVersion: apps/v1 kind: Deployment metadata: name: tei-rerank-deployment spec: replicas: 1 + strategy: + # RWO model-cache PVC can't be mounted by an old + new pod simultaneously. + type: Recreate selector: matchLabels: app: tei-rerank @@ -19,38 +46,70 @@ spec: metadata: labels: app: tei-rerank + annotations: + # Leaf service — keep it out of the mesh (see header). + sidecar.istio.io/inject: "false" spec: + # Tolerate the GPU pool's taint (gpu=true:NoSchedule). + tolerations: + - key: gpu + operator: Equal + value: "true" + effect: NoSchedule + initContainers: + # Prefetch the model into the PVC with the (robust) Python HF client. + # Idempotent: skips the download when the weights are already cached. + - name: fetch-model + image: python:3.11-slim + command: ["sh", "-c"] + args: + - | + set -e + if [ ! -f /data/model/model.safetensors ]; then + pip install --no-cache-dir huggingface_hub + python -c "from huggingface_hub import snapshot_download; snapshot_download('BAAI/bge-reranker-v2-m3', local_dir='/data/model')" + else + echo "model already cached at /data/model" + fi + volumeMounts: + - name: model-cache + mountPath: /data containers: - name: tei-rerank - # Our own image with BAAI/bge-reranker-v2-m3 baked in (see Dockerfile). - # No runtime HuggingFace download → fast, reliable, offline; sidesteps - # the stock-image download bug and the re-download-on-restart problem. - image: sfbrdevhelmweacr.azurecr.io/danswer/tei-reranker:bge-v2-m3-2 + # Upstream TEI GPU image, pinned. `turing-1.5` == NVIDIA T4 (compute + # cap 7.5). For an A10 node switch to `86-1.5`; A100/H100 → `1.5`. + image: ghcr.io/huggingface/text-embeddings-inference:turing-1.5 args: - # Local path — the ONNX model is baked into the image at /model. + # Local path — prefetched by the init container; no runtime download. - "--model-id" - - "/model" + - "/data/model" - "--dtype" - - "float32" # full precision: no accuracy loss vs GPU + - "float16" # GPU: fp16, ~no quality loss vs fp32 for this model - "--port" - "80" - # Bound memory/load so a burst queues instead of OOMing. - "--max-concurrent-requests" - "64" - "--max-batch-tokens" - "16384" - "--max-client-batch-size" - "32" + env: + - name: HF_HUB_OFFLINE + value: "1" # belt-and-suspenders: never touch the hub at runtime ports: - containerPort: 80 protocol: TCP resources: requests: - cpu: "4" - memory: 4Gi + cpu: "2" + memory: 6Gi limits: - cpu: "4" # request==limit cpu → predictable rerank latency - memory: 8Gi + # Requesting a GPU forces this pod onto the GPU node pool. + nvidia.com/gpu: 1 + memory: 12Gi + volumeMounts: + - name: model-cache + mountPath: /data readinessProbe: httpGet: path: /health @@ -61,14 +120,17 @@ spec: path: /health port: 80 periodSeconds: 30 - # Model is baked in (loaded from the image's /data), so startup is - # just model load — but keep headroom for CPU model init. + # First boot prefetches the model (init) then loads it; generous. startupProbe: httpGet: path: /health port: 80 periodSeconds: 10 - failureThreshold: 30 + failureThreshold: 60 + volumes: + - name: model-cache + persistentVolumeClaim: + claimName: tei-rerank-model-cache --- apiVersion: v1 kind: Service From 26c16009d59a693420b4fe1b441c612fe1bc54f6 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Tue, 9 Jun 2026 18:34:04 +0530 Subject: [PATCH 011/115] k8s(prod): make source-diversity config explicit in env.properties MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PROTECTED_SOURCES and SOURCE_DIVERSITY_RESERVED_SLOTS were running on code defaults (web,sfkbarticles / 2), so the always-on diversity logic was active but invisible in the configmap. Pin them to the current defaults for visibility and tunability — no behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- k8s/overlays/prod/env.properties | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/k8s/overlays/prod/env.properties b/k8s/overlays/prod/env.properties index 10c3fa2e776..c46e274586b 100644 --- a/k8s/overlays/prod/env.properties +++ b/k8s/overlays/prod/env.properties @@ -81,6 +81,12 @@ RERANK_ENABLED=true RERANK_SERVER_URL=http://tei-rerank-service:80 # LLM relevance filter (LLM-based, no GPU). Independent of reranking. LLM_RELEVANCE_FILTER_ENABLED=true +# Source diversity at final doc selection: reserve up to N front slots for the +# highest-ranked docs from PROTECTED_SOURCES so curated content isn't crowded +# out of the prompt by a chatty source (e.g. Slack). These match the code +# defaults — set explicitly here for visibility/tuning (set SLOTS=0 to disable). +PROTECTED_SOURCES=web,sfkbarticles +SOURCE_DIVERSITY_RESERVED_SLOTS=2 # --- LLM --- GEN_AI_MODEL_PROVIDER=custom From adb4dac94334373c491e5d19cd359a895d54942f Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Wed, 10 Jun 2026 15:52:18 +0530 Subject: [PATCH 012/115] fix(auth): authorize valid X-API-Key requests under enforced auth (OIDC) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These api keys are service credentials for automation and intentionally don't map to a User. Enabling OIDC (AUTH_TYPE=oidc) flipped DISABLE_AUTH off, so current_user started 403'ing api-key requests (no session, and the keys don't resolve to a user) — bouncing automation into the SSO login flow. current_user now authorizes a request carrying a valid X-API-Key as an anonymous service caller (user=None, which endpoints already handle). Browser requests without a session still 403 (SSO gate intact), and a key alone does not grant admin (current_admin_user still requires an admin user). Adds request_has_valid_api_key() mirroring validate_api_key's lookup+cache, plus integration tests locking both the SSO and api-key flows. Co-Authored-By: Claude Opus 4.8 (1M context) --- backend/danswer/auth/api_key.py | 32 ++++ backend/danswer/auth/users.py | 12 ++ .../integration/test_auth_api_key_and_sso.py | 153 ++++++++++++++++++ 3 files changed, 197 insertions(+) create mode 100644 backend/tests/integration/test_auth_api_key_and_sso.py diff --git a/backend/danswer/auth/api_key.py b/backend/danswer/auth/api_key.py index 15b372cc73e..6b4dd778d33 100644 --- a/backend/danswer/auth/api_key.py +++ b/backend/danswer/auth/api_key.py @@ -40,3 +40,35 @@ def validate_api_key(request: Request, db_session: Session = Depends(get_session # Cache it for future requests cache[api_key_value] = True return None + + +def request_has_valid_api_key(request: Request, db_session: Session) -> bool: + """Return True if the request carries a valid X-API-Key. + + These keys are service credentials for automation (they intentionally do NOT + map to a browser `User`). `current_user` uses this to authorize an api-key + request as an anonymous service caller instead of 403'ing it into the SSO + flow once AUTH_TYPE enforces auth (e.g. OIDC). Mirrors `validate_api_key`'s + lookup + cache exactly, so the two stay consistent. + + NOTE: `db_session` is passed in (not a Depends) because the caller already + holds a session. + """ + if _API_KEY_HEADER not in request.headers: + return False + + api_key_value = request.headers.get(_API_KEY_HEADER) + if not api_key_value: + return False + + if api_key_value in cache: + return True + + api_key = db_session.scalar( + select(ApiKey).where(ApiKey.hashed_api_key == api_key_value) + ) + if api_key is None: + return False + + cache[api_key_value] = True + return True diff --git a/backend/danswer/auth/users.py b/backend/danswer/auth/users.py index 5605fdbde35..ebc4ec92339 100644 --- a/backend/danswer/auth/users.py +++ b/backend/danswer/auth/users.py @@ -26,6 +26,7 @@ from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase from sqlalchemy.orm import Session +from danswer.auth.api_key import request_has_valid_api_key from danswer.auth.invited_users import get_invited_users from danswer.auth.schemas import UserCreate from danswer.auth.schemas import UserRole @@ -354,8 +355,19 @@ async def double_check_user( async def current_user( + request: Request, user: User | None = Depends(optional_user), + db_session: Session = Depends(get_session), ) -> User | None: + # API keys are service credentials for automation (not browser users) and + # intentionally don't map to a User. A request carrying a valid key is + # authorized as an anonymous service caller (user stays None — endpoints + # already handle that), rather than being 403'd into the SSO flow once + # AUTH_TYPE enforces auth (e.g. OIDC). Browser requests without a session + # still 403. Admin-only routes (current_admin_user) still require an admin + # user, so a key alone does not grant admin access. + if user is None and request_has_valid_api_key(request, db_session): + return None return await double_check_user(user) diff --git a/backend/tests/integration/test_auth_api_key_and_sso.py b/backend/tests/integration/test_auth_api_key_and_sso.py new file mode 100644 index 00000000000..7990f328265 --- /dev/null +++ b/backend/tests/integration/test_auth_api_key_and_sso.py @@ -0,0 +1,153 @@ +"""Integration tests for the auth gate — both flows must keep working: + + 1. SSO / session: a browser request with no session is rejected (403), which + is what drives the OIDC login flow. This must NOT be weakened. + 2. API key: these are service credentials for *automation* and intentionally + do NOT map to a `User`. A request carrying a valid `X-API-Key` is authorized + as an anonymous service caller (`user=None`, which endpoints already handle) + instead of being 403'd into the SSO flow. + +Regression context: enabling OIDC (AUTH_TYPE=oidc) flipped DISABLE_AUTH off, so +`current_user` started 403'ing api-key requests (they have no session and the +keys don't resolve to a user). These tests lock the contract so future changes +to auth or SSO don't silently break either flow. + +Style matches the other tests in this dir: drive the real dependency functions +directly with stubs (no live server / DB). +""" +import asyncio +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from starlette.datastructures import Headers + +from danswer.auth import api_key as api_key_mod +from danswer.auth import users as users_mod +from danswer.auth.api_key import request_has_valid_api_key +from danswer.auth.schemas import UserRole + + +def _run(coro): # avoid a hard pytest-asyncio dependency + return asyncio.run(coro) + + +def _request(headers: dict[str, str]) -> SimpleNamespace: + # Real Starlette Headers => case-insensitive lookup, exactly like a request + # (the Postman collection sends lowercase "x-api-key"). + return SimpleNamespace(headers=Headers(headers)) + + +class _StubDB: + """Stands in for the Session: `.scalar()` returns a row iff `found`.""" + + def __init__(self, found: bool) -> None: + self._found = found + + def scalar(self, *_args, **_kwargs): # type: ignore[no-untyped-def] + return SimpleNamespace(id=1, user_id="00000000-0000-0000-0000-000000000000") if self._found else None + + +def _user(*, verified: bool = True, role: UserRole = UserRole.BASIC) -> SimpleNamespace: + return SimpleNamespace(is_verified=verified, role=role, email="svc@example.com") + + +@pytest.fixture(autouse=True) +def _clear_api_key_cache(): + # validate_api_key / request_has_valid_api_key share a module-level TTLCache; + # clear it so tests don't leak validity into each other. + api_key_mod.cache.clear() + yield + api_key_mod.cache.clear() + + +@pytest.fixture +def auth_enforced(monkeypatch): + """Force the 'auth enabled' world (e.g. OIDC) deterministically. + + `double_check_user`'s `optional` default is captured from DISABLE_AUTH at + import time, so we pin it to False here regardless of the test env's + AUTH_TYPE. `current_admin_user` reads DISABLE_AUTH at call time, so patch + that too. + """ + monkeypatch.setattr(users_mod.double_check_user, "__defaults__", (False,)) + monkeypatch.setattr(users_mod, "DISABLE_AUTH", False) + + +# --------------------------------------------------------------------------- # +# request_has_valid_api_key — the validity check itself +# --------------------------------------------------------------------------- # +def test_valid_api_key_header_is_accepted(): + assert request_has_valid_api_key(_request({"X-API-Key": "k"}), _StubDB(found=True)) is True + + +def test_lowercase_header_is_accepted(): + # Postman sends "x-api-key"; header lookup must be case-insensitive. + assert request_has_valid_api_key(_request({"x-api-key": "k"}), _StubDB(found=True)) is True + + +def test_unknown_api_key_is_rejected(): + assert request_has_valid_api_key(_request({"X-API-Key": "nope"}), _StubDB(found=False)) is False + + +def test_missing_header_is_not_valid(): + assert request_has_valid_api_key(_request({}), _StubDB(found=True)) is False + + +def test_empty_header_value_is_not_valid(): + assert request_has_valid_api_key(_request({"X-API-Key": ""}), _StubDB(found=True)) is False + + +# --------------------------------------------------------------------------- # +# current_user — the gate endpoints actually depend on (auth enforced) +# --------------------------------------------------------------------------- # +def test_valid_api_key_authorizes_as_anonymous_service_caller(auth_enforced): + # The fix: valid key, no session -> authorized with user=None (no 403/SSO). + result = _run(current_user_call(_request({"x-api-key": "k"}), user=None, db=_StubDB(found=True))) + assert result is None + + +def test_invalid_api_key_is_rejected(auth_enforced): + with pytest.raises(HTTPException) as exc: + _run(current_user_call(_request({"x-api-key": "bad"}), user=None, db=_StubDB(found=False))) + assert exc.value.status_code == 403 + + +def test_no_session_and_no_api_key_is_rejected_so_sso_still_triggers(auth_enforced): + # The SSO guard: browser request, no session, no key -> 403 (drives login). + with pytest.raises(HTTPException) as exc: + _run(current_user_call(_request({}), user=None, db=_StubDB(found=False))) + assert exc.value.status_code == 403 + + +def test_session_user_still_authenticates(auth_enforced): + # A real (verified) OIDC/session user must still pass, key or not. + u = _user(verified=True) + assert _run(current_user_call(_request({}), user=u, db=_StubDB(found=False))) is u + + +# --------------------------------------------------------------------------- # +# current_admin_user — a key alone must NOT grant admin +# --------------------------------------------------------------------------- # +def test_api_key_does_not_grant_admin(auth_enforced): + # api-key request resolves to user=None -> admin gate must still 403. + with pytest.raises(HTTPException) as exc: + _run(users_mod.current_admin_user(user=None)) + assert exc.value.status_code == 403 + + +def test_admin_user_passes_admin_gate(auth_enforced): + admin = _user(role=UserRole.ADMIN) + assert _run(users_mod.current_admin_user(user=admin)) is admin + + +def test_basic_user_blocked_from_admin_gate(auth_enforced): + with pytest.raises(HTTPException) as exc: + _run(users_mod.current_admin_user(user=_user(role=UserRole.BASIC))) + assert exc.value.status_code == 403 + + +# helper: call current_user with explicit args (bypassing FastAPI DI defaults) +async def current_user_call(request, user, db): # noqa: ANN001 + return await users_mod.current_user(request=request, user=user, db_session=db) From 8f7a702a984bbf23b28a3628717571d94406eba2 Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Wed, 10 Jun 2026 15:55:19 +0530 Subject: [PATCH 013/115] k8s(prod): bump backend vha-149 (api-key auth under OIDC fix) Deployed + validated live: GET /api/persona with a valid x-api-key returns 200; without a key still 403s (SSO gate intact). Co-Authored-By: Claude Opus 4.8 (1M context) --- k8s/overlays/prod/kustomization.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/k8s/overlays/prod/kustomization.yaml b/k8s/overlays/prod/kustomization.yaml index 06767df7e85..46425b4977a 100644 --- a/k8s/overlays/prod/kustomization.yaml +++ b/k8s/overlays/prod/kustomization.yaml @@ -28,7 +28,7 @@ namespace: darwin images: - name: danswer-backend newName: sfbrdevhelmweacr.azurecr.io/danswer/danswer-backend - newTag: vha-148 + newTag: vha-149 - name: danswer-web-server newName: sfbrdevhelmweacr.azurecr.io/danswer/danswer-web-server newTag: vha-78 From 8c5f4f876c6bd8457a319ea21bcce2f5f4bbd92c Mon Sep 17 00:00:00 2001 From: rajiv chodisetti Date: Wed, 10 Jun 2026 20:06:03 +0530 Subject: [PATCH 014/115] web: app-wide dark mode (default) + chat UX overhaul + token theming - Make design-system color tokens CSS-variable-driven so they flip under a `.dark` ancestor (light values unchanged via fallbacks). Dark is now the app-wide default via a no-flash init script on ; users opt into light via the sidebar toggle (persisted app-wide to localStorage). - Typography: replace Inter with IBM Plex Sans (body) + Fraunces (display); empty-state headline uses the display font with a staggered entrance. - Chat polish: input bar elevated with an accent focus-glow (fixes the previously-broken focus ring); subtle atmosphere glow on the chat canvas; reusable da-fade-up motion (reduced-motion guarded). - Assistant picker ("Choose Assistant") gains a live search box (name + description) with an empty state. - Consistency sweep: map raw neutral colors to semantic tokens across shared components + search/admin so they flip correctly in dark. Followup judgment-call stragglers + per-page UX land in a separate commit. Co-Authored-By: Claude Opus 4.8 (1M context) --- web/src/app/admin/tools/ToolEditor.tsx | 2 +- .../assistants/gallery/AssistantsGallery.tsx | 2 +- web/src/app/chat/ChatIntro.tsx | 16 +++-- web/src/app/chat/ChatPage.tsx | 5 +- web/src/app/chat/ChatThemeToggle.tsx | 49 ++++++++++++++ web/src/app/chat/StarterMessage.tsx | 2 +- .../documentSidebar/ChatDocumentDisplay.tsx | 2 +- web/src/app/chat/input/ChatInputBar.tsx | 13 ++-- .../modal/configuration/AssistantsTab.tsx | 31 ++++++++- .../sessionSidebar/ChatSessionDisplay.tsx | 8 +-- .../app/chat/sessionSidebar/ChatSidebar.tsx | 5 ++ web/src/app/globals.css | 66 +++++++++++++++++++ web/src/app/layout.tsx | 36 ++++++---- web/src/components/EditableValue.tsx | 2 +- web/src/components/Spinner.tsx | 2 +- web/src/components/search/DocumentDisplay.tsx | 2 +- web/src/components/search/QAFeedback.tsx | 4 +- web/tailwind-themes/tailwind.config.js | 45 +++++++------ 18 files changed, 233 insertions(+), 59 deletions(-) create mode 100644 web/src/app/chat/ChatThemeToggle.tsx diff --git a/web/src/app/admin/tools/ToolEditor.tsx b/web/src/app/admin/tools/ToolEditor.tsx index 89046d21f1f..4b6b34e19a9 100644 --- a/web/src/app/admin/tools/ToolEditor.tsx +++ b/web/src/app/admin/tools/ToolEditor.tsx @@ -143,7 +143,7 @@ function ToolForm({

Available methods

- +
diff --git a/web/src/app/assistants/gallery/AssistantsGallery.tsx b/web/src/app/assistants/gallery/AssistantsGallery.tsx index cfae8b122b0..4016ab0e17c 100644 --- a/web/src/app/assistants/gallery/AssistantsGallery.tsx +++ b/web/src/app/assistants/gallery/AssistantsGallery.tsx @@ -225,7 +225,7 @@ function GalleryCard({ assistant, user, isAdded, onAdd, onRemove }: CardProps) { )} {/* Footer row: author (or built-in subtle text) + Add/Remove */} -
+
{isBuiltIn ? ( Bundled assistant diff --git a/web/src/app/chat/ChatIntro.tsx b/web/src/app/chat/ChatIntro.tsx index b4cba7d4dde..2633e9b515e 100644 --- a/web/src/app/chat/ChatIntro.tsx +++ b/web/src/app/chat/ChatIntro.tsx @@ -118,22 +118,26 @@ export function ChatIntro({
-
+
-
+
{selectedPersona?.name || "How can I help you today?"}
{selectedPersona && ( -
{selectedPersona.description}
+
+ {selectedPersona.description} +
)}
{setConfigModalActiveTab && ( - +
+ +
)} {selectedPersona && selectedPersona.num_chunks !== 0 && ( diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx index 9fbfeadd41a..f44530b6f35 100644 --- a/web/src/app/chat/ChatPage.tsx +++ b/web/src/app/chat/ChatPage.tsx @@ -1149,7 +1149,10 @@ export function ChatPage({ Only used in the EE version of the app. */} -
+
, which drives the CSS-variable token palette across + * every page. + */ +export function ChatThemeToggle() { + const [dark, setDark] = useState(true); + + useEffect(() => { + setDark(isDark()); + }, []); + + const toggle = () => { + const next = !dark; + setDark(next); + localStorage.setItem(STORAGE_KEY, next ? "dark" : "light"); + applyDark(next); + }; + + return ( + +
+ {dark ? ( + + ) : ( + + )} + {dark ? "Light mode" : "Dark mode"} +
+ + ); +} diff --git a/web/src/app/chat/StarterMessage.tsx b/web/src/app/chat/StarterMessage.tsx index f344739148f..15d67eb6a09 100644 --- a/web/src/app/chat/StarterMessage.tsx +++ b/web/src/app/chat/StarterMessage.tsx @@ -10,7 +10,7 @@ export function StarterMessage({ return (
diff --git a/web/src/app/chat/documentSidebar/ChatDocumentDisplay.tsx b/web/src/app/chat/documentSidebar/ChatDocumentDisplay.tsx index 3243c15896f..0734b1d47d2 100644 --- a/web/src/app/chat/documentSidebar/ChatDocumentDisplay.tsx +++ b/web/src/app/chat/documentSidebar/ChatDocumentDisplay.tsx @@ -57,7 +57,7 @@ export function ChatDocumentDisplay({ {isAIPick && (
} + mainContent={} popupContent={
diff --git a/web/src/app/chat/input/ChatInputBar.tsx b/web/src/app/chat/input/ChatInputBar.tsx index 460400a94dd..da10f859c8a 100644 --- a/web/src/app/chat/input/ChatInputBar.tsx +++ b/web/src/app/chat/input/ChatInputBar.tsx @@ -228,7 +228,7 @@ export function ChatInputBar({ ref={suggestionsRef} className="text-sm absolute inset-x-0 top-0 w-full transform -translate-y-full" > -
+
{filteredPersonas.map((currentPersona, index) => (
Name