diff --git a/CHANGELOG.md b/CHANGELOG.md index 557fd6d..bc2f6dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ and this project adheres to - ✨(backend) add deletion endpoint - ✨(backend) add path filter - ✨(backend) add search_type param +- ✨(backend) implement reranker #47 ## Changed diff --git a/src/backend/core/services/reranking.py b/src/backend/core/services/reranking.py new file mode 100644 index 0000000..dcef278 --- /dev/null +++ b/src/backend/core/services/reranking.py @@ -0,0 +1,73 @@ +"""Reranking utilities using rerankers library.""" + +import logging +from functools import cache +from rerankers import Reranker, Document + +from django.conf import settings + +from core.services.indexing import format_document +from core.utils import get_language_value + +logger = logging.getLogger(__name__) + + +@cache +def get_reranker(): + """Get or create the reranker instance (singleton pattern for efficiency).""" + try: + logger.info("Initializing reranker model: %s", settings.RERANKER_MODEL_NAME) + reranker = Reranker("mixedbread-ai/mxbai-rerank-large-v1", model_type="cross-encoder") + logger.info("Reranker initialized successfully") + except Exception as e: # noqa: BLE001 + logger.error("Failed to initialize reranker: %s", str(e)) + return None + + return reranker + + +def rerank(query, results): + """ + Rerank search results using the configured reranker model. + + Args: + query: The search query string + results: List of OpenSearch hit objects with _source containing title and content + + Returns: + List of reranked results in the same format as input + """ + if not settings.RERANKER_ENABLED: + logger.debug("Reranker disabled, returning original results") + return results + + reranker = get_reranker() + if reranker is None: + logger.warning("Reranker not available, returning original results") + return results + + try: + documents = [] + for hit in results: + title = get_language_value(hit["_source"], "title") + content = get_language_value(hit["_source"], "content") + documents.append(format_document(title, content)) + + logger.info("Reranking %d results for query: %s", len(results), query) + reranked = reranker.rank(query=query, docs=documents) + + reranked_results = [] + for reranked_result in reranked.results: + # doc_id is the index position in the original list + original_hit = results[reranked_result.doc_id] + # Add reranking score as metadata + original_hit["_rerank_score"] = reranked_result.score + reranked_results.append(original_hit) + + logger.info("Reranking completed, returned %d results", len(reranked_results)) + + return reranked_results + + except Exception as e: # noqa: BLE001 + logger.error("Reranking failed: %s, returning original results", str(e)) + return results diff --git a/src/backend/core/services/search.py b/src/backend/core/services/search.py index e151b0d..0278c63 100644 --- a/src/backend/core/services/search.py +++ b/src/backend/core/services/search.py @@ -9,6 +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 logger = logging.getLogger(__name__) @@ -40,7 +41,7 @@ def search( # noqa : PLR0913 path=path, search_type=search_type, ) - return opensearch_client().search( # pylint: disable=unexpected-keyword-arg + response = opensearch_client().search( # pylint: disable=unexpected-keyword-arg index=",".join(search_indices), body={ "_source": enums.SOURCE_FIELDS, # limit the fields to return @@ -61,6 +62,15 @@ def search( # noqa : PLR0913 # ignore_unavailable is not in the method declaration ignore_unavailable=True, ) + + # Apply reranking if enabled + if settings.RERANKER_ENABLED and q != "*": + response["hits"]["hits"] = rerank( + query=q, + results=response["hits"]["hits"], + ) + + return response # pylint: disable=too-many-arguments, too-many-positional-arguments diff --git a/src/backend/find/settings.py b/src/backend/find/settings.py index 70caf77..f9686b9 100755 --- a/src/backend/find/settings.py +++ b/src/backend/find/settings.py @@ -319,6 +319,24 @@ class Base(Configuration): default=1024, environ_name="EMBEDDING_DIMENSION", environ_prefix=None ) + # Reranking + RERANKER_ENABLED = values.BooleanValue( + default=False, environ_name="RERANKER_ENABLED", environ_prefix=None + ) + RERANKER_MODEL_NAME = values.Value( + default=None, + environ_name="RERANKER_MODEL_NAME", + environ_prefix=None, + ) + RERANKER_MODEL_TYPE = values.Value( + default=None, + environ_name="RERANKER_MODEL_TYPE", + environ_prefix=None, + ) + RERANKER_API_KEY = values.Value( + default=None, environ_name="RERANKER_API_KEY", environ_prefix=None + ) + # CORS CORS_ALLOW_CREDENTIALS = True CORS_ALLOW_ALL_ORIGINS = values.BooleanValue(True)