Merge pull request #414 from sahilds1/380-track-chatbot-costs

[#380] Track chatbot semantic search costs and performance

Author: taichan03

frontend/src/api/apiClient.ts

@@ -289,7 +289,9 @@ const sendAssistantMessage = async (
   previousResponseId?: string,
 ) => {
   try {
-    const response = await publicApi.post(`/v1/api/assistant`, {
+    // The adminApi interceptor doesn't gracefully omit the JWT token if you're not authenticated
+    const api = localStorage.getItem("access") ? adminApi : publicApi;
+    const response = await api.post(`/v1/api/assistant`, {
       message,
       previous_response_id: previousResponseId,
     });

server/api/migrations/0015_semanticsearchusage.py

@@ -0,0 +1,39 @@
+# Generated by Django 4.2.3 on 2025-11-26 21:02
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('api', '0014_alter_medrule_rule_type'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='SemanticSearchUsage',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('guid', models.UUIDField(default=uuid.uuid4, editable=False, unique=True)),
+                ('timestamp', models.DateTimeField(auto_now_add=True)),
+                ('query_text', models.TextField(blank=True, help_text='The search query text', null=True)),
+                ('document_name', models.TextField(blank=True, help_text='Document name filter if used', null=True)),
+                ('document_guid', models.UUIDField(blank=True, help_text='Document GUID filter if used', null=True)),
+                ('num_results_requested', models.IntegerField(default=10, help_text='Number of results requested')),
+                ('encoding_time', models.FloatField(help_text='Time to encode query in seconds')),
+                ('db_query_time', models.FloatField(help_text='Time for database query in seconds')),
+                ('num_results_returned', models.IntegerField(help_text='Number of results returned')),
+                ('min_distance', models.FloatField(blank=True, help_text='Minimum L2 distance (null if no results)', null=True)),
+                ('max_distance', models.FloatField(blank=True, help_text='Maximum L2 distance (null if no results)', null=True)),
+                ('median_distance', models.FloatField(blank=True, help_text='Median L2 distance (null if no results)', null=True)),
+                ('user', models.ForeignKey(blank=True, help_text='User who performed the search (null for unauthenticated users)', null=True, on_delete=django.db.models.deletion.CASCADE, related_name='semantic_searches', to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'ordering': ['-timestamp'],
+                'indexes': [models.Index(fields=['-timestamp'], name='api_semanti_timesta_0b5730_idx'), models.Index(fields=['user', '-timestamp'], name='api_semanti_user_id_e11ecb_idx')],
+            },
+        ),
+    ]

server/api/models/model_search_usage.py

@@ -0,0 +1,42 @@
+import uuid
+
+from django.db import models
+from django.conf import settings
+
+class SemanticSearchUsage(models.Model):
+    """
+    Tracks performance metrics and usage data for embedding searches.
+    """
+    guid = models.UUIDField(unique=True, default=uuid.uuid4, editable=False)
+    timestamp = models.DateTimeField(auto_now_add=True)
+    query_text = models.TextField(blank=True, null=True, help_text="The search query text")
+    document_name = models.TextField(blank=True, null=True, help_text="Document name filter if used")
+    document_guid = models.UUIDField(blank=True, null=True, help_text="Document GUID filter if used")
+    num_results_requested = models.IntegerField(default=10, help_text="Number of results requested")
+    user = models.ForeignKey(
+        settings.AUTH_USER_MODEL,
+        on_delete=models.CASCADE,
+        related_name='semantic_searches',
+        null=True,
+        blank=True,
+        help_text="User who performed the search (null for unauthenticated users)"
+    )
+    encoding_time = models.FloatField(help_text="Time to encode query in seconds")
+    db_query_time = models.FloatField(help_text="Time for database query in seconds")
+    num_results_returned = models.IntegerField(help_text="Number of results returned")
+    min_distance = models.FloatField(null=True, blank=True, help_text="Minimum L2 distance (null if no results)")
+    max_distance = models.FloatField(null=True, blank=True, help_text="Maximum L2 distance (null if no results)")
+    median_distance = models.FloatField(null=True, blank=True, help_text="Median L2 distance (null if no results)")
+
+
+    class Meta:
+        ordering = ['-timestamp']
+        indexes = [
+            models.Index(fields=['-timestamp']),
+            models.Index(fields=['user', '-timestamp']),
+        ]
+
+    def __str__(self):
+        total_time = self.encoding_time + self.db_query_time
+        user_display = self.user.email if self.user else "Anonymous"
+        return f"Search by {user_display} at {self.timestamp} ({total_time:.3f}s)"

server/api/services/embedding_services.py

@@ -1,11 +1,15 @@
+import time
+import logging
+from statistics import median
+
 from django.db.models import Q
 from pgvector.django import L2Distance
 
 from .sentencetTransformer_model import TransformerModel
-
-# Adjust import path as needed
 from ..models.model_embeddings import Embeddings
+from ..models.model_search_usage import SemanticSearchUsage
 
+logger = logging.getLogger(__name__)
 
 def get_closest_embeddings(
     user, message_data, document_name=None, guid=None, num_results=10
@@ -38,9 +42,14 @@ def get_closest_embeddings(
         - file_id: GUID of the source file
     """
 
+    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 = (
@@ -62,18 +71,19 @@ def get_closest_embeddings(
             .order_by("distance")
         )
 
-    # Filter by GUID if provided, otherwise filter by document name if provided
+    # Filtering to a document GUID takes precedence over a document name
     if guid:
         closest_embeddings_query = closest_embeddings_query.filter(
             upload_file__guid=guid
         )
     elif document_name:
         closest_embeddings_query = closest_embeddings_query.filter(name=document_name)
 
-    # Slice the results to limit to num_results
+    # Slicing is equivalent to SQL's LIMIT clause
     closest_embeddings_query = closest_embeddings_query[:num_results]
 
-    # Format the results to be returned
+    # Iterating evaluates the QuerySet and hits the database
+    # TODO: Research improving the query evaluation performance
     results = [
         {
             "name": obj.name,
@@ -86,4 +96,42 @@ def get_closest_embeddings(
         for obj in closest_embeddings_query
     ]
 
+    db_query_time = time.time() - db_query_start
+
+    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(
+                query_text=message_data,
+                user=user if (user and user.is_authenticated) else None,
+                document_guid=guid,
+                document_name=document_name,
+                num_results_requested=num_results,
+                encoding_time=encoding_time,
+                db_query_time=db_query_time,
+                num_results_returned=len(results),
+                max_distance=max(distances),
+                median_distance=median(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,
+                document_guid=guid,
+                document_name=document_name,
+                num_results_requested=num_results,
+                encoding_time=encoding_time,
+                db_query_time=db_query_time,
+                num_results_returned=0,
+                max_distance=None,
+                median_distance=None,
+                min_distance=None
+            )
+    except Exception as e:
+        logger.error(f"Failed to create semantic search usage database record: {e}")
+
     return results