Merge pull request #461 from sahilds1/441-embedding-models

[#441] [IMPROVE] Preload embedding model at startup

Author: sahilds1

README.md

@@ -73,6 +73,11 @@ df = pd.read_sql(query, engine)
 
 #### Django REST
 - The email and password are set in `server/api/management/commands/createsu.py`
+- Backend tests can be run using `pytest` by running the below command inside the running backend container:
+
+```
+docker compose exec backend pytest api/ -v
+```
 
 ## API Documentation
 

server/api/apps.py

@@ -4,3 +4,38 @@
 class ApiConfig(AppConfig):
     default_auto_field = 'django.db.models.BigAutoField'
     name = 'api'
+
+    def ready(self):
+        
+        try:
+            import os
+            import sys
+    
+            # ready() runs in every Django process: migrate, test, shell, runserver, etc.
+            # Only preload the model when we're actually going to serve requests.
+            # Dev (docker-compose.yml) runs `manage.py runserver 0.0.0.0:8000`.
+            # Prod (Dockerfile.prod CMD) runs `manage.py runserver 0.0.0.0:8000 --noreload`.
+            # entrypoint.prod.sh also runs migrate, createsu, and populatedb before exec'ing
+            # runserver — the guard below correctly skips model loading for those commands too.
+            if sys.argv[1:2] != ['runserver']:
+                return
+    
+            # Dev's autoreloader spawns two processes: a parent file-watcher and a child
+            # server. ready() runs in both, but only the child (RUN_MAIN=true) serves
+            # requests. Skip the parent to avoid loading the model twice on each file change.
+            # Prod uses --noreload so RUN_MAIN is never set; 'noreload' in sys.argv handles that case.
+            if os.environ.get('RUN_MAIN') != 'true' and '--noreload' not in sys.argv:
+                return
+    
+            # Note: paraphrase-MiniLM-L6-v2 (~80MB) is downloaded from HuggingFace on first
+            # use and cached to ~/.cache/torch/sentence_transformers/ inside the container.
+            # That cache is ephemeral — every container rebuild re-downloads the model unless
+            # a volume is mounted at that path.
+            from .services.sentencetTransformer_model import TransformerModel
+            TransformerModel.get_instance()
+        except Exception:
+            # TransformerModel._instance stays None on failure, so the first actual request 
+            # that calls get_instance() will attempt to load the model again.
+            import logging
+            logger = logging.getLogger(__name__)
+            logger.exception("Failed to preload the embedding model at startup")

server/api/services/embedding_services.py

@@ -2,6 +2,7 @@
 import logging
 from statistics import median
 
+# Use Q objects to express OR conditions in Django queries
 from django.db.models import Q
 from pgvector.django import L2Distance
 
@@ -11,18 +12,17 @@
 
 logger = logging.getLogger(__name__)
 
-def get_closest_embeddings(
-    user, message_data, document_name=None, guid=None, num_results=10
-):
+
+def build_query(user, embedding_vector, document_name=None, guid=None, num_results=10):
     """
-    Find the closest embeddings to a given message for a specific user.
+    Build an unevaluated QuerySet for the closest embeddings.
 
     Parameters
     ----------
     user : User
         The user whose uploaded documents will be searched
-    message_data : str
-        The input message to find similar embeddings for
+    embedding_vector : array-like
+        Pre-computed embedding vector to compare against
     document_name : str, optional
         Filter results to a specific document name
     guid : str, optional
@@ -32,59 +32,52 @@ def get_closest_embeddings(
 
     Returns
     -------
-    list[dict]
-        List of dictionaries containing embedding results with keys:
-        - name: document name
-        - text: embedded text content
-        - page_number: page number in source document
-        - chunk_number: chunk number within the document
-        - distance: L2 distance from query embedding
-        - file_id: GUID of the source file
+    QuerySet
+        Unevaluated Django QuerySet ordered by L2 distance, sliced to num_results
     """
-
-    encoding_start = time.time()
-    transformerModel = TransformerModel.get_instance().model
-    embedding_message = transformerModel.encode(message_data)
-    encoding_time = time.time() - encoding_start
-
-    db_query_start = time.time()
-
     # Django QuerySets are lazily evaluated
     if user.is_authenticated:
         # User sees their own files + files uploaded by superusers
-        closest_embeddings_query = (
-            Embeddings.objects.filter(
-                Q(upload_file__uploaded_by=user) | Q(upload_file__uploaded_by__is_superuser=True)
-            )
-            .annotate(
-                distance=L2Distance("embedding_sentence_transformers", embedding_message)
-            )
-            .order_by("distance")
+        queryset = Embeddings.objects.filter(
+            Q(upload_file__uploaded_by=user) | Q(upload_file__uploaded_by__is_superuser=True)
         )
     else:
         # Unauthenticated users only see superuser-uploaded files
-        closest_embeddings_query = (
-            Embeddings.objects.filter(upload_file__uploaded_by__is_superuser=True)
-            .annotate(
-                distance=L2Distance("embedding_sentence_transformers", embedding_message)
-            )
-            .order_by("distance")
-        )
+        queryset = Embeddings.objects.filter(upload_file__uploaded_by__is_superuser=True)
+
+    queryset = (
+        queryset
+        .annotate(distance=L2Distance("embedding_sentence_transformers", embedding_vector))
+        .order_by("distance")
+    )
 
     # Filtering to a document GUID takes precedence over a document name
     if guid:
-        closest_embeddings_query = closest_embeddings_query.filter(
-            upload_file__guid=guid
-        )
+        queryset = queryset.filter(upload_file__guid=guid)
     elif document_name:
-        closest_embeddings_query = closest_embeddings_query.filter(name=document_name)
+        queryset = queryset.filter(name=document_name)
 
     # Slicing is equivalent to SQL's LIMIT clause
-    closest_embeddings_query = closest_embeddings_query[:num_results]
+    return queryset[:num_results]
+
+
+def evaluate_query(queryset):
+    """
+    Evaluate a QuerySet and return a list of result dicts.
+
+    Parameters
+    ----------
+    queryset : iterable
+        Iterable of Embeddings objects (or any objects with the expected attributes)
 
+    Returns
+    -------
+    list[dict]
+        List of dicts with keys: name, text, page_number, chunk_number, distance, file_id
+    """
     # Iterating evaluates the QuerySet and hits the database
     # TODO: Research improving the query evaluation performance
-    results = [
+    return [
         {
             "name": obj.name,
             "text": obj.text,
@@ -93,13 +86,36 @@ def get_closest_embeddings(
             "distance": obj.distance,
             "file_id": obj.upload_file.guid if obj.upload_file else None,
         }
-        for obj in closest_embeddings_query
+        for obj in queryset
     ]
 
-    db_query_time = time.time() - db_query_start
 
+def log_usage(
+    results, message_data, user, guid, document_name, num_results, encoding_time, db_query_time
+):
+    """
+    Create a SemanticSearchUsage record. Swallows exceptions so search isn't interrupted.
+
+    Parameters
+    ----------
+    results : list[dict]
+        The search results, each containing a "distance" key
+    message_data : str
+        The original search query text
+    user : User
+        The user who performed the search
+    guid : str or None
+        Document GUID filter used in the search
+    document_name : str or None
+        Document name filter used in the search
+    num_results : int
+        Number of results requested
+    encoding_time : float
+        Time in seconds to encode the query
+    db_query_time : float
+        Time in seconds for the database query
+    """
     try:
-        # Handle user having no uploaded docs or doc filtering returning no matches
         if results:
             distances = [r["distance"] for r in results]
             SemanticSearchUsage.objects.create(
@@ -113,11 +129,10 @@ def get_closest_embeddings(
                 num_results_returned=len(results),
                 max_distance=max(distances),
                 median_distance=median(distances),
-                min_distance=min(distances)
+                min_distance=min(distances),
             )
         else:
             logger.warning("Semantic search returned no results")
-
             SemanticSearchUsage.objects.create(
                 query_text=message_data,
                 user=user if (user and user.is_authenticated) else None,
@@ -129,9 +144,58 @@ def get_closest_embeddings(
                 num_results_returned=0,
                 max_distance=None,
                 median_distance=None,
-                min_distance=None
+                min_distance=None,
             )
-    except Exception as e:
-        logger.error(f"Failed to create semantic search usage database record: {e}")
+    except Exception:
+        logger.exception("Failed to create semantic search usage database record")
+
+
+def get_closest_embeddings(
+    user, message_data, document_name=None, guid=None, num_results=10
+):
+    """
+    Find the closest embeddings to a given message for a specific user.
+
+    Parameters
+    ----------
+    user : User
+        The user whose uploaded documents will be searched
+    message_data : str
+        The input message to find similar embeddings for
+    document_name : str, optional
+        Filter results to a specific document name
+    guid : str, optional
+        Filter results to a specific document GUID (takes precedence over document_name)
+    num_results : int, default 10
+        Maximum number of results to return
+
+    Returns
+    -------
+    list[dict]
+        List of dictionaries containing embedding results with keys:
+        - name: document name
+        - text: embedded text content
+        - page_number: page number in source document
+        - chunk_number: chunk number within the document
+        - distance: L2 distance from query embedding
+        - file_id: GUID of the source file
+
+    Notes
+    -----
+    Creates a SemanticSearchUsage record. Swallows exceptions so search isn't interrupted.
+    """
+    encoding_start = time.time()
+    model = TransformerModel.get_instance().model
+    embedding_vector = model.encode(message_data)
+    encoding_time = time.time() - encoding_start
+
+    db_query_start = time.time()
+    queryset = build_query(user, embedding_vector, document_name, guid, num_results)
+    results = evaluate_query(queryset)
+    db_query_time = time.time() - db_query_start
+
+    log_usage(
+        results, message_data, user, guid, document_name, num_results, encoding_time, db_query_time
+    )
 
     return results