local-llm-server/llm_server/routes/helpers/http.py

import json
from functools import wraps
from typing import Union

import flask
import requests
from flask import make_response, Request
from flask import request, jsonify

from llm_server import opts
from llm_server.database.database import is_valid_api_key


def cache_control(seconds):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            resp = make_response(f(*args, **kwargs))
            if seconds > 0:
                resp.headers['Cache-Control'] = f'public, max-age={seconds}'
            else:
                resp.headers['Cache-Control'] = f'no-store'
            return resp

        return decorated_function

    return decorator


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


def validate_json(data: Union[str, flask.Request, requests.models.Response, flask.Response, dict]):
    if isinstance(data, dict):
        return True, data
    try:
        if isinstance(data, (Request, flask.Response)):
            data = data.json
            return True, data
        elif isinstance(data, requests.models.Response):
            data = data.json()
            return True, data
    except Exception as e:
        return False, e
    try:
        j = json.loads(str(data))
        return True, j
    except Exception as e:
        return False, e