feat(api): assistants endpoint (#424)

* Adds assistants endpoints for CRUD operations and DB migration
* Adds tests for assistants and files (requires Supabase)
* Adds Supabase config to API
---------

Co-authored-by: gharvey <gato.harvey@defenseunicorns.com>

.gitignore

@@ -23,6 +23,7 @@ build/
 .ruff_cache
 .branches
 .temp
+src/leapfrogai_api/config.yaml
 
 # local model and tokenizer files
 *.bin

packages/api/chart/templates/api/deployment.yaml

@@ -45,6 +45,14 @@ spec:
             value: "*.toml"
           - name: PORT
             value: "{{ .Values.api.port }}"
+          - name: SUPABASE_URL
+            value: "{{ .Values.supabase.url }}"
+          - name: SUPABASE_SERVICE_KEY
+            valueFrom:
+              secretKeyRef:
+                name: supabase-bootstrap-jwt
+                key: service-key
+                optional: true
           ports:
             - containerPort: 8080
           livenessProbe:

packages/api/chart/values.yaml

@@ -4,6 +4,9 @@ image:
   # x-release-please-end
   kiwigridTag: 1.23.3
 
+supabase:
+  url: "https://supabase-kong.###ZARF_VAR_HOSTED_DOMAIN###"
+
 api:
   replicas: 1
   port: 8080

packages/api/zarf.yaml

@@ -18,6 +18,8 @@ variables:
   - name: EXPOSE_OPENAPI_SCHEMA
     default: "false"
     description: "Flag to expose the OpenAPI schema for debugging."
+  - name: HOSTED_DOMAIN
+    default: "uds.dev"
 
 components:
   - name: leapfrogai

src/leapfrogai_api/Makefile

@@ -4,3 +4,6 @@ install:
 dev:
 	make install
 	python -m uvicorn main:app --port 3000 --reload
+
+test-integration:
+	cd ../../ && python -m pytest tests/integration/api

src/leapfrogai_api/README.md

@@ -23,4 +23,21 @@ Setup environment variables:
 ``` bash
 export SUPABASE_URL="http://localhost:54321" # or whatever you configured it as in your Supabase config.toml
 export SUPABASE_SERVICE_KEY="<YOUR_KEY>" # supabase status will show you the keys
-```
+```
+
+## Integration Tests
+
+The integration tests serve to identify any mismatches between components:
+
+- Check all API routes
+- Validate Request/Response types
+- DB CRUD operations
+- Schema mismatches
+
+Integration tests require a Supabase instance and environment variables configured (see [Local Development](#local-development)).
+
+From this directory run:
+
+``` bash
+make test-integration
+```

src/leapfrogai_api/backend/types.py

@@ -4,6 +4,9 @@
 from typing import Literal
 from pydantic import BaseModel
 from fastapi import UploadFile, Form, File
+from openai.types import FileObject
+from openai.types.beta import Assistant, AssistantTool
+from openai.types.beta.assistant import ToolResources
 
 ##########
 # GENERIC
@@ -229,6 +232,13 @@ def as_form(
         return cls(file=file, purpose=purpose)
 
 
+class ListFilesResponse(BaseModel):
+    """Response object for listing files."""
+
+    object: str = Literal["list"]
+    data: list[FileObject] = []
+
+
 #############
 # ASSISTANTS
 #############
@@ -241,10 +251,8 @@ class CreateAssistantRequest(BaseModel):
     name: str | None = "Froggy Assistant"
     description: str | None = "A helpful assistant."
     instructions: str | None = "You are a helpful assistant."
-    tools: list[dict[Literal["type"], Literal["file_search"]]] | None = [
-        {"type": "file_search"}
-    ]  # This is all we support right now
-    tool_resources: object | None = {}
+    tools: list[AssistantTool] | None = []  # This is all we support right now
+    tool_resources: ToolResources | None = ToolResources()
     metadata: object | None = {}
     temperature: float | None = 1.0
     top_p: float | None = 1.0
@@ -253,3 +261,10 @@ class CreateAssistantRequest(BaseModel):
 
 class ModifyAssistantRequest(CreateAssistantRequest):
     """Request object for modifying an assistant."""
+
+
+class ListAssistantsResponse(BaseModel):
+    """Response object for listing files."""
+
+    object: str = Literal["list"]
+    data: list[Assistant] = []

src/leapfrogai_api/data/crud_assistant_object.py

@@ -0,0 +1,34 @@
+"""CRUD Operations for Assistant."""
+
+from openai.types.beta import Assistant
+from supabase_py_async import AsyncClient
+from leapfrogai_api.data.crud_base import CRUDBase
+
+
+class CRUDAssistant(CRUDBase[Assistant]):
+    """CRUD Operations for Assistant"""
+
+    def __init__(self, model: type[Assistant], table_name: str = "assistant_objects"):
+        super().__init__(model=model, table_name=table_name)
+
+    async def create(self, db: AsyncClient, object_: Assistant) -> Assistant | None:
+        """Create a new assistant."""
+        return await super().create(db=db, object_=object_)
+
+    async def get(self, id_: str, db: AsyncClient) -> Assistant | None:
+        """Get an assistant by its ID."""
+        return await super().get(db=db, id_=id_)
+
+    async def list(self, db: AsyncClient) -> list[Assistant] | None:
+        """List all assistants."""
+        return await super().list(db=db)
+
+    async def update(
+        self, id_: str, db: AsyncClient, object_: Assistant
+    ) -> Assistant | None:
+        """Update an assistant by its ID."""
+        return await super().update(id_=id_, db=db, object_=object_)
+
+    async def delete(self, id_: str, db: AsyncClient) -> bool:
+        """Delete an assistant by its ID."""
+        return await super().delete(id_=id_, db=db)

src/leapfrogai_api/data/crud_base.py

@@ -0,0 +1,75 @@
+"""CRUD Operations for VectorStore."""
+
+from typing import Generic, TypeVar
+from supabase_py_async import AsyncClient
+from pydantic import BaseModel
+
+ModelType = TypeVar("ModelType", bound=BaseModel)
+
+
+class CRUDBase(Generic[ModelType]):
+    """CRUD Operations"""
+
+    def __init__(self, model: type[ModelType], table_name: str):
+        self.model = model
+        self.table_name = table_name
+
+    async def create(self, db: AsyncClient, object_: ModelType) -> ModelType | None:
+        """Create new row."""
+        dict_ = object_.model_dump()
+        del dict_["id"]  # Ensure this is created by the database
+        del dict_["created_at"]  # Ensure this is created by the database
+        data, _count = await db.table(self.table_name).insert(dict_).execute()
+
+        _, response = data
+
+        if response:
+            return self.model(**response[0])
+        return None
+
+    async def get(self, id_: str, db: AsyncClient) -> ModelType | None:
+        """Get row by ID."""
+        data, _count = (
+            await db.table(self.table_name).select("*").eq("id", id_).execute()
+        )
+
+        _, response = data
+
+        if response:
+            return self.model(**response[0])
+        return None
+
+    async def list(self, db: AsyncClient) -> list[ModelType] | None:
+        """List all rows."""
+        data, _count = await db.table(self.table_name).select("*").execute()
+
+        _, response = data
+
+        if response:
+            return [self.model(**item) for item in response]
+        return None
+
+    async def update(
+        self, id_: str, db: AsyncClient, object_: ModelType
+    ) -> ModelType | None:
+        """Update a vector store by its ID."""
+        data, _count = (
+            await db.table(self.table_name)
+            .update(object_.model_dump())
+            .eq("id", id_)
+            .execute()
+        )
+
+        _, response = data
+
+        if response:
+            return self.model(**response[0])
+        return None
+
+    async def delete(self, id_: str, db: AsyncClient) -> bool:
+        """Delete a vector store by its ID."""
+        data, _count = await db.table(self.table_name).delete().eq("id", id_).execute()
+
+        _, response = data
+
+        return bool(response)

src/leapfrogai_api/data/crud_file_object.py

@@ -1,77 +1,34 @@
 """CRUD Operations for FileObject"""
 
+from openai.types import FileObject
 from supabase_py_async import AsyncClient
-from openai.types import FileObject, FileDeleted
+from leapfrogai_api.data.crud_base import CRUDBase
 
 
-class CRUDFileObject:
+class CRUDFileObject(CRUDBase[FileObject]):
     """CRUD Operations for FileObject"""
 
-    def __init__(self, model: type[FileObject]):
-        self.model = model
+    def __init__(self, model: type[FileObject], table_name: str = "file_objects"):
+        super().__init__(model=model, table_name=table_name)
 
-    async def create(
-        self, client: AsyncClient, file_object: FileObject
-    ) -> FileObject | None:
+    async def create(self, db: AsyncClient, object_: FileObject) -> FileObject | None:
         """Create a new file object."""
-        file_object_dict = file_object.model_dump()
-        if file_object_dict.get("id") == "":
-            del file_object_dict["id"]
-        data, _count = (
-            await client.table("file_objects").insert(file_object_dict).execute()
-        )
-
-        _, response = data
-
-        if response:
-            return self.model(**response[0])
-        return None
+        return await super().create(db=db, object_=object_)
 
-    async def get(self, client: AsyncClient, file_id: str) -> FileObject | None:
+    async def get(self, id_: str, db: AsyncClient) -> FileObject | None:
         """Get a file object by its ID."""
-        data, _count = (
-            await client.table("file_objects").select("*").eq("id", file_id).execute()
-        )
+        return await super().get(db=db, id_=id_)
 
-        _, response = data
-
-        if response:
-            return self.model(**response[0])
-        return None
-
-    async def list(self, client: AsyncClient) -> list[FileObject] | None:
+    async def list(self, db: AsyncClient) -> list[FileObject] | None:
         """List all file objects."""
-        data, _count = await client.table("file_objects").select("*").execute()
-
-        _, response = data
-
-        if response:
-            return [self.model(**item) for item in response]
-        return None
+        return await super().list(db=db)
 
     async def update(
-        self, client: AsyncClient, file_id: str, file_object: FileObject
+        self, id_: str, db: AsyncClient, object_: FileObject
     ) -> FileObject | None:
         """Update a file object by its ID."""
-        data, _count = (
-            await client.table("file_objects")
-            .update(file_object.model_dump())
-            .eq("id", file_id)
-            .execute()
-        )
-
-        _, response = data
-
-        if response:
-            return self.model(**response[0])
-        return None
+        return await super().update(id_=id_, db=db, object_=object_)
 
-    async def delete(self, client: AsyncClient, file_id: str) -> FileDeleted:
+    async def delete(self, id_: str, db: AsyncClient) -> bool:
         """Delete a file object by its ID."""
-        data, _count = (
-            await client.table("file_objects").delete().eq("id", file_id).execute()
-        )
-
-        _, response = data
-
-        return FileDeleted(id=file_id, deleted=bool(response), object="file")
+        return await super().delete(id_=id_, db=db)