Feat/add OpenAI reranking

conf/config.yaml

@@ -68,9 +68,13 @@ rdb:
 # --- Reranker ---
 # Env: RERANKER_ENABLED, RERANKER_MODEL, RERANKER_TOP_K, RERANKER_BASE_URL, RERANKER_PORT
 reranker:
-  enable: true
+  provider: infinity
+  enabled: true
   model_name: Alibaba-NLP/gte-multilingual-reranker-base
   top_k: 10
+  api_key: EMPTY
+  timeout: 60.0
+  semaphore: 5
   base_url: ""   # Default built from RERANKER_PORT if empty
 
 # --- Map-Reduce ---

docker-compose.yaml

@@ -1,7 +1,7 @@
 include:
   - vdb/milvus.yaml
   - ${CHAINLIT_DATALAYER_COMPOSE:-extern/dummy.yaml}
-  - extern/infinity.yaml
+  - extern/reranker/${RERANKER_PROVIDER:-infinity}.yaml
   - ${TRANSCRIBER_COMPOSE:-extern/dummy.yaml}
 
 x-openrag: &openrag_template

docs/content/docs/documentation/env_vars.md

@@ -229,19 +229,24 @@ The retriever fetches relevant documents from the vector database based on query
 
 ### Reranker Configuration
 
-The reranker enhances search quality by re-scoring and reordering retrieved documents according to their relevance to the user's query. Currently, the system uses [Infinity server](https://github.com/michaelfeil/infinity) for reranking functionality.
-
-:::info[Future Improvements]
-The current Infinity server interface is not OpenAI-compatible, which limits integration flexibility. We plan to improve this by supporting OpenAI-compatible reranker interfaces in future releases.
-:::
+The reranker enhances search quality by re-scoring and reordering retrieved documents according to their relevance to the user's query. Two providers are supported: **Infinity** (default) and **OpenAI-compatible** endpoints.
 
 | Variable | Type | Default | Description |
 |----------|------|---------|-------------|
 | `RERANKER_ENABLED` | `bool` | true | Enable or disable the reranking mechanism |
+| `RERANKER_PROVIDER` | `str` | `infinity` | Reranker backend to use. Accepted values: `infinity`, `openai` |
 | `RERANKER_MODEL` | `str` | Alibaba-NLP/gte-multilingual-reranker-base | Model used for reranking documents.|
-| `RERANKER_TOP_K` | `int` | 5 | Number of top documents to return after reranking. Increase to 8 for better results if your LLM has a wider context window |
+| `RERANKER_TOP_K` | `int` | 10 | Number of top documents to return after reranking. Increase for better results if your LLM has a wider context window |
 | `RERANKER_BASE_URL` | `str` | `http://reranker:7997` | Base URL of the reranker service |
