feat: implement lazy loading for faiss and sentence_transformers in RAG services

Author: sanjeevafk

api/services/rag/embeddings.py

@@ -10,8 +10,6 @@
 from google.genai import types as genai_types
 from openai import AsyncOpenAI
 from tenacity import retry, stop_after_attempt, wait_exponential
-
-from sentence_transformers import SentenceTransformer
 from api.config import get_settings
 
 logger = structlog.get_logger(__name__)
@@ -48,6 +46,8 @@ def reload_clients(self):
         elif self.provider == "local_bge":
             # Forcing CPU for stability as requested
             device = "cpu"
+            from sentence_transformers import SentenceTransformer
+
             self.local_model = SentenceTransformer(str(self.model), device=device)
             logger.info("local_model_loaded", model=self.model, device=device)
         else:

api/services/rag/filesystem_rag_store.py

@@ -11,14 +11,20 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
-import faiss
 import numpy as np
 import structlog
 from filelock import FileLock
 from rank_bm25 import BM25Okapi
 
 logger = structlog.get_logger(__name__)
 
+
+def _load_faiss() -> Any:
+    """Import faiss only when filesystem RAG operations need it."""
+    import faiss
+
+    return faiss
+
 @dataclass
 class RetrievalResult:
     chunk_id: str
@@ -86,6 +92,7 @@ async def ingest(
 
             # 3. Update FAISS Index
             dim = len(embeddings[0]) if embeddings else 768
+            faiss = _load_faiss()
             if paths["vectors"].exists():
                 index = faiss.read_index(str(paths["vectors"]))
             else:
@@ -141,6 +148,7 @@ def load_namespace(self, namespace: str):
         with open(paths["chunks"], "r", encoding="utf-8") as f:
             chunks = json.load(f)
 
+        faiss = _load_faiss()
         index = faiss.read_index(str(paths["vectors"]))
 
         with open(paths["bm25"], "rb") as f:
@@ -189,6 +197,7 @@ def _bootstrap_indices_from_chunks(self, namespace: str) -> None:
                 return
 
             dim = len(embeddings[0])
+            faiss = _load_faiss()
             index = faiss.IndexHNSWFlat(dim, 32)
             index.hnsw.efConstruction = 200
             vectors = np.array(embeddings).astype("float32")
@@ -242,6 +251,7 @@ async def retrieve(
             bm25 = data["bm25"]
 
             # 1. Vector Search (FAISS)
+            faiss = _load_faiss()
             xq = np.array([query_embedding]).astype("float32")
             faiss.normalize_L2(xq) # Assuming cosine similarity if index is inner product, 
                                    # but IndexHNSWFlat uses L2 distance by default.

api/tests/test_lazy_imports.py

@@ -0,0 +1,31 @@
+import builtins
+import importlib
+import sys
+
+
+def test_embeddings_module_import_does_not_require_sentence_transformers(monkeypatch):
+    original_import = builtins.__import__
+
+    def guarded_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "sentence_transformers":
+            raise ModuleNotFoundError("blocked for test")
+        return original_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", guarded_import)
+    sys.modules.pop("api.services.rag.embeddings", None)
+    module = importlib.import_module("api.services.rag.embeddings")
+    assert module is not None
+
+
+def test_filesystem_rag_store_module_import_does_not_require_faiss(monkeypatch):
+    original_import = builtins.__import__
+
+    def guarded_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "faiss":
+            raise ModuleNotFoundError("blocked for test")
+        return original_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", guarded_import)
+    sys.modules.pop("api.services.rag.filesystem_rag_store", None)
+    module = importlib.import_module("api.services.rag.filesystem_rag_store")
+    assert module is not None

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "sentry-sdk[fastapi]>=2.20.0",
     "slowapi>=0.1.9",
     "standardwebhooks>=1.0.0",
+    "jsonschema>=4.0.0",
     "tiktoken>=0.7.0",
     "faiss-cpu>=1.8.0",
     "rank-bm25>=0.2.2",