feat(api): adds migration for tracking vector store indexing status

.github/actions/lfai-core/action.yaml

@@ -18,6 +18,7 @@ runs:
       id: set-env-var
       run: |
         echo "ANON_KEY=$(uds zarf tools kubectl get secret supabase-bootstrap-jwt -n leapfrogai -o jsonpath='{.data.anon-key}' | base64 -d)" >> "$GITHUB_ENV"
+        echo "SERVICE_KEY=$(uds zarf tools kubectl get secret supabase-bootstrap-jwt -n leapfrogai -o jsonpath='{.data.service-key}' | base64 -d)" >> "$GITHUB_ENV"
 
     - name: Deploy LFAI-API
       shell: bash

.github/workflows/e2e-playwright.yaml

@@ -1,4 +1,4 @@
-# End-to-end testing that deploys and tests Supabase, API, UI, and VLLM
+# End-to-end testing that deploys and tests Supabase, API, and UI
 
 name: e2e-playwright
 on:

.github/workflows/e2e-text-backend-full-cpu.yaml

@@ -0,0 +1,101 @@
+# End-to-end testing that deploys and tests Supabase, API, llama-cpp-python, and text-embeddings
+
+name: e2e-text-backend-full-cpu
+on:
+  pull_request:
+    types:
+      - opened            # default trigger
+      - reopened          # default trigger
+      - synchronize       # default trigger
+      - ready_for_review  # don't run on draft PRs
+      - milestoned        # allows us to trigger on bot PRs
+    paths:
+      # Catch-all
+      - "**"
+
+      # Ignore updates to the .github directory, unless it's this current file
+      - "!.github/**"
+      - ".github/workflows/e2e-text-backend-full-cpu.yaml"
+      - ".github/actions/uds-cluster/action.yaml"
+
+      # Ignore docs and website things
+      - "!**.md"
+      - "!docs/**"
+      - "!adr/**"
+      - "!website/**"
+      - "!netlify.toml"
+
+      # Ignore updates to generic github metadata files
+      - "!CODEOWNERS"
+      - "!.gitignore"
+      - "!LICENSE"
+
+      # Ignore local development files
+      - "!.pre-commit-config.yaml"
+
+      # Ignore non e2e tests changes
+      - "!tests/pytest/**"
+
+      # Ignore LFAI-UI source code changes
+      - "!src/leapfrogai_ui/**"
+
+      # Ignore changes to unrelated packages
+      - "!packages/k3d-gpu/**"
+      - "!packages/repeater/**"
+      - "!packages/ui/**"
+      - "!packages/vllm/**"
+      - "!packages/whisper/**"
+
+
+
+concurrency:
+  group: e2e-text-backend-full-cpu-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  e2e_text_backend_full_cpu:
+    runs-on: ai-ubuntu-big-boy-8-core
+    if: ${{ !github.event.pull_request.draft }}
+
+    steps:
+        - name: Checkout Repo
+          uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+
+        - name: Setup Python
+          uses: ./.github/actions/python
+
+        - name: Setup UDS Cluster
+          uses: ./.github/actions/uds-cluster
+          with:
+            registry1Username: ${{ secrets.IRON_BANK_ROBOT_USERNAME }}
+            registry1Password: ${{ secrets.IRON_BANK_ROBOT_PASSWORD }}
+
+        - name: Setup LFAI-API and Supabase
+          uses: ./.github/actions/lfai-core
+
+        ##########
+        # text-embeddings
+        ##########
+        - name: Deploy text-embeddings
+          run: |
+            make build-text-embeddings LOCAL_VERSION=e2e-test
+            docker image prune -af
+            uds zarf package deploy packages/text-embeddings/zarf-package-text-embeddings-amd64-e2e-test.tar.zst -l=trace --confirm
+            rm packages/text-embeddings/zarf-package-text-embeddings-amd64-e2e-test.tar.zst
+
+        ##########
+        # llama
+        ##########
+        - name: Deploy llama-cpp-python
+          run: |
+            make build-llama-cpp-python LOCAL_VERSION=e2e-test
+            docker image prune -af
+            uds zarf package deploy packages/llama-cpp-python/zarf-package-llama-cpp-python-amd64-e2e-test.tar.zst -l=trace --confirm
+            rm packages/llama-cpp-python/zarf-package-llama-cpp-python-amd64-e2e-test.tar.zst
+
+        ##########
+        # Test
+        ##########
+        - name: Test Text Backend
+          run: |
+            python -m pytest ./tests/e2e/test_text_backend_full.py -v

packages/api/supabase/migrations/20240827103100_v0.11.1_indexing_status.sql