-| `RERANKER_PORT` | `int` | 7997 | Port on which the reranker service listens |
+| `RERANKER_API_KEY` | `str` | `EMPTY` | API key for the reranker service. Required when using the `openai` provider |
+| `RERANKER_SEMAPHORE` | `int` | 40 | Maximum number of concurrent reranking requests. Adjust based on your server capacity |
+
+#### Reranker Providers
+
+| Provider | `RERANKER_PROVIDER` value | Description |
+|----------|--------------------------|-------------|
+| **Infinity** | `infinity` | Uses the [Infinity server](https://github.com/michaelfeil/infinity) via its native client. Default port: `7997` |
+| **OpenAI-compatible** | `openai` | Uses any OpenAI-compatible reranker endpoint (e.g. vLLM, LiteLLM, TEI). Default port: `8000` |
 
 ## Extra
 ### Prompts

extern/infinity.yaml → extern/reranker/infinity.yaml

@@ -6,7 +6,7 @@ x-reranker: &reranker_template
   volumes:
     - ${VLLM_CACHE:-/root/.cache/huggingface}:/app/.cache/huggingface # Model weights for RAG
   # ports:
-  #   - ${RERANKER_PORT:-7997}:${RERANKER_PORT:-7997}
+  #   - ${RERANKER_PORT:-7997}:7997
 
 services:
   reranker:
@@ -23,7 +23,8 @@ services:
     command: >
       v2
       --model-id ${RERANKER_MODEL:-Alibaba-NLP/gte-multilingual-reranker-base}
-      --port ${RERANKER_PORT:-7997}
+      --api-key ${RERANKER_API_KEY:-"EMPTY"}
+      --port 7997
     profiles:
       - ''
 
@@ -35,7 +36,8 @@ services:
       v2
       --engine torch
       --model-id ${RERANKER_MODEL:-Alibaba-NLP/gte-multilingual-reranker-base}
-      --port ${RERANKER_PORT:-7997}
+      --api-key ${RERANKER_API_KEY:-"EMPTY"}
+      --port 7997
     profiles:
       - 'cpu'
 

extern/reranker/openai.yaml

@@ -0,0 +1,63 @@
+x-vllm-env: &vllm_env
+  HUGGING_FACE_HUB_TOKEN:
+  VLLM_SLEEP_WHEN_IDLE: 1  # Avoid 100% CPU usage when idle
+
+x-reranker: &reranker_template
+  networks:
+    default:
+      aliases:
+        - reranker
+  # restart: on-failure
+  environment:
+    - HUGGING_FACE_HUB_TOKEN
+    - VLLM_SLEEP_WHEN_IDLE=1  # Avoid 100% CPU usage when idle
+  ipc: "host"
+  volumes:
+    - ${VLLM_CACHE:-/root/.cache/huggingface}:/root/.cache/huggingface
+  command: >
+    --model ${RERANKER_MODEL:-BAAI/bge-reranker-v2-m3}
+    --trust-remote-code
+    --api-key ${RERANKER_API_KEY:-"EMPTY"}
+    --gpu_memory_utilization 0.3
+  healthcheck:
+    test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+    interval: 20s
+    timeout: 5s
+    retries: 4
+    start_period: 90s
+  # ports:
+  #     - ${RERANKER_PORT:-8000}:8000
+
+services:
+  reranker-gpu:
+    <<: *reranker_template
+    image: vllm/vllm-openai:v0.17.1
+    environment:
+      <<: *vllm_env
+      NVIDIA_VISIBLE_DEVICES: all
+      NVIDIA_DRIVER_CAPABILITIES: compute,utility
+    runtime: nvidia
+    profiles:
+      - ""
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+
+  reranker-cpu:
+    <<: *reranker_template
+    image: vllm/vllm-openai-cpu:v0.17.1
+    deploy: {}
+    environment:
+      <<: *vllm_env
+      VLLM_CPU_KVCACHE_SPACE: 8
+    command: >
+      --model ${RERANKER_MODEL:-BAAI/bge-reranker-v2-m3}
+      --trust-remote-code
+      --api-key ${RERANKER_API_KEY:-"EMPTY"}
+      --dtype float32
+    profiles:
+      - "cpu"

openrag/app_front.py

@@ -125,6 +125,7 @@ async def chat_profile(current_user: cl.User):
                     name=m.id,
                     markdown_description=description_template.format(name=m.id, partition=partition),
                     icon="/public/favicon.svg",
+                    default=m.id == "openrag-all",
                 )
             )
         return chat_profiles

openrag/components/pipeline.py

@@ -20,7 +20,7 @@
 
 from .llm import LLM
 from .map_reduce import RAGMapReduce
-from .reranker import Reranker
+from .reranker import BaseReranker, RerankerFactory
 from .retriever import BaseRetriever, RetrieverFactory
 from .utils import SOURCE_SEPARATOR
 
@@ -49,9 +49,9 @@ def __init__(self) -> None:
         self.retriever: BaseRetriever = RetrieverFactory.create_retriever(config=config)
 
         # reranker
