From 5c722635afc40c0c438353bf1f4c69ce337ea544 Mon Sep 17 00:00:00 2001 From: charles Date: Tue, 17 Feb 2026 12:02:09 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(backend)=20add=20rerank=20param=20on?= =?UTF-8?q?=20search=20route?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit We want to be able to disable the reranker at run time even if it is enabled in settings --- src/backend/core/schemas.py | 1 + src/backend/core/services/reranking.py | 24 ++++++- src/backend/core/services/search.py | 5 +- .../core/tests/test_api_documents_search.py | 34 ++++++++++ src/backend/core/tests/test_search.py | 65 ++++++++++++------- src/backend/core/views.py | 4 ++ 6 files changed, 105 insertions(+), 28 deletions(-) diff --git a/src/backend/core/schemas.py b/src/backend/core/schemas.py index fe2e00d..d44197f 100644 --- a/src/backend/core/schemas.py +++ b/src/backend/core/schemas.py @@ -120,6 +120,7 @@ class SearchQueryParametersSchema(BaseModel): order_direction: Optional[Literal["asc", "desc"]] = Field(default="desc") nb_results: Optional[conint(ge=1, le=300)] = Field(default=50) search_type: Optional[SearchTypeEnum] = Field(default=None) + rerank: Optional[bool] = Field(default=None) class DeleteDocumentsSchema(BaseModel): diff --git a/src/backend/core/services/reranking.py b/src/backend/core/services/reranking.py index 15fb7f2..f3cab74 100644 --- a/src/backend/core/services/reranking.py +++ b/src/backend/core/services/reranking.py @@ -1,7 +1,6 @@ """Reranking utilities using rerankers library.""" import logging -from enum import Enum from functools import cache from django.conf import settings @@ -50,7 +49,9 @@ def get_reranker() -> BaseRanker | None: try: logger.info("Initializing reranker model: %s", settings.RERANKER_MODEL_NAME) return Reranker( - settings.RERANKER_MODEL_NAME, model_type=settings.RERANKER_MODEL_TYPE, api_key=settings.RERANKER_API_KEY + settings.RERANKER_MODEL_NAME, + model_type=settings.RERANKER_MODEL_TYPE, + api_key=settings.RERANKER_API_KEY, ) except Exception as e: # noqa: BLE001# pylint: disable=broad-exception-caught logger.error("Failed to initialize reranker: %s", str(e)) @@ -106,3 +107,22 @@ def prepare_rerank_data(original_hits: list[dict]) -> tuple[list[str], list[str] doc_ids.append(hit["_id"]) return docs, doc_ids + + +def should_rerank(is_rerank_requested: bool | None) -> bool: + """ + Determine whether to perform reranking based on the input parameter and settings. + logs warning if reranking was explicitly requested but the reranker is disabled in settings. + falls back to settings.RERANKER_ENABLED if is_rerank_requested is None. + """ + if is_rerank_requested and not settings.RERANKER_ENABLED: + logger.warning( + "Reranking was explicitly requested but the reranker " + "is disabled in settings. Reranking skipped." + ) + return False + + if is_rerank_requested is None: + return settings.RERANKER_ENABLED + + return is_rerank_requested and settings.RERANKER_ENABLED diff --git a/src/backend/core/services/search.py b/src/backend/core/services/search.py index d9661e9..bfa2e48 100644 --- a/src/backend/core/services/search.py +++ b/src/backend/core/services/search.py @@ -9,7 +9,7 @@ from core.enums import SearchTypeEnum from .embedding import embed_text from .opensearch import check_hybrid_search_enabled, opensearch_client -from .reranking import rerank +from .reranking import rerank, should_rerank logger = logging.getLogger(__name__) @@ -28,6 +28,7 @@ def search( # noqa : PLR0913 tags, search_type, path=None, + rerank_requested=None, ): """Perform an OpenSearch search""" query = get_query( @@ -63,7 +64,7 @@ def search( # noqa : PLR0913 ignore_unavailable=True, ) - if settings.RERANKER_ENABLED and q != "*": + if should_rerank(rerank_requested) and q != "*": response["hits"]["hits"] = rerank( query=q, hits=response["hits"]["hits"], diff --git a/src/backend/core/tests/test_api_documents_search.py b/src/backend/core/tests/test_api_documents_search.py index a79c503..6bc4def 100644 --- a/src/backend/core/tests/test_api_documents_search.py +++ b/src/backend/core/tests/test_api_documents_search.py @@ -8,6 +8,7 @@ of documents is slow and better done only once. import logging import operator import random +from unittest import mock import pytest import responses @@ -993,3 +994,36 @@ def test_api_documents_search_filtering_by_path(settings): assert len(response.json()) == 2 for hit in response.json(): assert hit["_source"]["path"].startswith(path_filter) + + +@responses.activate +@mock.patch("core.services.reranking.rerank") +def test_api_documents_rerank_blocked( + mock_rerank, + settings, + caplog, +): + """Test can disable rerank at run time even when enabled in settings""" + settings.RERANKER_ENABLED = True + + setup_oicd_resource_server(responses, settings, sub="user_sub") + service = factories.ServiceFactory() + + documents = bulk_create_documents([{"title": "Fox"}, {"title": "Wolf"}]) + prepare_index(service.index_name, documents) + mock_rerank.return_value = documents + + with caplog.at_level(logging.INFO): + response = APIClient().post( + "/api/v1.0/documents/search/", + { + "q": "programming", + "rerank": False, + "visited": [doc["id"] for doc in documents], + }, + format="json", + HTTP_AUTHORIZATION=f"Bearer {build_authorization_bearer()}", + ) + + assert response.status_code == 200 + assert not mock_rerank.called diff --git a/src/backend/core/tests/test_search.py b/src/backend/core/tests/test_search.py index 6018546..9233ef7 100644 --- a/src/backend/core/tests/test_search.py +++ b/src/backend/core/tests/test_search.py @@ -715,10 +715,40 @@ def test_search_filtering_by_query_path_and_tag(): assert returned_ids == expected_ids +# pylint: disable=too-many-arguments, too-many-positional-arguments +@responses.activate @patch("core.services.search.rerank") -def test_reranker_disabled(mocked_rerank, settings): - """Test that reranker is not called when disabled""" - settings.RERANKER_ENABLED = False +@pytest.mark.parametrize( + "rerank_requested, reranker_enabled, rerank_called, caplog_messages", + [ + (True, True, True, []), + ( + True, + False, + False, + [ + "Reranking was explicitly requested but the reranker is " + "disabled in settings. Reranking skipped." + ], + ), + (False, True, False, []), + (False, False, False, []), + (None, False, False, []), + (None, True, True, []), + ], +) +def test_search_rerank( # noqa : PLR0913 + mock_rerank, + rerank_requested, + reranker_enabled, + rerank_called, + caplog_messages, + settings, + caplog, +): + """Test that rerank activation according to args and settings""" + settings.RERANKER_ENABLED = reranker_enabled + documents = bulk_create_documents( [ {"title": "wolf", "content": "wolves live in packs and hunt together"}, @@ -726,28 +756,15 @@ def test_reranker_disabled(mocked_rerank, settings): {"title": "cat", "content": "cats are curious and independent pets"}, ] ) + mock_rerank.return_value = documents service = factories.ServiceFactory(name=SERVICE_NAME) prepare_index(service.index_name, documents) - search(q="canine pet", **search_params(service)) + with caplog.at_level(logging.INFO): + search( + q="canine pet", rerank_requested=rerank_requested, **search_params(service) + ) - mocked_rerank.assert_not_called() - - -@patch("core.services.search.rerank") -def test_reranker_enabled(mocked_rerank, settings): - """Test that reranker is called when enabled""" - settings.RERANKER_ENABLED = True - documents = bulk_create_documents( - [ - {"title": "wolf", "content": "wolves live in packs and hunt together"}, - {"title": "dog", "content": "dogs are loyal domestic animals"}, - {"title": "cat", "content": "cats are curious and independent pets"}, - ] - ) - service = factories.ServiceFactory(name=SERVICE_NAME) - prepare_index(service.index_name, documents) - - search(q="canine pet", **search_params(service)) - - mocked_rerank.assert_called_once() + assert mock_rerank.called == rerank_called + for message in caplog_messages: + assert message in caplog.text diff --git a/src/backend/core/views.py b/src/backend/core/views.py index 5a4aa26..635c1ec 100644 --- a/src/backend/core/views.py +++ b/src/backend/core/views.py @@ -354,6 +354,9 @@ class SearchDocumentView(ResourceServerMixin, views.APIView): - 'full_text': Uses only full-text search, even if hybrid search is enabled on the server. if the not specified, the server will use hybrid search when enabled + rerank : bool, optional + Enable or disable reranking of results. If not specified, falls back to + the RERANKER_ENABLED setting. Returns: -------- @@ -394,6 +397,7 @@ class SearchDocumentView(ResourceServerMixin, views.APIView): else SearchTypeEnum.HYBRID if check_hybrid_search_enabled() else SearchTypeEnum.FULL_TEXT, + rerank_requested=params.rerank, )["hits"]["hits"] logger.info("found %d results", len(result)) logger.debug("results %s", result)