@@ -0,0 +1,24 @@
+-- Update the vector_store_file table to add an updated_at column
+ALTER TABLE vector_store_file ADD COLUMN updated_at timestamp DEFAULT timezone('utc', now()) NOT NULL;
+
+-- Add an index on user_id for faster queries
+CREATE INDEX idx_vector_store_file_user_id ON vector_store_file(user_id);
+
+-- Create a function to update the updated_at column
+CREATE OR REPLACE FUNCTION update_modified_column()
+RETURNS TRIGGER AS $$
+BEGIN
+    NEW.updated_at = timezone('utc', now());
+    RETURN NEW;
+END;
+$$ language 'plpgsql';
+
+-- Create a trigger to automatically update the updated_at column
+CREATE TRIGGER update_vector_store_file_modtime
+BEFORE UPDATE ON vector_store_file
+FOR EACH ROW
+EXECUTE FUNCTION update_modified_column();
+
+-- Enable Supabase realtime for the vector_store_file table
+alter publication supabase_realtime
+add table vector_store_file;

packages/text-embeddings/main.py

@@ -17,7 +17,10 @@
 
 class InstructorEmbedding:
     async def CreateEmbedding(self, request: EmbeddingRequest, context: GrpcContext):
-        embeddings = model.encode(request.inputs)
+        # Run the CPU-intensive encoding in a separate thread
+        embeddings = await asyncio.to_thread(
+            model.encode, sentences=request.inputs, show_progress_bar=True
+        )
 
         embeddings = [Embedding(embedding=inner_list) for inner_list in embeddings]
         return EmbeddingResponse(embeddings=embeddings)

pyproject.toml

@@ -27,6 +27,8 @@ dev = [
     "requests-toolbelt",
     "pytest",
     "huggingface_hub[cli,hf_transfer]",
+    "fastapi",
+    "supabase == 2.6.0"
 ]
 
 dev-whisper = ["ctranslate2 == 4.1.0", "transformers[torch] == 4.39.3"]

src/leapfrogai_api/backend/rag/index.py

@@ -3,7 +3,7 @@
 import logging
 import tempfile
 import time
-from fastapi import HTTPException, UploadFile, status
+from fastapi import HTTPException, UploadFile, status, BackgroundTasks
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from openai.types.beta.vector_store import FileCounts, VectorStore
@@ -89,7 +89,7 @@ async def index_file(self, vector_store_id: str, file_id: str) -> VectorStoreFil
                     ),
                     object="vector_store.file",
                     status=VectorStoreFileStatus.FAILED.value,
-                    usage_bytes=0,
+                    usage_bytes=0,  # Leave blank to have Postgres calculate the document bytes
                     vector_store_id=vector_store_id,
                 )
                 return await crud_vector_store_file.create(object_=vector_store_file)
@@ -100,7 +100,7 @@ async def index_file(self, vector_store_id: str, file_id: str) -> VectorStoreFil
                 last_error=None,
                 object="vector_store.file",
                 status=VectorStoreFileStatus.IN_PROGRESS.value,
-                usage_bytes=0,
+                usage_bytes=0,  # Leave blank to have Postgres calculate the document bytes
                 vector_store_id=vector_store_id,
             )
 
@@ -154,59 +154,90 @@ async def index_files(
         return responses
 
     async def create_new_vector_store(
-        self, request: CreateVectorStoreRequest
+        self,
+        request: CreateVectorStoreRequest,
+        background_tasks: BackgroundTasks | None = None,
     ) -> VectorStore:
         """Create a new vector store given a set of file ids"""
         crud_vector_store = CRUDVectorStore(db=self.db)
 
-        last_active_at = int(time.time())
-
-        expires_after, expires_at = request.get_expiry(last_active_at)
+        current_time = int(time.time())
+        expires_after, expires_at = request.get_expiry(current_time)
 
         try:
-            vector_store = VectorStore(
+            # Create a placeholder vector store
+            placeholder_vector_store = VectorStore(
                 id="",  # Leave blank to have Postgres generate a UUID
-                usage_bytes=0,  # Automatically calculated by DB
+                name=request.name or "",
+                status=VectorStoreStatus.IN_PROGRESS.value,
+                object="vector_store",
                 created_at=0,  # Leave blank to have Postgres generate a timestamp
+                last_active_at=current_time,
                 file_counts=FileCounts(
                     cancelled=0, completed=0, failed=0, in_progress=0, total=0
                 ),
-                last_active_at=last_active_at,  # Set to current time
+                usage_bytes=0,  # Leave blank to have Postgres calculate the document bytes
                 metadata=request.metadata,
-                name=request.name or "",
-                object="vector_store",
-                status=VectorStoreStatus.IN_PROGRESS.value,
                 expires_after=expires_after,
                 expires_at=expires_at,
             )
-            new_vector_store = await crud_vector_store.create(object_=vector_store)
 
-            if request.file_ids != []:
-                responses = await self.index_files(
-                    new_vector_store.id, request.file_ids
-                )
+            # Save the placeholder to the database
+            saved_placeholder = await crud_vector_store.create(
+                object_=placeholder_vector_store
+            )
 
-                for response in responses:
-                    await self._increment_vector_store_file_status(
-                        new_vector_store, response
-                    )
+            if saved_placeholder is None:
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail="Unable to create vector store",
+                )
 
-            new_vector_store.status = VectorStoreStatus.COMPLETED.value
+            # Split the files, convert the chunks into vectors, and insert them into the db
+            if background_tasks:
+                # Perform the indexing in the background
+                background_tasks.add_task(
+                    self._complete_vector_store_creation,
+                    saved_placeholder.id,
+                    request,
+                )
+            else:
+                await self._complete_vector_store_creation(
+                    saved_placeholder.id, request
+                )
 
-            return await crud_vector_store.update(
-                id_=new_vector_store.id,
-                object_=new_vector_store,
-            )
+            return saved_placeholder
         except Exception as exc:
+            logging.error(exc)
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail="Unable to parse vector store request",
             ) from exc
 
