feat: hybrid search

Dockerfile

@@ -8,8 +8,9 @@ ARG USE_CUDA_VER=cu121
 # any sentence transformer model; models to use can be found at https://huggingface.co/models?library=sentence-transformers
 # Leaderboard: https://huggingface.co/spaces/mteb/leaderboard 
 # for better performance and multilangauge support use "intfloat/multilingual-e5-large" (~2.5GB) or "intfloat/multilingual-e5-base" (~1.5GB)
-# IMPORTANT: If you change the default model (sentence-transformers/all-MiniLM-L6-v2) and vice versa, you aren't able to use RAG Chat with your previous documents loaded in the WebUI! You need to re-embed them.
+# IMPORTANT: If you change the embedding model (sentence-transformers/all-MiniLM-L6-v2) and vice versa, you aren't able to use RAG Chat with your previous documents loaded in the WebUI! You need to re-embed them.
 ARG USE_EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2
+ARG USE_RERANKING_MODEL=BAAI/bge-reranker-base
 
 ######## WebUI frontend ########
 FROM --platform=$BUILDPLATFORM node:21-alpine3.19 as build
@@ -30,6 +31,7 @@ ARG USE_CUDA
 ARG USE_OLLAMA
 ARG USE_CUDA_VER
 ARG USE_EMBEDDING_MODEL
+ARG USE_RERANKING_MODEL
 
 ## Basis ##
 ENV ENV=prod \
@@ -38,7 +40,8 @@ ENV ENV=prod \
     USE_OLLAMA_DOCKER=${USE_OLLAMA} \
     USE_CUDA_DOCKER=${USE_CUDA} \
     USE_CUDA_DOCKER_VER=${USE_CUDA_VER} \
-    USE_EMBEDDING_MODEL_DOCKER=${USE_EMBEDDING_MODEL}
+    USE_EMBEDDING_MODEL_DOCKER=${USE_EMBEDDING_MODEL} \
+    USE_RERANKING_MODEL_DOCKER=${USE_RERANKING_MODEL}
 
 ## Basis URL Config ##
 ENV OLLAMA_BASE_URL="/ollama" \
@@ -62,7 +65,7 @@ ENV WHISPER_MODEL="base" \
 
 ## RAG Embedding model settings ##
 ENV RAG_EMBEDDING_MODEL="$USE_EMBEDDING_MODEL_DOCKER" \
-    RAG_EMBEDDING_MODEL_DIR="/app/backend/data/cache/embedding/models" \
+    RAG_RERANKING_MODEL="$USE_RERANKING_MODEL_DOCKER" \
     SENTENCE_TRANSFORMERS_HOME="/app/backend/data/cache/embedding/models"
 #### Other models ##########################################################
 
@@ -99,11 +102,13 @@ RUN pip3 install uv && \
         pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/$USE_CUDA_DOCKER_VER --no-cache-dir && \
         uv pip install --system -r requirements.txt --no-cache-dir && \
         python -c "import os; from sentence_transformers import SentenceTransformer; SentenceTransformer(os.environ['RAG_EMBEDDING_MODEL'], device='cpu')" && \
+        python -c "import os; from sentence_transformers import CrossEncoder; CrossEncoder(os.environ['RAG_RERANKING_MODEL'], device='cpu')" && \
         python -c "import os; from faster_whisper import WhisperModel; WhisperModel(os.environ['WHISPER_MODEL'], device='cpu', compute_type='int8', download_root=os.environ['WHISPER_MODEL_DIR'])"; \
     else \
         pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --no-cache-dir && \
         uv pip install --system -r requirements.txt --no-cache-dir && \
         python -c "import os; from sentence_transformers import SentenceTransformer; SentenceTransformer(os.environ['RAG_EMBEDDING_MODEL'], device='cpu')" && \
+        python -c "import os; from sentence_transformers import CrossEncoder; CrossEncoder(os.environ['RAG_RERANKING_MODEL'], device='cpu')" && \
         python -c "import os; from faster_whisper import WhisperModel; WhisperModel(os.environ['WHISPER_MODEL'], device='cpu', compute_type='int8', download_root=os.environ['WHISPER_MODEL_DIR'])"; \
     fi
 

