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

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;

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"]

tests/e2e/test_supabase.py

@@ -1,6 +1,24 @@
+import asyncio
+import io
+import threading
+import uuid
+from fastapi import UploadFile
 import requests
+from openai.types.beta.vector_stores import VectorStoreFile
+from openai.types.beta import VectorStore
+from openai.types.beta.vector_store import FileCounts
+import _thread
 
-from .utils import ANON_KEY
+from supabase import AClient as AsyncClient, acreate_client
+from realtime import Socket
+from leapfrogai_api.data.crud_file_bucket import CRUDFileBucket
+from leapfrogai_api.data.crud_file_object import CRUDFileObject
+from leapfrogai_api.data.crud_vector_store import CRUDVectorStore
+
+from leapfrogai_api.data.crud_vector_store_file import CRUDVectorStoreFile
+
+from .utils import ANON_KEY, create_test_user, SERVICE_KEY
+from openai.types import FileObject
 
 health_urls = {
     "auth_health_url": "http://supabase-kong.uds.dev/auth/v1/health",
@@ -12,9 +30,151 @@
 def test_studio():
     try:
         for url_name in health_urls:
-            response = requests.get(health_urls[url_name], headers={"apikey": ANON_KEY})
-            response.raise_for_status()
+            resp = requests.get(health_urls[url_name], headers={"apikey": ANON_KEY})
+            resp.raise_for_status()
     except requests.exceptions.RequestException as e:
-        print(f"Error: Request failed with status code {response.status_code}")
+        print(f"Error: Request failed with status code {resp.status_code}")
         print(e)
         exit(1)
+
+
+def test_supabase_realtime_vector_store_indexing():
+    class TestCompleteException(Exception):
+        pass
+
+    def timeout_handler():
+        print("Test timed out after 10 seconds")
+        # This is necessary to stop the thread from hanging forever
+        _thread.interrupt_main()
+
+    async def postgres_db_changes():
+        """
+        This function is responsible for creating a vector store and uploading a file to it.
+        """
+        client: AsyncClient = await acreate_client(
+            supabase_key=ANON_KEY,
+            supabase_url="https://supabase-kong.uds.dev",
+        )
+        await client.auth.set_session(access_token=access_token, refresh_token="dummy")
+
+        upload_file_id = await upload_file(client)
+        assert upload_file_id is not None, "Failed to upload file"
+
+        vector_store = VectorStore(
+            id=str(uuid.uuid4()),
+            created_at=0,
+            file_counts=FileCounts(
+                cancelled=0,
+                completed=0,
+                failed=0,
+                in_progress=0,
+                total=0,
+            ),
+            name="test_vector_store",
+            object="vector_store",
+            status="completed",
+            usage_bytes=0,
+        )
+
+        await CRUDVectorStore(client).create(vector_store)
+
+        vector_store_file = VectorStoreFile(
+            id=upload_file_id,
+            vector_store_id=vector_store.id,
+            created_at=0,
+            object="vector_store.file",
+            status="completed",
+            usage_bytes=0,
+        )
+
+        await CRUDVectorStoreFile(client).create(vector_store_file)
+
+    def postgres_changes_callback(payload):
+        """
+        This function is responsible for listening for changes to the vector store file and signaling success if the file triggers realtime successfully.
+        """
+        expected_record = {
+            "object": "vector_store.file",
+            "status": "completed",
+            "usage_bytes": 0,
+        }
+
+        all_records_match = all(
+            payload.get("record", {}).get(key) == value
+            for key, value in expected_record.items()
+        )
+        event_information_match = (
+            payload.get("table") == "vector_store_file"
+            and payload.get("type") == "INSERT"
+        )
+
+        if event_information_match and all_records_match:
+            raise TestCompleteException("Test completed successfully")
+
+    async def upload_file(client: AsyncClient) -> str:
+        """
+        This function is responsible for uploading a file to the file bucket.
+        """
+        id_ = str(uuid.uuid4())
+
+        empty_file_object = FileObject(
+            id=id_,
+            bytes=0,
+            created_at=0,
+            filename="",
+            object="file",
+            purpose="assistants",
+            status="uploaded",
+            status_details=None,
+        )
+
+        crud_file_object = CRUDFileObject(client)
+
+        file_object = await crud_file_object.create(object_=empty_file_object)
+        assert file_object is not None, "Failed to create file object"
+
+        crud_file_bucket = CRUDFileBucket(db=client, model=UploadFile)
+        await crud_file_bucket.upload(
+            file=UploadFile(filename="", file=io.BytesIO(b"")), id_=file_object.id
+        )
+        return id_
+
+    def run_postgres_db_changes():
+        """
+        This function is responsible for running the postgres_db_changes function.
+        """
+        asyncio.run(postgres_db_changes())
+
+    timeout_timer = None
+    try:
+        random_name = str(uuid.uuid4())
+        access_token = create_test_user(email=f"{random_name}@fake.com")
+
+        # Schedule postgres_db_changes to run after 5 seconds
+        threading.Timer(5.0, run_postgres_db_changes).start()
+
+        # Set a timeout of 10 seconds
+        timeout_timer = threading.Timer(10.0, timeout_handler)
+        timeout_timer.start()
+
+        # Listening socket
+        # The service key is needed for proper permission to listen to realtime events
+        # At the time of writing this, the Supabase realtime library does not support RLS
+        URL = f"wss://supabase-kong.uds.dev/realtime/v1/websocket?apikey={SERVICE_KEY}&vsn=1.0.0"
+        s = Socket(URL)
+        s.connect()
+
+        # Set channel to listen for changes to the vector_store_file table
+        channel_1 = s.set_channel("realtime:public:vector_store_file")
+        # Listen for all events on the channel ex: INSERT, UPDATE, DELETE
+        channel_1.join().on("*", postgres_changes_callback)
+
+        # Start listening
+        s.listen()
+    except TestCompleteException:
+        if timeout_timer is not None:
+            timeout_timer.cancel()  # Cancel the timeout timer if test completes successfully
+
+        assert True
+    except Exception:
+        assert False

tests/e2e/utils.py

@@ -6,6 +6,7 @@
 
 # This is the anon_key for supabase, it provides access to the endpoints that would otherwise be inaccessible
 ANON_KEY = os.environ["ANON_KEY"]
+SERVICE_KEY = os.environ["SERVICE_KEY"]
 
 DEFAULT_TEST_EMAIL = "[email protected]"
 DEFAULT_TEST_PASSWORD = "password"