+    async def _complete_vector_store_creation(
+        self, vector_store_id: str, request: CreateVectorStoreRequest
+    ):
+        """Complete the vector store creation process in the background."""
+        crud_vector_store = CRUDVectorStore(db=self.db)
+        vector_store = await crud_vector_store.get(
+            filters=FilterVectorStore(id=vector_store_id)
+        )
+
+        if request.file_ids != []:
+            responses = await self.index_files(vector_store_id, request.file_ids)
+            for response in responses:
+                await self._increment_vector_store_file_status(vector_store, response)
+
+        vector_store.status = VectorStoreStatus.COMPLETED.value
+        vector_store.last_active_at = int(time.time())
+
+        await crud_vector_store.update(id_=vector_store_id, object_=vector_store)
+
     async def modify_existing_vector_store(
         self,
         vector_store_id: str,
         request: ModifyVectorStoreRequest,
+        background_tasks: BackgroundTasks | None = None,
     ) -> VectorStore:
         """Modify an existing vector store given its id."""
         crud_vector_store = CRUDVectorStore(db=self.db)
@@ -236,43 +267,65 @@ async def modify_existing_vector_store(
                 expires_at=old_vector_store.expires_at,
             )
 
-            await crud_vector_store.update(
+            # Update the vector store with the new information and set status to in_progress for the duration of this function
+            updated_vector_store = await crud_vector_store.update(
                 id_=vector_store_id,
                 object_=new_vector_store,
-            )  # Sets status to in_progress for the duration of this function
+            )
 
-            if request.file_ids:
-                responses = await self.index_files(
-                    new_vector_store.id, request.file_ids
+            if updated_vector_store is None:
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail="Unable to modify vector store",
                 )
-                for response in responses:
-                    await self._increment_vector_store_file_status(
-                        new_vector_store, response
-                    )
-
-            new_vector_store.status = VectorStoreStatus.COMPLETED.value
 
-            last_active_at = int(time.time())
-            new_vector_store.last_active_at = (
-                last_active_at  # Update after indexing files
-            )
-            expires_after, expires_at = request.get_expiry(last_active_at)
-
-            if expires_at and expires_at:
-                new_vector_store.expires_after = expires_after
-                new_vector_store.expires_at = expires_at
+            # Split the files, convert the chunks into vectors, and insert them into the db
+            if request.file_ids:
+                if background_tasks:
+                    # Perform the indexing in the background
+                    background_tasks.add_task(
+                        self._complete_vector_store_modification,
+                        vector_store_id,
+                        request,
+                    )
+                else:
+                    await self._complete_vector_store_modification(
+                        vector_store_id, request
+                    )
 
-            return await crud_vector_store.update(
-                id_=vector_store_id,
-                object_=new_vector_store,
-            )
+            return updated_vector_store
         except Exception as exc:
             logging.error(exc)
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail="Unable to parse vector store request",
             ) from exc
 
+    async def _complete_vector_store_modification(
+        self, vector_store_id: str, request: ModifyVectorStoreRequest
+    ):
+        """Complete the vector store modification process in the background."""
+        crud_vector_store = CRUDVectorStore(db=self.db)
+        vector_store = await crud_vector_store.get(
+            filters=FilterVectorStore(id=vector_store_id)
+        )
+
+        if request.file_ids:
+            responses = await self.index_files(vector_store_id, request.file_ids)
+            for response in responses:
+                await self._increment_vector_store_file_status(vector_store, response)
+
+        vector_store.status = VectorStoreStatus.COMPLETED.value
+        last_active_at = int(time.time())
+        vector_store.last_active_at = last_active_at  # Update after indexing files
+
+        expires_after, expires_at = request.get_expiry(last_active_at)
+        if expires_after and expires_at:
+            vector_store.expires_after = expires_after
+            vector_store.expires_at = expires_at
+
+        await crud_vector_store.update(id_=vector_store_id, object_=vector_store)
+
     async def file_ids_are_valid(self, file_ids: str | list[str]) -> bool:
         """Check if the provided file ids exist"""
         crud_file_object = CRUDFileObject(db=self.db)