backend/apps/rag/main.py

@@ -49,8 +49,8 @@ from apps.web.models.documents import (
 
 from apps.rag.utils import (
     query_embeddings_doc,
+    query_embeddings_function,
     query_embeddings_collection,
-    generate_openai_embeddings,
 )
 
 from utils.misc import (
@@ -67,6 +67,8 @@ from config import (
     RAG_EMBEDDING_ENGINE,
     RAG_EMBEDDING_MODEL,
     RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
+    RAG_RERANKING_MODEL,
+    RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
     RAG_OPENAI_API_BASE_URL,
     RAG_OPENAI_API_KEY,
     DEVICE_TYPE,
@@ -91,6 +93,7 @@ app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
 
 app.state.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
 app.state.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
+app.state.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
 app.state.RAG_TEMPLATE = RAG_TEMPLATE
 
 app.state.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
@@ -105,6 +108,12 @@ if app.state.RAG_EMBEDDING_ENGINE == "":
         trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
     )
 
+app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
+    app.state.RAG_RERANKING_MODEL,
+    device=DEVICE_TYPE,
+    trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
+)
+
 
 origins = ["*"]
 
@@ -134,6 +143,7 @@ async def get_status():
         "template": app.state.RAG_TEMPLATE,
         "embedding_engine": app.state.RAG_EMBEDDING_ENGINE,
         "embedding_model": app.state.RAG_EMBEDDING_MODEL,
+        "reranking_model": app.state.RAG_RERANKING_MODEL,
     }
 
 
@@ -150,6 +160,11 @@ async def get_embedding_config(user=Depends(get_admin_user)):
     }
 
 
+@app.get("/reranking")
+async def get_reraanking_config(user=Depends(get_admin_user)):
+    return {"status": True, "reranking_model": app.state.RAG_RERANKING_MODEL}
+
+
 class OpenAIConfigForm(BaseModel):
     url: str
     key: str
@@ -205,6 +220,36 @@ async def update_embedding_config(
         )
 
 
