local-llm-server/llm_server/routes/v1/generate.py

import traceback

from flask import jsonify, request

from . import bp
from ..helpers.client import format_sillytavern_err
from ..helpers.http import validate_json
from ..ooba_request_handler import OobaRequestHandler


@bp.route('/generate', methods=['POST'])
def generate():
    disable_st_error_formatting = request.headers.get('LLM-ST-Errors', False) == 'true'
    request_valid_json, request_json_body = validate_json(request)
    if not request_valid_json or not request_json_body.get('prompt'):
        return jsonify({'code': 400, 'msg': 'Invalid JSON'}), 400
    else:
        try:
            return OobaRequestHandler(request).handle_request()
        except Exception as e:
            print(f'EXCEPTION on {request.url}!!!')
            print(traceback.format_exc())
            if disable_st_error_formatting:
                return '500', 500
            else:
                return format_sillytavern_err(f'Server encountered exception.', 'error'), 500