Merge pull request #64 from tjmlabs/gpu-on

feat: always on gpu for faster queries
tjmlabs · Nov 6, 2024 · 728bcb5 · 728bcb5
2 parents 3950ae6 + c7f70bf
commit 728bcb5
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/web/api/views.py b/web/api/views.py
@@ -138,7 +138,10 @@ async def create_collection(
             name=payload.name, owner=request.auth, metadata=payload.metadata
         )
         return 201, CollectionOut(
-            id=collection.id, name=collection.name, metadata=collection.metadata, num_documents=0
+            id=collection.id,
+            name=collection.name,
+            metadata=collection.metadata,
+            num_documents=0,
         )
     except IntegrityError:
         return 409, GenericError(
@@ -165,7 +168,12 @@ async def list_collections(request: Request) -> List[CollectionOut]:
         HTTPException: If there is an issue with the request or authentication.
     """
     collections = [
-        CollectionOut(id=c.id, name=c.name, metadata=c.metadata, num_documents=await c.document_count())
+        CollectionOut(
+            id=c.id,
+            name=c.name,
+            metadata=c.metadata,
+            num_documents=await c.document_count(),
+        )
         async for c in Collection.objects.filter(owner=request.auth)
     ]
     return collections
@@ -207,7 +215,10 @@ async def get_collection(
             name=collection_name, owner=request.auth
         )
         return 200, CollectionOut(
-            id=collection.id, name=collection.name, metadata=collection.metadata, num_documents=await collection.document_count()
+            id=collection.id,
+            name=collection.name,
+            metadata=collection.metadata,
+            num_documents=await collection.document_count(),
         )
     except Collection.DoesNotExist:
         return 404, GenericError(detail=f"Collection: {collection_name} doesn't exist")
@@ -252,7 +263,10 @@ async def partial_update_collection(
 
     await collection.asave()
     return 200, CollectionOut(
-        id=collection.id, name=collection.name, metadata=collection.metadata, num_documents=await collection.document_count()
+        id=collection.id,
+        name=collection.name,
+        metadata=collection.metadata,
+        num_documents=await collection.document_count(),
     )
 
 
@@ -947,7 +961,7 @@ async def search(
 
 
 async def get_query_embeddings(query: str) -> List:
-    EMBEDDINGS_URL = settings.EMBEDDINGS_URL
+    EMBEDDINGS_URL = settings.ALWAYS_ON_EMBEDDINGS_URL
     embed_token = settings.EMBEDDINGS_URL_TOKEN
     headers = {"Authorization": f"Bearer {embed_token}"}
     payload = {
@@ -1098,7 +1112,7 @@ async def embeddings(
     Raises:
         HttpError: If the documents cannot be embedded.
     """
-    EMBEDDINGS_URL = settings.EMBEDDINGS_URL
+    EMBEDDINGS_URL = settings.ALWAYS_ON_EMBEDDINGS_URL
     embed_token = settings.EMBEDDINGS_URL_TOKEN
     headers = {"Authorization": f"Bearer {embed_token}"}
     task = payload.task

diff --git a/web/config/settings.py b/web/config/settings.py
@@ -231,6 +231,8 @@
 
 # EMEDDING Service
 EMBEDDINGS_URL = env("EMBEDDINGS_URL")
+# Queries need to be fast, so we use a separate service for embeddings.
+ALWAYS_ON_EMBEDDINGS_URL = env("ALWAYS_ON_EMBEDDINGS_URL", default=EMBEDDINGS_URL)
 EMBEDDINGS_URL_TOKEN = env("EMBEDDINGS_URL_TOKEN")
 
 # Gotenberg