b625bface9
Models can be loaded and unloaded via the API. Also add authentication to use the API and for administrator tasks. Both types of authorization use different keys. Also fix the unload function to properly free all used vram. Signed-off-by: kingbri <bdashore3@proton.me>
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
import secrets
|
|
import yaml
|
|
from fastapi import Header, HTTPException
|
|
from typing import Optional
|
|
|
|
"""
|
|
This method of authorization is pretty insecure, but since TabbyAPI is a local
|
|
application, it should be fine.
|
|
"""
|
|
|
|
class AuthKeys:
|
|
api_key: str
|
|
admin_key: str
|
|
|
|
def __init__(self, api_key: str, admin_key: str):
|
|
self.api_key = api_key
|
|
self.admin_key = admin_key
|
|
|
|
auth_keys: Optional[AuthKeys] = None
|
|
|
|
def load_auth_keys():
|
|
global auth_keys
|
|
try:
|
|
with open("api_tokens.yml", "r") as auth_file:
|
|
auth_keys = yaml.safe_load(auth_file)
|
|
except:
|
|
new_auth_keys = AuthKeys(
|
|
api_key = secrets.token_hex(16),
|
|
admin_key = secrets.token_hex(16)
|
|
)
|
|
auth_keys = new_auth_keys
|
|
|
|
with open("api_tokens.yml", "w") as auth_file:
|
|
yaml.dump(auth_keys, auth_file)
|
|
|
|
def check_api_key(x_api_key: str = Header(None), authorization: str = Header(None)):
|
|
if x_api_key and x_api_key == auth_keys.api_key:
|
|
return x_api_key
|
|
elif authorization:
|
|
split_key = authorization.split(" ")
|
|
if split_key[0].lower() == "bearer" and split_key[1] == auth_keys.api_key:
|
|
return authorization
|
|
else:
|
|
raise HTTPException(401, "Invalid API key")
|
|
|
|
def check_admin_key(x_admin_key: str = Header(None), authorization: str = Header(None)):
|
|
if x_admin_key and x_admin_key == auth_keys.admin_key:
|
|
return x_admin_key
|
|
elif authorization:
|
|
split_key = authorization.split(" ")
|
|
if split_key[0].lower() == "bearer" and split_key[1] == auth_keys.admin_key:
|
|
return authorization
|
|
else:
|
|
raise HTTPException(401, "Invalid admin key")
|