2023-08-21 21:28:52 -06:00
|
|
|
import json
|
2023-09-12 16:40:09 -06:00
|
|
|
from functools import wraps
|
2023-08-21 21:28:52 -06:00
|
|
|
from typing import Union
|
|
|
|
|
2023-09-12 16:40:09 -06:00
|
|
|
import flask
|
|
|
|
import requests
|
|
|
|
from flask import make_response, Request
|
2023-08-21 21:28:52 -06:00
|
|
|
from flask import request, jsonify
|
|
|
|
|
|
|
|
from llm_server import opts
|
2023-09-20 20:30:31 -06:00
|
|
|
from llm_server.database.database import is_valid_api_key
|
2023-08-21 21:28:52 -06:00
|
|
|
|
|
|
|
|
2023-08-21 22:49:44 -06:00
|
|
|
def cache_control(seconds):
|
|
|
|
def decorator(f):
|
|
|
|
@wraps(f)
|
|
|
|
def decorated_function(*args, **kwargs):
|
|
|
|
resp = make_response(f(*args, **kwargs))
|
2023-08-23 12:40:13 -06:00
|
|
|
if seconds > 0:
|
2023-08-21 22:49:44 -06:00
|
|
|
resp.headers['Cache-Control'] = f'public, max-age={seconds}'
|
|
|
|
else:
|
|
|
|
resp.headers['Cache-Control'] = f'no-store'
|
|
|
|
return resp
|
|
|
|
|
|
|
|
return decorated_function
|
|
|
|
|
|
|
|
return decorator
|
|
|
|
|
|
|
|
|
2023-08-21 21:28:52 -06:00
|
|
|
def require_api_key():
|
|
|
|
if not opts.auth_required:
|
|
|
|
return
|
|
|
|
elif 'X-Api-Key' in request.headers:
|
|
|
|
if is_valid_api_key(request.headers['X-Api-Key']):
|
|
|
|
return
|
|
|
|
else:
|
|
|
|
return jsonify({'code': 403, 'message': 'Invalid API key'}), 403
|
|
|
|
else:
|
|
|
|
return jsonify({'code': 401, 'message': 'API key required'}), 401
|
|
|
|
|
|
|
|
|
2023-09-23 17:57:23 -06:00
|
|
|
def validate_json(data: Union[str, flask.Request, requests.models.Response, flask.Response, dict]):
|
|
|
|
if isinstance(data, dict):
|
|
|
|
return True, data
|
2023-09-12 16:40:09 -06:00
|
|
|
try:
|
|
|
|
if isinstance(data, (Request, flask.Response)):
|
|
|
|
data = data.json
|
|
|
|
return True, data
|
|
|
|
elif isinstance(data, requests.models.Response):
|
2023-08-21 21:28:52 -06:00
|
|
|
data = data.json()
|
|
|
|
return True, data
|
2023-09-12 16:40:09 -06:00
|
|
|
except Exception as e:
|
|
|
|
return False, e
|
2023-08-21 21:28:52 -06:00
|
|
|
try:
|
2023-09-12 16:40:09 -06:00
|
|
|
j = json.loads(str(data))
|
2023-08-21 21:28:52 -06:00
|
|
|
return True, j
|
|
|
|
except Exception as e:
|
2023-09-12 16:40:09 -06:00
|
|
|
return False, e
|