feat(api): long lived api keys

packages/api/supabase/migrations/20240618163044_api_keys.sql

@@ -0,0 +1,144 @@
+-- Initialize api_keys table
+create table api_keys (
+    name text,
+    id uuid primary key default uuid_generate_v4(),
+    user_id uuid references auth.users not null,
+    api_key text not null unique,
+    created_at bigint default extract(epoch from now()) not null,
+    expires_at bigint
+);
+
+--- RLS for api_keys table
+
+alter table api_keys enable row level security;
+
+create policy "Read only if API key matches and is current" ON api_keys for
+    select using (
+        api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+        and (expires_at is null or expires_at > extract(epoch from now()))
+    );
+
+create policy "Individuals can view their own api_keys." on api_keys for
+    select using (auth.uid() = user_id);
+create policy "Individuals can create api_keys." on api_keys for
+    insert with check (auth.uid() = user_id);
+create policy "Individuals can delete their own api_keys." on api_keys for
+    delete using (auth.uid() = user_id);
+
+-- API Key Policies
+
+create policy "Individuals can CRUD their own assistant_objects via API key."
+    on assistant_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = assistant_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own thread_objects via API key."
+    on thread_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = thread_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own message_objects via API key."
+    on message_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = message_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own file_objects via API key."
+    on file_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = file_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD file_bucket via API key."
+    on storage.buckets for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+        )
+    );
+
+create policy "Individuals can CRUD their own run_objects via API key."
+    on run_objects for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = run_objects.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own vector_store via API key."
+    on vector_store for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = vector_store.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own vector_store_file via API key."
+    on vector_store_file for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = vector_store_file.user_id
+        )
+    );
+
+create policy "Individuals can CRUD their own vector_content via API key."
+    on vector_content for all
+    to anon
+    using
+    (
+        exists (
+            select 1
+            from api_keys
+            where api_keys.api_key = current_setting('request.headers')::json->>'x-custom-api-key'
+            and api_keys.user_id = vector_content.user_id
+        )
+    );

packages/api/zarf.yaml

@@ -17,7 +17,7 @@ constants:
 variables:
   - name: EXPOSE_OPENAPI_SCHEMA
     default: "false"
-    description: "Flag to expose the OpenAPI schema for debugging."
+    description: "Flag to expose the API service"
   - name: HOSTED_DOMAIN
     default: "uds.dev"
   - name: DEFAULT_EMBEDDINGS_MODEL

src/leapfrogai_api/backend/rag/index.py

@@ -8,7 +8,7 @@
 from langchain_core.embeddings import Embeddings
 from openai.types.beta.vector_stores import VectorStoreFile
 from openai.types.beta.vector_stores.vector_store_file import LastError
-from supabase_py_async import AsyncClient
+from supabase import AClient as AsyncClient
 from leapfrogai_api.backend.rag.document_loader import load_file, split
 from leapfrogai_api.backend.rag.leapfrogai_embeddings import LeapfrogAIEmbeddings
 from leapfrogai_api.backend.types import VectorStoreFileStatus
@@ -184,7 +184,7 @@ async def aadd_documents(
                 metadata=document.metadata,
                 embedding=embedding,
             )
-            ids.append(response.data[0]["id"])
+            ids.append(response[0]["id"])
 
         return ids
 
@@ -273,5 +273,12 @@ async def _aadd_vector(
             "metadata": metadata,
             "embedding": embedding,
         }
-        response = await self.db.from_(self.table_name).insert(row).execute()
+        data, _count = await self.db.from_(self.table_name).insert(row).execute()
+
+        _, response = data
+
+        for item in response:
+            if "user_id" in item:
+                del item["user_id"]
+
         return response

src/leapfrogai_api/backend/rag/query.py

@@ -1,6 +1,6 @@
 """Service for querying the RAG model."""
 
-from supabase_py_async import AsyncClient
+from supabase import AClient as AsyncClient
 from leapfrogai_api.backend.rag.index import IndexingService
 from postgrest.base_request_builder import SingleAPIResponse
 

src/leapfrogai_api/backend/security/api_key.py