+class RerankingModelUpdateForm(BaseModel):
+    reranking_model: str
+    
+
+@app.post("/reranking/update")
+async def update_reranking_config(
+    form_data: RerankingModelUpdateForm, user=Depends(get_admin_user)
+):
+    log.info(
+        f"Updating reranking model: {app.state.RAG_RERANKING_MODEL} to {form_data.reranking_model}"
+    )
+    try:
+        app.state.RAG_RERANKING_MODEL = form_data.reranking_model
+        app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
+            app.state.RAG_RERANKING_MODEL,
+            device=DEVICE_TYPE,
+        )
+
+        return {
+            "status": True,
+            "reranking_model": app.state.RAG_RERANKING_MODEL,
+        }
+    except Exception as e:
+        log.exception(f"Problem updating reranking model: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=ERROR_MESSAGES.DEFAULT(e),
+        )
+
+
 @app.get("/config")
 async def get_rag_config(user=Depends(get_admin_user)):
     return {
@@ -286,34 +331,21 @@ def query_doc_handler(
     user=Depends(get_current_user),
 ):
     try:
-        if app.state.RAG_EMBEDDING_ENGINE == "":
-            query_embeddings = app.state.sentence_transformer_ef.encode(
-                form_data.query
-            ).tolist()
-        elif app.state.RAG_EMBEDDING_ENGINE == "ollama":
-            query_embeddings = generate_ollama_embeddings(
-                GenerateEmbeddingsForm(
-                    **{
-                        "model": app.state.RAG_EMBEDDING_MODEL,
-                        "prompt": form_data.query,
-                    }
-                )
-            )
-        elif app.state.RAG_EMBEDDING_ENGINE == "openai":
-            query_embeddings = generate_openai_embeddings(
-                model=app.state.RAG_EMBEDDING_MODEL,
-                text=form_data.query,
-                key=app.state.OPENAI_API_KEY,
-                url=app.state.OPENAI_API_BASE_URL,
-            )
+        embeddings_function = query_embeddings_function(
+            app.state.RAG_EMBEDDING_ENGINE,
+            app.state.RAG_EMBEDDING_MODEL,
+            app.state.sentence_transformer_ef,
+            app.state.OPENAI_API_KEY,
+            app.state.OPENAI_API_BASE_URL,
+        )
 
         return query_embeddings_doc(
             collection_name=form_data.collection_name,
             query=form_data.query,
-            query_embeddings=query_embeddings,
             k=form_data.k if form_data.k else app.state.TOP_K,
+            embeddings_function=embeddings_function,
+            reranking_function=app.state.sentence_transformer_rf,
         )
-
     except Exception as e:
         log.exception(e)
         raise HTTPException(
@@ -334,33 +366,21 @@ def query_collection_handler(
     user=Depends(get_current_user),
 ):
     try:
-        if app.state.RAG_EMBEDDING_ENGINE == "":
-            query_embeddings = app.state.sentence_transformer_ef.encode(
-                form_data.query
-            ).tolist()
-        elif app.state.RAG_EMBEDDING_ENGINE == "ollama":
-            query_embeddings = generate_ollama_embeddings(
-                GenerateEmbeddingsForm(
-                    **{
-                        "model": app.state.RAG_EMBEDDING_MODEL,
-                        "prompt": form_data.query,
-                    }
-                )
-            )
-        elif app.state.RAG_EMBEDDING_ENGINE == "openai":
-            query_embeddings = generate_openai_embeddings(
-                model=app.state.RAG_EMBEDDING_MODEL,
-                text=form_data.query,
-                key=app.state.OPENAI_API_KEY,
-                url=app.state.OPENAI_API_BASE_URL,
-            )
+        embeddings_function = embeddings_function(
+            app.state.RAG_EMBEDDING_ENGINE,
+            app.state.RAG_EMBEDDING_MODEL,
+            app.state.sentence_transformer_ef,
+            app.state.OPENAI_API_KEY,
+            app.state.OPENAI_API_BASE_URL,
+        )
 
         return query_embeddings_collection(
             collection_names=form_data.collection_names,
-            query_embeddings=query_embeddings,
+            query=form_data.query,
             k=form_data.k if form_data.k else app.state.TOP_K,
+            embeddings_function=embeddings_function,
+            reranking_function=app.state.sentence_transformer_rf,
         )
-
     except Exception as e:
         log.exception(e)
         raise HTTPException(
@@ -427,8 +447,6 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
     log.info(f"store_docs_in_vector_db {docs} {collection_name}")
 
     texts = [doc.page_content for doc in docs]
-    texts = list(map(lambda x: x.replace("\n", " "), texts))
-
     metadatas = [doc.metadata for doc in docs]
 
     try:
@@ -440,26 +458,20 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
 
         collection = CHROMA_CLIENT.create_collection(name=collection_name)
 
+        embedding_func = query_embeddings_function(
+            app.state.RAG_EMBEDDING_ENGINE,
+            app.state.RAG_EMBEDDING_MODEL,
+            app.state.sentence_transformer_ef,
+            app.state.OPENAI_API_KEY,
+            app.state.OPENAI_API_BASE_URL,
+        )
+
+        embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))
         if app.state.RAG_EMBEDDING_ENGINE == "":
-            embeddings = app.state.sentence_transformer_ef.encode(texts).tolist()
-        elif app.state.RAG_EMBEDDING_ENGINE == "ollama":
-            embeddings = [
-                generate_ollama_embeddings(
-                    GenerateEmbeddingsForm(
-                        **{"model": app.state.RAG_EMBEDDING_MODEL, "prompt": text}
-                    )
-                )
-                for text in texts
-            ]
-        elif app.state.RAG_EMBEDDING_ENGINE == "openai":
+            embeddings = embedding_func(embedding_texts)
+        else:
             embeddings = [
-                generate_openai_embeddings(
-                    model=app.state.RAG_EMBEDDING_MODEL,
-                    text=text,
-                    key=app.state.OPENAI_API_KEY,
-                    url=app.state.OPENAI_API_BASE_URL,
-                )
-                for text in texts
+                embedding_func(embedding_texts) for text in texts
             ]
 
         for batch in create_batches(

backend/apps/rag/utils.py

@@ -1,5 +1,8 @@
 import logging
 import requests
+import operator
+
+import sentence_transformers
 
 from typing import List
 
@@ -8,6 +11,11 @@ from apps.ollama.main import (
     GenerateEmbeddingsForm,
 )
 
+from langchain.retrievers import (
+    BM25Retriever,
+    EnsembleRetriever,
+)
+
 from config import SRC_LOG_LEVELS, CHROMA_CLIENT
 
 
@@ -15,60 +23,96 @@ log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
 
 
-def query_embeddings_doc(collection_name: str, query: str, query_embeddings, k: int):
+def query_embeddings_doc(
+    collection_name: str,
+    query: str,
+    k: int,
+    embeddings_function,
+    reranking_function,
+):
     try:
         # if you use docker use the model from the environment variable
-        log.info(f"query_embeddings_doc {query_embeddings}")
         collection = CHROMA_CLIENT.get_collection(name=collection_name)
 
-        result = collection.query(
-            query_embeddings=[query_embeddings],
-            n_results=k,
+        # keyword search
+        documents = collection.get() # get all documents
+        bm25_retriever = BM25Retriever.from_texts(
+            texts=documents.get("documents"),
+            metadatas=documents.get("metadatas"),
+        )
+        bm25_retriever.k = k
+
+        # semantic search (vector)
+        chroma_retriever = ChromaRetriever(
+            collection=collection,
+            k=k,
+            embeddings_function=embeddings_function,
+        )
+
+        # hybrid search (ensemble)
+        ensemble_retriever = EnsembleRetriever(
+            retrievers=[bm25_retriever, chroma_retriever],
+            weights=[0.6, 0.4]
         )
 
-        log.info(f"query_embeddings_doc:result {result}")
+        documents = ensemble_retriever.invoke(query)
+        result = query_results_rank(
+            query=query,
+            documents=documents,
+            k=k,
+            reranking_function=reranking_function,
+        )
+        result = {
+            "distances": [[d[1].item() for d in result]],
+            "documents": [[d[0].page_content for d in result]],
+            "metadatas": [[d[0].metadata for d in result]],
+        }
+
         return result
     except Exception as e:
         raise e
 
 
+def query_results_rank(query: str, documents, k: int, reranking_function):
+    scores = reranking_function.predict([(query, doc.page_content) for doc in documents])
+    docs_with_scores = list(zip(documents, scores))
+    result = sorted(docs_with_scores, key=operator.itemgetter(1), reverse=True)
+    return result[: k]
+
+
 def merge_and_sort_query_results(query_results, k):
     # Initialize lists to store combined data
-    combined_ids = []
     combined_distances = []
-    combined_metadatas = []
     combined_documents = []
+    combined_metadatas = []
 
     # Combine data from each dictionary
     for data in query_results:
-        combined_ids.extend(data["ids"][0])
         combined_distances.extend(data["distances"][0])
-        combined_metadatas.extend(data["metadatas"][0])
         combined_documents.extend(data["documents"][0])
+        combined_metadatas.extend(data["metadatas"][0])
 
-    # Create a list of tuples (distance, id, metadata, document)
+    # Create a list of tuples (distance, document, metadata)
     combined = list(
-        zip(combined_distances, combined_ids, combined_metadatas, combined_documents)
+        zip(combined_distances, combined_documents, combined_metadatas)
     )
 
     # Sort the list based on distances
     combined.sort(key=lambda x: x[0])
 
     # Unzip the sorted list
-    sorted_distances, sorted_ids, sorted_metadatas, sorted_documents = zip(*combined)
+    sorted_distances, sorted_documents, sorted_metadatas = zip(*combined)
 
     # Slicing the lists to include only k elements
     sorted_distances = list(sorted_distances)[:k]
-    sorted_ids = list(sorted_ids)[:k]
-    sorted_metadatas = list(sorted_metadatas)[:k]
     sorted_documents = list(sorted_documents)[:k]
+    sorted_metadatas = list(sorted_metadatas)[:k]
 
     # Create the output dictionary
     merged_query_results = {
-        "ids": [sorted_ids],
         "distances": [sorted_distances],
-        "metadatas": [sorted_metadatas],
         "documents": [sorted_documents],
+        "metadatas": [sorted_metadatas],
         "embeddings": None,
         "uris": None,
         "data": None,
@@ -78,19 +122,23 @@ def merge_and_sort_query_results(query_results, k):
 
 
 def query_embeddings_collection(
-    collection_names: List[str], query: str, query_embeddings, k: int
+    collection_names: List[str],
+    query: str,
+    k: int,
+    embeddings_function,
+    reranking_function,
 ):
 
     results = []
-    log.info(f"query_embeddings_collection {query_embeddings}")
 
     for collection_name in collection_names:
         try:
             result = query_embeddings_doc(
                 collection_name=collection_name,
                 query=query,
-                query_embeddings=query_embeddings,
                 k=k,
+                embeddings_function=embeddings_function,
+                reranking_function=reranking_function,
             )
             results.append(result)
         except:
@@ -105,6 +153,33 @@ def rag_template(template: str, context: str, query: str):
     return template
 
 
+def query_embeddings_function(
+    embedding_engine,
+    embedding_model,
+    embedding_function,
+    openai_key,
+    openai_url,
+):
+    if embedding_engine == "":
+        return lambda query: embedding_function.encode(query).tolist()
+    elif embedding_engine == "ollama":
+        return lambda query: generate_ollama_embeddings(
+            GenerateEmbeddingsForm(
+                **{
+                    "model": embedding_model,
+                    "prompt": query,
+                }
+            )
+        )
+    elif embedding_engine == "openai":
+        return lambda query: generate_openai_embeddings(
+            model=embedding_model,
+            text=query,
+            key=openai_key,
+            url=openai_url,
+        )
+
+
 def rag_messages(
     docs,
     messages,
@@ -113,11 +188,12 @@ def rag_messages(
     embedding_engine,
     embedding_model,
     embedding_function,
+    reranking_function,
     openai_key,
     openai_url,
 ):
     log.debug(
-        f"docs: {docs} {messages} {embedding_engine} {embedding_model} {embedding_function} {openai_key} {openai_url}"
+        f"docs: {docs} {messages} {embedding_engine} {embedding_model} {embedding_function} {reranking_function} {openai_key} {openai_url}"
     )
 
     last_user_message_idx = None
@@ -155,38 +231,29 @@ def rag_messages(
             if doc["type"] == "text":
                 context = doc["content"]
             else:
-                if embedding_engine == "":
-                    query_embeddings = embedding_function.encode(query).tolist()
-                elif embedding_engine == "ollama":
-                    query_embeddings = generate_ollama_embeddings(
-                        GenerateEmbeddingsForm(
-                            **{
-                                "model": embedding_model,
-                                "prompt": query,
-                            }
-                        )
-                    )
-                elif embedding_engine == "openai":
-                    query_embeddings = generate_openai_embeddings(
-                        model=embedding_model,
-                        text=query,
-                        key=openai_key,
-                        url=openai_url,
-                    )
+                embeddings_function = query_embeddings_function(
+                    embedding_engine,
+                    embedding_model,
+                    embedding_function,
+                    openai_key,
+                    openai_url,
+                )
 
                 if doc["type"] == "collection":
                     context = query_embeddings_collection(
                         collection_names=doc["collection_names"],
                         query=query,
-                        query_embeddings=query_embeddings,
                         k=k,
+                        embeddings_function=embeddings_function,
+                        reranking_function=reranking_function,
                     )
                 else:
                     context = query_embeddings_doc(
                         collection_name=doc["collection_name"],
                         query=query,
-                        query_embeddings=query_embeddings,
                         k=k,
+                        embeddings_function=embeddings_function,
+                        reranking_function=reranking_function,
                     )
 
         except Exception as e:
@@ -250,3 +317,41 @@ def generate_openai_embeddings(
     except Exception as e:
         print(e)
         return None
+
+
+from typing import Any
+
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
+from langchain_core.documents import Document
+from langchain_core.retrievers import BaseRetriever
+
+
+class ChromaRetriever(BaseRetriever):
+    collection: Any
+    k: int
+    embeddings_function: Any
+
+    def _get_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+    ) -> List[Document]:
+        query_embeddings = self.embeddings_function(query)
+
+        results = self.collection.query(
+            query_embeddings=[query_embeddings],
+            n_results=self.k,
+        )
+
+        ids = results["ids"][0]
+        metadatas = results["metadatas"][0]
+        documents = results["documents"][0]
+
+        return [
+            Document(
+                metadata=metadatas[idx],
+                page_content=documents[idx],
+            )
+            for idx in range(len(ids))
+        ]

backend/config.py

@@ -424,6 +424,15 @@ RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE = (
     os.environ.get("RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE", "").lower() == "true"
 )
 
+RAG_RERANKING_MODEL = os.environ.get(
+    "RAG_RERANKING_MODEL", "BAAI/bge-reranker-v2-m3"
+)
+log.info(f"Reranking model set: {RAG_RERANKING_MODEL}"),
+
+RAG_RERANKING_MODEL_TRUST_REMOTE_CODE = (
+    os.environ.get("RAG_RERANKING_MODEL_TRUST_REMOTE_CODE", "").lower() == "true"
+)
+
 # device type embedding models - "cpu" (default), "cuda" (nvidia gpu required) or "mps" (apple silicon) - choosing this right can lead to better performance
 USE_CUDA = os.environ.get("USE_CUDA_DOCKER", "false")
 

backend/main.py

@@ -117,6 +117,7 @@ class RAGMiddleware(BaseHTTPMiddleware):
                     rag_app.state.RAG_EMBEDDING_ENGINE,
                     rag_app.state.RAG_EMBEDDING_MODEL,
                     rag_app.state.sentence_transformer_ef,
+                    rag_app.state.sentence_transformer_rf,
                     rag_app.state.RAG_OPENAI_API_KEY,
                     rag_app.state.RAG_OPENAI_API_BASE_URL,
                 )

src/lib/apis/rag/index.ts

@@ -413,3 +413,64 @@ export const updateEmbeddingConfig = async (token: string, payload: EmbeddingMod
 
 	return res;
 };
+
+export const getRerankingConfig = async (token: string) => {
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/reranking`, {
+		method: 'GET',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			error = err.detail;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
+type RerankingModelUpdateForm = {
+	reranking_model: string;
+};
+
+export const updateRerankingConfig = async (token: string, payload: RerankingModelUpdateForm) => {
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/reranking/update`, {
+		method: 'POST',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		},
+		body: JSON.stringify({
+			...payload
+		})
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			error = err.detail;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};

src/lib/components/documents/Settings/General.svelte

@@ -8,7 +8,9 @@
 		updateQuerySettings,
 		resetVectorDB,
 		getEmbeddingConfig,
-		updateEmbeddingConfig
+		updateEmbeddingConfig,
+		getRerankingConfig,
+		updateRerankingConfig
 	} from '$lib/apis/rag';
 
 	import { documents, models } from '$lib/stores';
@@ -23,11 +25,13 @@
 
 	let scanDirLoading = false;
 	let updateEmbeddingModelLoading = false;
+	let updateRerankingModelLoading = false;
 
 	let showResetConfirm = false;
 
 	let embeddingEngine = '';
 	let embeddingModel = '';
+	let rerankingModel = '';
 
 	let OpenAIKey = '';
 	let OpenAIUrl = '';
@@ -115,6 +119,29 @@
 		}
 	};
 
+	const rerankingModelUpdateHandler = async () => {
+		console.log('Update reranking model attempt:', rerankingModel);
+
+		updateRerankingModelLoading = true;
+		const res = await updateRerankingConfig(localStorage.token, {
+			reranking_model: rerankingModel,
+		}).catch(async (error) => {
+			toast.error(error);
+			await setRerankingConfig();
+			return null;
+		});
+		updateRerankingModelLoading = false;
+
+		if (res) {
+			console.log('rerankingModelUpdateHandler:', res);
+			if (res.status === true) {
+				toast.success($i18n.t('Reranking model set to "{{reranking_model}}"', res), {
+					duration: 1000 * 10
+				});
+			}
+		}
+	};
+
 	const submitHandler = async () => {
 		const res = await updateRAGConfig(localStorage.token, {
 			pdf_extract_images: pdfExtractImages,
@@ -138,6 +165,14 @@
 		}
 	};
 
+	const setRerankingConfig = async () => {
+		const rerankingConfig = await getRerankingConfig(localStorage.token);
+
+		if (rerankingConfig) {
+			rerankingModel = rerankingConfig.reranking_model;
+		}
+	};
+
 	onMount(async () => {
 		const res = await getRAGConfig(localStorage.token);
 
@@ -149,6 +184,7 @@
 		}
 
 		await setEmbeddingConfig();
+		await setRerankingConfig();
 
 		querySettings = await getQuerySettings(localStorage.token);
 	});
@@ -349,6 +385,73 @@
 
 				<hr class=" dark:border-gray-700 my-3" />
 
+				<div class=" ">
+					<div class=" mb-2 text-sm font-medium">{$i18n.t('Update Reranking Model')}</div>
+
+					<div class="flex w-full">
+						<div class="flex-1 mr-2">
+							<input
+								class="w-full rounded-lg py-2 px-4 text-sm dark:text-gray-300 dark:bg-gray-850 outline-none"
+								placeholder={$i18n.t('Update reranking model (e.g. {{model}})', {
+									model: rerankingModel.slice(-40)
+								})}
+								bind:value={rerankingModel}
+							/>
+						</div>
+						<button
+							class="px-2.5 bg-gray-100 hover:bg-gray-200 text-gray-800 dark:bg-gray-850 dark:hover:bg-gray-800 dark:text-gray-100 rounded-lg transition"
+							on:click={() => {
+								rerankingModelUpdateHandler();
+							}}
+							disabled={updateRerankingModelLoading}
+						>
+							{#if updateRerankingModelLoading}
+								<div class="self-center">
+									<svg
+										class=" w-4 h-4"
+										viewBox="0 0 24 24"
+										fill="currentColor"
+										xmlns="http://www.w3.org/2000/svg"
+										><style>
+											.spinner_ajPY {
+												transform-origin: center;
+												animation: spinner_AtaB 0.75s infinite linear;
+											}
+											@keyframes spinner_AtaB {
+												100% {
+													transform: rotate(360deg);
+												}
+											}
+										</style><path
+											d="M12,1A11,11,0,1,0,23,12,11,11,0,0,0,12,1Zm0,19a8,8,0,1,1,8-8A8,8,0,0,1,12,20Z"
+											opacity=".25"
+										/><path
+											d="M10.14,1.16a11,11,0,0,0-9,8.92A1.59,1.59,0,0,0,2.46,12,1.52,1.52,0,0,0,4.11,10.7a8,8,0,0,1,6.66-6.61A1.42,1.42,0,0,0,12,2.69h0A1.57,1.57,0,0,0,10.14,1.16Z"
+											class="spinner_ajPY"
+										/></svg
+									>
+								</div>
+							{:else}
+								<svg
+									xmlns="http://www.w3.org/2000/svg"
+									viewBox="0 0 16 16"
+									fill="currentColor"
+									class="w-4 h-4"
+								>
+									<path
+										d="M8.75 2.75a.75.75 0 0 0-1.5 0v5.69L5.03 6.22a.75.75 0 0 0-1.06 1.06l3.5 3.5a.75.75 0 0 0 1.06 0l3.5-3.5a.75.75 0 0 0-1.06-1.06L8.75 8.44V2.75Z"
+									/>
+									<path
+										d="M3.5 9.75a.75.75 0 0 0-1.5 0v1.5A2.75 2.75 0 0 0 4.75 14h6.5A2.75 2.75 0 0 0 14 11.25v-1.5a.75.75 0 0 0-1.5 0v1.5c0 .69-.56 1.25-1.25 1.25h-6.5c-.69 0-1.25-.56-1.25-1.25v-1.5Z"
+									/>
+								</svg>
+							{/if}
+						</button>
+					</div>
+				</div>
+
+				<hr class=" dark:border-gray-700 my-3" />
+
 				<div class="  flex w-full justify-between">
 					<div class=" self-center text-xs font-medium">
 						{$i18n.t('Scan for documents from {{path}}', { path: '/data/docs' })}