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

import traceback

from flask import jsonify, request

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


@bp.route('/v1/generate', methods=['POST'])
@bp.route('/<model_name>/v1/generate', methods=['POST'])
def generate(model_name=None):
    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:
        handler = OobaRequestHandler(request, selected_model=model_name)
        try:
            return handler.handle_request()
        except Exception:
            traceback.print_exc()
            return handler.handle_error('Server encountered exception.', 'exception')[0], 500