-        self.reranker_enabled = config.reranker.enable
-        self.reranker = Reranker(logger, config)
-        logger.debug("Reranker", enabled=self.reranker_enabled)
+        self.reranker_enabled = config.reranker.enabled
+        self.reranker: BaseReranker = RerankerFactory.get_reranker(config)
+        logger.debug("Reranker", enabled=self.reranker_enabled, provider=config.reranker.provider)
         self.reranker_top_k = config.reranker.top_k
 
     async def retrieve_docs(

openrag/components/reranker/__init__.py

@@ -0,0 +1,20 @@
+from .base import BaseReranker
+from .infinity import InfinityReranker
+from .openai import OpenAIReranker
+
+RERANKER_MAPPING = {
+    "infinity": InfinityReranker,
+    "openai": OpenAIReranker,
+}
+
+
+class RerankerFactory:
+    @staticmethod
+    def get_reranker(config: dict) -> BaseReranker:
+        provider = config.reranker.get("provider")
+        reranker_class = RERANKER_MAPPING.get(provider, None)
+
+        if not reranker_class:
+            raise ValueError(f"Unsupported reranker provider: {provider}")
+
+        return reranker_class(config)

openrag/components/reranker/base.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+
+from langchain_core.documents.base import Document
+
+
+class BaseReranker(ABC):
+    @abstractmethod
+    async def rerank(self, query: str, documents: list[Document], top_k: int | None = None) -> list[Document]:
+        """Rerank a list of documents based on a query and an optional top_k parameter"""
+
+    @staticmethod
+    def rrf_reranking(doc_lists: list[list[Document]], k: int = 60) -> list[Document]:
+        """Reciprocal_rank_fusion that takes multiple lists of ranked documents
+        and an optional parameter k used in the RRF formula
+        RRF formula: \\sum_{i=1}^{n} \frac{1}{k + rank_i}
+         where rank_i is the rank of the document in the i-th list and n is the number of lists.
+
+        k small: High sensitivity to top ranks
+        k large: More balanced sensitivity across ranks
+        k = 60 a common and balanced choice in practice.
+        """
+
+        if len(doc_lists) == 1:
+            return doc_lists[0]
+
+        # Initialize a dictionary to hold fused scores for each unique document
+        fused_scores = {}
+
+        for doc_list in doc_lists:
+            doc_list: list[Document]
+            for rank, doc in enumerate(doc_list, start=1):
+                doc_id = doc.metadata.get("_id")
+
+                score, d = fused_scores.get(doc_id, (0, doc))
+                fused_scores[doc_id] = (score + 1 / (rank + k), d)
+
+        # sort the docs
+        reranked_docs = [doc for _, doc in sorted(fused_scores.values(), key=lambda x: x[0], reverse=True)]
+        return reranked_docs

openrag/components/reranker/infinity.py

@@ -0,0 +1,55 @@
+import asyncio
+
+from infinity_client import Client
+from infinity_client.api.default import rerank
+from infinity_client.models import RerankInput, ReRankResult
+from langchain_core.documents.base import Document
+from utils.logger import get_logger
+
+from .base import BaseReranker
+
+logger = get_logger()
+
+
+class InfinityReranker(BaseReranker):
+    def __init__(self, config):
+        self.model_name = config.reranker.model_name
+        self.client = Client(
+            base_url=config.reranker.base_url,
+            timeout=config.reranker.timeout,
+            headers={"Authorization": f"Bearer {config.reranker.api_key}"},
+        )
+        self.semaphore = asyncio.Semaphore(config.reranker.semaphore)
+        logger.debug("Reranker initialized", model_name=self.model_name)
+
+    async def rerank(self, query: str, documents: list[Document], top_k: int | None = None) -> list[Document]:
+        async with self.semaphore:
+            logger.debug("Reranking documents", documents_count=len(documents), top_k=top_k)
+            top_k = min(top_k, len(documents)) if top_k is not None else len(documents)
+            rerank_input = RerankInput.from_dict(
+                {
+                    "model": self.model_name,
+                    "query": query,
+                    "documents": [doc.page_content for doc in documents],
+                    "top_n": top_k,
+                    "return_documents": True,
+                    "raw_scores": True,  # Normalized score between 0 and 1
+                }
+            )
+            try:
+                rerank_result: ReRankResult = await rerank.asyncio(client=self.client, body=rerank_input)
+                output = []
+                for rerank_res in rerank_result.results:
+                    doc = documents[rerank_res.index]
+                    doc.metadata["relevance_score"] = rerank_res.relevance_score
+                    output.append(doc)
+                return output
+
+            except Exception as e:
+                logger.error(
+                    "Reranking failed",
+                    error=str(e),
+                    model_name=self.model_name,
+                    documents_count=len(documents),
+                )
+                return documents[:top_k]

openrag/components/reranker/openai.py

@@ -0,0 +1,57 @@
+import asyncio
+
+import httpx
+from langchain_core.documents.base import Document
+from utils.logger import get_logger
+
+from .base import BaseReranker
+
+logger = get_logger()
+
+
+class OpenAIReranker(BaseReranker):
+    def __init__(self, config):
+        self.model_name = config.reranker.model_name
+        self.api_key = config.reranker.api_key
+        base_url = config.reranker.base_url.rstrip("/")
+        self.rerank_url = f"{base_url}/rerank"
+        semaphore = config.reranker.semaphore
+        self.semaphore = asyncio.Semaphore(semaphore)
+        self.timeout = config.reranker.timeout
+        logger.debug("OpenAI Reranker initialized", model_name=self.model_name)
+
+    async def rerank(self, query: str, documents: list[Document], top_k: int | None = None) -> list[Document]:
+        async with self.semaphore:
+            logger.debug("Reranking documents", documents_count=len(documents), top_k=top_k)
+            top_k = min(top_k, len(documents)) if top_k is not None else len(documents)
+            try:
+                async with httpx.AsyncClient() as client:
+                    response = await client.post(
+                        self.rerank_url,
+                        headers={"Authorization": f"Bearer {self.api_key}"},
+                        json={
+                            "model": self.model_name,
+                            "query": query,
+                            "documents": [doc.page_content for doc in documents],
+                            "top_n": top_k,
+                        },
+                        timeout=self.timeout,
+                    )
+                    response.raise_for_status()
+                    data = response.json()
+
+                output = []
+                for result in data["results"]:
+                    doc = documents[result["index"]]
+                    doc.metadata["relevance_score"] = result["relevance_score"]
+                    output.append(doc)
+                return output
+
+            except Exception as e:
+                logger.error(
+                    "Reranking failed",
+                    error=str(e),
+                    model_name=self.model_name,
+                    documents_count=len(documents),
+                )
+                return documents[:top_k]

openrag/components/test_rrf_reranking.py → openrag/components/reranker/test_rrf_reranking.py

@@ -1,8 +1,9 @@
 """Tests for BaseReranker.rrf_reranking static method."""
 
-from components.reranker import BaseReranker
 from langchain_core.documents.base import Document
 
+from .base import BaseReranker
+
 
 def make_doc(doc_id: str, content: str = "", **metadata) -> Document:
     return Document(page_content=content, metadata={"_id": doc_id, **metadata})

openrag/config/loader.py

@@ -53,10 +53,14 @@
     ("POSTGRES_PASSWORD", "rdb.password", str),
     ("DEFAULT_FILE_QUOTA", "rdb.default_file_quota", int),
     # Reranker
-    ("RERANKER_ENABLED", "reranker.enable", bool),
+    ("RERANKER_PROVIDER", "reranker.provider", str),
+    ("RERANKER_ENABLED", "reranker.enabled", bool),
     ("RERANKER_MODEL", "reranker.model_name", str),
     ("RERANKER_TOP_K", "reranker.top_k", int),
     ("RERANKER_BASE_URL", "reranker.base_url", str),
+    ("RERANKER_API_KEY", "reranker.api_key", str),
+    ("RERANKER_TIMEOUT", "reranker.timeout", float),
+    ("RERANKER_SEMAPHORE", "reranker.semaphore", int),
     # Map-Reduce
     ("MAP_REDUCE_INITIAL_BATCH_SIZE", "map_reduce.initial_batch_size", int),
     ("MAP_REDUCE_EXPANSION_BATCH_SIZE", "map_reduce.expansion_batch_size", int),