@@ -0,0 +1,100 @@
+"""Basic API key security functions."""
+
+import secrets
+import hashlib
+
+KEY_PREFIX = "lfai"
+KEY_BYTES = 32
+CHECKSUM_LENGTH = 4
+
+
+def generate_api_key() -> tuple[str, str]:
+    """Generate an API key.
+
+    returns:
+        read_once_token: str - in the format: {prefix}_{unique_key}_{checksum}
+        hashed_token: str - in the format: {prefix}_{hashed_key}_{checksum}
+    """
+    unique_key = secrets.token_bytes(KEY_BYTES).hex()
+    hashed_token = _encode_unique_key(unique_key)
+    checksum = parse_api_key(hashed_token)[2]
+
+    read_once_token = f"{KEY_PREFIX}_{unique_key}_{checksum}"
+    hashed_token = _encode_unique_key(unique_key)
+
+    return read_once_token, hashed_token
+
+
+def validate_and_encode_api_key(api_key: str) -> tuple[bool, str]:
+    """
+    Validate and encode an API key.
+
+    Should be in the form: `{prefix}_{unique_key}_{checksum}`
+
+    returns:
+        valid: bool
+        encoded_key: str
+    """
+
+    valid = validate_api_key(api_key)
+    encoded_key = ""
+
+    if valid:
+        encoded_key = _encode_unique_key(parse_api_key(api_key)[1])
+
+    return valid, encoded_key
+
+
+def _encode_unique_key(unique_key: str) -> str:
+    """Hashes and encodes an API key as a string.
+
+    returns:
+        api_key: str # in the format {prefix}_{one_way_hash}_{checksum}
+    """
+    one_way_hash = hashlib.sha256(unique_key.encode()).hexdigest()
+    checksum = hashlib.sha256(unique_key.encode()).hexdigest()[:CHECKSUM_LENGTH]
+    return f"{KEY_PREFIX}_{one_way_hash}_{checksum}"
+
+
+def _validate_checksum(unique_key: str, checksum: str):
+    """Validate the checksum of an API key."""
+    return hashlib.sha256(unique_key.encode()).hexdigest()[:CHECKSUM_LENGTH] == checksum
+
+
+def validate_api_key(api_key: str) -> bool:
+    """Validate an API key.
+
+    returns:
+        bool: True if the key is valid, False otherwise
+    """
+    try:
+        _prefix, key, checksum = parse_api_key(api_key)
+    except ValueError:
+        return False
+
+    return _validate_checksum(key, checksum)
+
+
+def parse_api_key(api_key: str) -> tuple[str, str, str]:
+    """Parse an API key into its components.
+
+    key format: {prefix}_{unique_key}_{checksum}
+
+    Note: Does not validate prefix, checksum, or key.
+
+    returns:
+        prefix: str
+        key: str
+        checksum: str
+    """
+
+    prefix, key, checksum = (
+        api_key.split("_")[0],
+        "_".join(api_key.split("_")[1:-1]),
+        api_key.split("_")[-1],
+    )
+
+    if not prefix or not key or not checksum:
+        raise ValueError("Invalid API key format")
+
+    return prefix, key, checksum

src/leapfrogai_api/data/crud_assistant.py

@@ -1,39 +1,26 @@
 """CRUD Operations for Assistant."""
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 from openai.types.beta import Assistant
-from supabase_py_async import AsyncClient
+from supabase import AClient as AsyncClient
 from leapfrogai_api.data.crud_base import CRUDBase
 
 
-class AuthAssistant(Assistant):
-    """A wrapper for the Assistant that includes a user_id for auth"""
-
-    user_id: str = Field(default="")
-
-
 class FilterAssistant(BaseModel):
     """Validation for Assistant filter."""
 
     id: str
 
 
-class CRUDAssistant(CRUDBase[AuthAssistant]):
+class CRUDAssistant(CRUDBase[Assistant]):
     """CRUD Operations for Assistant"""
 
     def __init__(
         self,
         db: AsyncClient,
         table_name: str = "assistant_objects",
     ):
-        super().__init__(db=db, model=AuthAssistant, table_name=table_name)
-
-    async def create(self, object_: Assistant) -> Assistant | None:
-        """Create a new assistant."""
-        user_id: str = (await self.db.auth.get_user()).user.id
-        return await super().create(
-            object_=AuthAssistant(user_id=user_id, **object_.model_dump())
-        )
+        super().__init__(db=db, model=Assistant, table_name=table_name)
 
     async def get(self, filters: FilterAssistant | None = None) -> Assistant | None:
         """Get assistant by filters."""
@@ -45,13 +32,6 @@ async def list(
         """List all assistants."""
         return await super().list(filters=filters.model_dump() if filters else None)
 
-    async def update(self, id_: str, object_: Assistant) -> Assistant | None:
-        """Update an assistant by its ID."""
-        user_id: str = (await self.db.auth.get_user()).user.id
-        return await super().update(
-            id_=id_, object_=AuthAssistant(user_id=user_id, **object_.model_dump())
-        )
-
     async def delete(self, filters: FilterAssistant | None = None) -> bool:
         """Delete an assistant by its ID."""
         return await super().delete(filters=filters.model_dump() if filters else None)