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

import json
from typing import Union

from flask import make_response
from requests import Response

from flask import request, jsonify
from functools import wraps

from llm_server import opts
from llm_server.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, Response]):
    if isinstance(data, Response):
        try:
            data = data.json()
            return True, data
        except Exception as e:
            return False, None
    try:
        j = json.loads(data)
        return True, j
    except Exception as e:
        return False, None