This repository has been archived on 2024-10-27. You can view files and clone it, but cannot push or open issues or pull requests.
local-llm-server/llm_server/routes/helpers/http.py

103 lines
3.6 KiB
Python
Raw Normal View History

2023-09-26 22:49:53 -06:00
import simplejson as json
import traceback
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 Request, make_response
from flask import jsonify, request
2023-08-21 21:28:52 -06:00
from llm_server import opts
from llm_server.database.database import is_valid_api_key
2023-09-26 22:09:11 -06:00
from llm_server.routes.auth import parse_token
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
# TODO:
# File "/srv/server/local-llm-server/llm_server/routes/request_handler.py", line 240, in before_request
# response = require_api_key()
# ^^^^^^^^^^^^^^^^^
# File "/srv/server/local-llm-server/llm_server/routes/helpers/http.py", line 50, in require_api_key
# if token.startswith('SYSTEM__') or opts.auth_required:
# ^^^^^^^^^^^^^^^^
# AttributeError: 'NoneType' object has no attribute 'startswith'
2023-09-26 22:49:53 -06:00
def require_api_key(json_body: dict = None):
if json_body:
request_json = json_body
elif request.headers.get('Content-Type') == 'application/json':
valid_json, request_json = validate_json(request.data)
if not valid_json:
request_json = None
else:
request_json = None
2023-09-26 22:09:11 -06:00
if 'X-Api-Key' in request.headers:
api_key = request.headers['X-Api-Key']
if api_key.startswith('SYSTEM__') or opts.auth_required:
if is_valid_api_key(api_key):
return
else:
return jsonify({'code': 403, 'message': 'Invalid API key'}), 403
elif 'Authorization' in request.headers:
token = parse_token(request.headers['Authorization'])
if (token and token.startswith('SYSTEM__')) or opts.auth_required:
2023-09-26 22:09:11 -06:00
if is_valid_api_key(token):
return
else:
2023-09-26 22:49:53 -06:00
return jsonify({'code': 403, 'message': 'Invalid API key'}), 403
2023-08-21 21:28:52 -06:00
else:
try:
# Handle websockets
2023-09-26 22:49:53 -06:00
if opts.auth_required and not request_json:
# If we didn't get any valid JSON, deny.
return jsonify({'code': 403, 'message': 'Invalid API key'}), 403
if request_json and request_json.get('X-API-KEY'):
api_key = request_json.get('X-API-KEY')
2023-09-26 22:09:11 -06:00
if api_key.startswith('SYSTEM__') or opts.auth_required:
if is_valid_api_key(api_key):
return
else:
return jsonify({'code': 403, 'message': 'Invalid API key'}), 403
except:
# TODO: remove this one we're sure this works as expected
traceback.print_exc()
2023-08-21 21:28:52 -06:00
2023-09-26 22:49:53 -06:00
def validate_json(data: Union[str, flask.Request, requests.models.Response, flask.Response, dict, bytes]):
2023-09-23 17:57:23 -06:00
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-26 22:49:53 -06:00
elif isinstance(data, bytes):
s = data.decode('utf-8')
return False, json.loads(s)
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-10-01 22:58:08 -06:00
return False, e