fix: share ML model instances to reduce startup time (#412)

* fix: share ML model instances to reduce startup time

The backend startup was slow because RetrieverTools.initialize() creates
6 retriever chains, and each one independently loaded its own copy of
the embedding model (thenlper/gte-large) and reranker model
(BAAI/bge-reranker-base). That meant 12 heavy model loads when only
2 are actually needed, since all chains use the same model config.

This fix creates both models once at the top of initialize() and passes
the shared instances down through HybridRetrieverChain,
SimilarityRetrieverChain, and FAISSVectorDatabase. Both models are
stateless (they only run encode/score inference) so sharing a single
instance across all chains is safe. Each chain still builds its own
independent FAISS index with its own documents.

Startup model loading goes from ~34s to ~7s on a local machine (4.9x).

Resolves #88

Signed-off-by: Harsh Kumar <harshkumar3446@gmail.com>

* refactor(backend): build every embedding model in one factory

RetrieverTools had its own copy of the embedding setup. That copy built
plain Gemini embeddings without the retry on temporary 503 errors, and
it passed the old model_name argument to VertexAIEmbeddings. The shared
and per-database models now come from create_embedding_model, so the
two paths cannot drift.

Signed-off-by: Vitor Bandeira <vvbandeira@precisioninno.com>

* fix(backend): share the HuggingFace reranker only when it is used

With RERANKER_TYPE=VERTEX_AI, the chains build a Vertex AI reranker.
Building the shared CrossEncoder anyway downloaded and loaded a model
that no chain used.

Signed-off-by: Vitor Bandeira <vvbandeira@precisioninno.com>

---------

Signed-off-by: Harsh Kumar <harshkumar3446@gmail.com>
Signed-off-by: Vitor Bandeira <vvbandeira@precisioninno.com>
Co-authored-by: Harsh Kumar <harshkumar3446@gmail.com>

Author: vvbandeira

backend/src/agents/retriever_tools.py

@@ -1,12 +1,15 @@
 import os
+import logging
 from typing import Tuple, Optional, Union
 from dotenv import load_dotenv
 
 from langchain_core.tools import tool
 from langchain_classic.retrievers import EnsembleRetriever
 from langchain_classic.retrievers import ContextualCompressionRetriever
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 
 from ..chains.hybrid_retriever_chain import HybridRetrieverChain
+from ..vectorstores.faiss import create_embedding_model
 from ..tools.format_docs import format_docs
 
 load_dotenv()
@@ -46,6 +49,19 @@ def initialize(
         use_cuda: bool = False,
         fast_mode: bool = False,
     ) -> None:
+        # Create shared model instances once
+        embedding_model = create_embedding_model(
+            embeddings_config["type"], embeddings_config["name"], use_cuda
+        )
+        logging.info("Shared embedding model created.")
+
+        # Only the HuggingFace reranker is a local model worth sharing. The
+        # chains build the Vertex AI reranker as an API client.
+        reranker_model: Optional[HuggingFaceCrossEncoder] = None
+        if os.getenv("RERANKER_TYPE", "HF").upper() != "VERTEX_AI":
+            reranker_model = HuggingFaceCrossEncoder(model_name=reranking_model_name)
+            logging.info("Shared reranker model created.")
+
         markdown_docs_map = {
             "general": [
                 "./data/markdown/OR_docs",
@@ -100,6 +116,8 @@ def initialize(
             contextual_rerank=True,
             search_k=search_k,
             chunk_size=chunk_size,
+            embedding_model=embedding_model,
+            reranker_model=reranker_model,
         )
         general_retriever_chain.create_hybrid_retriever()
         RetrieverTools.general_retriever = general_retriever_chain.retriever
@@ -115,6 +133,8 @@ def initialize(
             contextual_rerank=True,
             search_k=search_k,
             chunk_size=chunk_size,
+            embedding_model=embedding_model,
+            reranker_model=reranker_model,
         )
         install_retriever_chain.create_hybrid_retriever()
         RetrieverTools.install_retriever = install_retriever_chain.retriever
@@ -131,6 +151,8 @@ def initialize(
             contextual_rerank=True,
             search_k=search_k,
             chunk_size=chunk_size,
+            embedding_model=embedding_model,
+            reranker_model=reranker_model,
         )
         commands_retriever_chain.create_hybrid_retriever()
         RetrieverTools.commands_retriever = commands_retriever_chain.retriever
@@ -146,6 +168,8 @@ def initialize(
             contextual_rerank=True,
             search_k=search_k,
             chunk_size=chunk_size,
+            embedding_model=embedding_model,
+            reranker_model=reranker_model,
         )
         yosys_rtdocs_retriever_chain.create_hybrid_retriever()
         RetrieverTools.yosys_rtdocs_retriever = yosys_rtdocs_retriever_chain.retriever
@@ -161,6 +185,8 @@ def initialize(
             contextual_rerank=True,
             search_k=search_k,
             chunk_size=chunk_size,
+            embedding_model=embedding_model,
+            reranker_model=reranker_model,
         )
         klayout_retriever_chain.create_hybrid_retriever()
         RetrieverTools.klayout_retriever = klayout_retriever_chain.retriever
@@ -176,6 +202,8 @@ def initialize(
             contextual_rerank=True,
             search_k=search_k,
             chunk_size=chunk_size,
+            embedding_model=embedding_model,
+            reranker_model=reranker_model,
         )
         errinfo_retriever_chain.create_hybrid_retriever()
         RetrieverTools.errinfo_retriever = errinfo_retriever_chain.retriever

backend/src/chains/hybrid_retriever_chain.py

@@ -15,7 +15,7 @@
 from .similarity_retriever_chain import SimilarityRetrieverChain
 from .mmr_retriever_chain import MMRRetrieverChain
 from .bm25_retriever_chain import BM25RetrieverChain
-from ..vectorstores.faiss import FAISSVectorDatabase
+from ..vectorstores.faiss import EmbeddingModel, FAISSVectorDatabase
 
 
 class HybridRetrieverChain(BaseChain):
@@ -37,6 +37,8 @@ def __init__(
         weights: list[float] = [0.33, 0.33, 0.33],
         chunk_size: int = 500,
         contextual_rerank: bool = False,
+        embedding_model: Optional[EmbeddingModel] = None,
+        reranker_model: Optional[HuggingFaceCrossEncoder] = None,
     ):
         super().__init__(
             llm_model=llm_model,
@@ -47,6 +49,8 @@ def __init__(
 
         self.reranking_model_name: Optional[str] = reranking_model_name
         self.use_cuda: bool = use_cuda
+        self.embedding_model: Optional[EmbeddingModel] = embedding_model
+        self.reranker_model: Optional[HuggingFaceCrossEncoder] = reranker_model
 
         self.search_k: int = search_k
         self.weights: list[float] = weights
@@ -73,6 +77,7 @@ def create_hybrid_retriever(self) -> None:
             html_docs_path=self.html_docs_path,
             chunk_size=self.chunk_size,
             use_cuda=self.use_cuda,
+            embedding_model=self.embedding_model,
         )
         if self.vector_db is None:
             cur_path = os.path.abspath(__file__)
@@ -147,12 +152,10 @@ def create_hybrid_retriever(self) -> None:
                 )
                 logging.info("Using Vertex AI reranker")
             else:
-                compressor = CrossEncoderReranker(
-                    model=HuggingFaceCrossEncoder(
-                        model_name=self.reranking_model_name
-                    ),
-                    top_n=self.search_k,
+                reranker = self.reranker_model or HuggingFaceCrossEncoder(
+                    model_name=self.reranking_model_name
                 )
+                compressor = CrossEncoderReranker(model=reranker, top_n=self.search_k)
                 logging.info("Using HuggingFace CrossEncoder reranker")
 
             self.retriever = ContextualCompressionRetriever(

backend/src/chains/similarity_retriever_chain.py

@@ -3,11 +3,11 @@
 
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
 from langchain_core.documents import Document
-from langchain_google_vertexai import ChatVertexAI
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_google_vertexai import ChatVertexAI
 from langchain_ollama import ChatOllama
 
-from ..vectorstores.faiss import FAISSVectorDatabase
+from ..vectorstores.faiss import EmbeddingModel, FAISSVectorDatabase
 from .base_chain import BaseChain
 
 
@@ -28,6 +28,7 @@ def __init__(
         embeddings_config: Optional[dict[str, str]] = None,
         use_cuda: bool = False,
         chunk_size: int = 500,
+        embedding_model: Optional[EmbeddingModel] = None,
     ):
         super().__init__(
             llm_model=llm_model,
@@ -40,6 +41,7 @@ def __init__(
 
         self.embeddings_config: Optional[dict[str, str]] = embeddings_config
         self.use_cuda: bool = use_cuda
+        self.embedding_model: Optional[EmbeddingModel] = embedding_model
 
         self.markdown_docs_path: Optional[list[str]] = markdown_docs_path
         self.other_docs_path: Optional[list[str]] = other_docs_path
@@ -125,6 +127,7 @@ def create_vector_db(self) -> None:
                 embeddings_model_name=self.embeddings_config["name"],
                 embeddings_type=self.embeddings_config["type"],
                 use_cuda=self.use_cuda,
+                embedding_model=self.embedding_model,
             )
         else:
             raise ValueError("Embeddings model config not provided correctly.")

backend/src/vectorstores/faiss.py

@@ -72,6 +72,38 @@ def embed_query(self, text: str, **kwargs: Any) -> list[float]:
         return super().embed_query(text, **kwargs)
 
 
+EmbeddingModel = Union[
+    HuggingFaceEmbeddings, GoogleGenerativeAIEmbeddings, VertexAIEmbeddings
+]
+
+
+def create_embedding_model(
+    embeddings_type: str, model_name: str, use_cuda: bool = False
+) -> EmbeddingModel:
+    """Build the embedding model for an EMBEDDINGS_TYPE value.
+
+    Every caller uses this function, so a shared model has the same
+    configuration as one built by FAISSVectorDatabase.
+    """
+    if embeddings_type == "GOOGLE_GENAI":
+        logging.info("Using Google GenerativeAI embeddings...")
+        return _RetryingGoogleGenerativeAIEmbeddings(
+            model=model_name, task_type="retrieval_document"
+        )
+    if embeddings_type == "GOOGLE_VERTEXAI":
+        logging.info("Using Google VertexAI embeddings...")
+        return VertexAIEmbeddings(model=model_name)
+    if embeddings_type == "HF":
+        logging.info("Using HuggingFace embeddings...")
+        return HuggingFaceEmbeddings(
+            model_name=model_name,
+            multi_process=False,
+            encode_kwargs={"normalize_embeddings": True},
+            model_kwargs={"device": "cuda" if use_cuda else "cpu"},
+        )
+    raise ValueError("Invalid embeddings type specified.")
+
+
 class FAISSVectorDatabase:
     def __init__(
         self,
@@ -80,37 +112,17 @@ def __init__(
         distance_strategy: DistanceStrategy = DistanceStrategy.COSINE,
         debug: bool = False,
         use_cuda: bool = False,
+        embedding_model: Optional[EmbeddingModel] = None,
     ):
         self.embeddings_model_name = embeddings_model_name
 
-        model_kwargs = {"device": "cuda"} if use_cuda else {"device": "cpu"}
-
-        self.embedding_model: Union[
-            HuggingFaceEmbeddings, GoogleGenerativeAIEmbeddings, VertexAIEmbeddings
-        ]
-
-        if embeddings_type == "GOOGLE_GENAI":
-            self.embedding_model = _RetryingGoogleGenerativeAIEmbeddings(
-                model=self.embeddings_model_name,
-                task_type="retrieval_document",
-            )
-            logging.info("Using Google GenerativeAI embeddings...")
-
-        elif embeddings_type == "GOOGLE_VERTEXAI":
-            self.embedding_model = VertexAIEmbeddings(model=self.embeddings_model_name)
-            logging.info("Using Google VertexAI embeddings...")
-
-        elif embeddings_type == "HF":
-            self.embedding_model = HuggingFaceEmbeddings(
-                model_name=self.embeddings_model_name,
-                multi_process=False,
-                encode_kwargs={"normalize_embeddings": True},
-                model_kwargs=model_kwargs,
+        self.embedding_model: EmbeddingModel = (
+            embedding_model
+            if embedding_model is not None
+            else create_embedding_model(
+                embeddings_type, self.embeddings_model_name, use_cuda
             )
-            logging.info("Using HuggingFace embeddings...")
-
-        else:
-            raise ValueError("Invalid embdeddings type specified.")
+        )
 
         self.debug = debug
         self.distance_strategy = distance_strategy

backend/tests/test_faiss_vectorstore.py

@@ -61,7 +61,7 @@ def test_init_with_google_vertexai_embeddings(self):
 
     def test_init_with_invalid_embeddings_type(self):
         """Test initialization with invalid embeddings type raises error."""
-        with pytest.raises(ValueError, match="Invalid embdeddings type specified"):
+        with pytest.raises(ValueError, match="Invalid embeddings type specified"):
             FAISSVectorDatabase(
                 embeddings_type="INVALID", embeddings_model_name="test-model"
             )