65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
from flask import jsonify, request
|
|
|
|
from . import bp
|
|
from llm_server.routes.stats import concurrent_semaphore, proompts
|
|
from ..helpers.http import validate_json
|
|
from ... import opts
|
|
from ...database import log_prompt
|
|
|
|
if opts.mode == 'oobabooga':
|
|
from ...llm.oobabooga.generate import generate
|
|
|
|
generator = generate
|
|
elif opts.mode == 'hf-textgen':
|
|
from ...llm.hf_textgen.generate import generate
|
|
|
|
generator = generate
|
|
|
|
|
|
@bp.route('/generate', methods=['POST'])
|
|
def generate():
|
|
request_valid_json, request_json_body = validate_json(request.data)
|
|
if not request_valid_json:
|
|
return jsonify({'code': 400, 'error': 'Invalid JSON'}), 400
|
|
|
|
with concurrent_semaphore:
|
|
success, response, error_msg = generator(request_json_body)
|
|
if not success:
|
|
return jsonify({
|
|
'code': 500,
|
|
'error': 'failed to reach backend'
|
|
}), 500
|
|
response_valid_json, response_json_body = validate_json(response)
|
|
if response_valid_json:
|
|
proompts.increment()
|
|
|
|
# request.headers = {"host": "proxy.chub-archive.evulid.cc", "x-forwarded-proto": "https", "user-agent": "node-fetch/1.0 (+https://github.com/bitinn/node-fetch)", "cf-visitor": {"scheme": "https"}, "cf-ipcountry": "CH", "accept": "*/*", "accept-encoding": "gzip",
|
|
# "x-forwarded-for": "193.32.127.228", "cf-ray": "7fa72c6a6d5cbba7-FRA", "cf-connecting-ip": "193.32.127.228", "cdn-loop": "cloudflare", "content-type": "application/json", "content-length": "9039"}
|
|
|
|
if request.headers.get('cf-connecting-ip'):
|
|
client_ip = request.headers.get('cf-connecting-ip')
|
|
elif request.headers.get('x-forwarded-for'):
|
|
client_ip = request.headers.get('x-forwarded-for')
|
|
else:
|
|
client_ip = request.remote_addr
|
|
|
|
parameters = request_json_body.copy()
|
|
del parameters['prompt']
|
|
|
|
token = request.headers.get('X-Api-Key')
|
|
|
|
log_prompt(opts.database_path, client_ip, token, request_json_body['prompt'], response_json_body['results'][0]['text'], parameters, dict(request.headers))
|
|
return jsonify({
|
|
**response_json_body
|
|
}), 200
|
|
else:
|
|
return jsonify({
|
|
'code': 500,
|
|
'error': 'failed to reach backend'
|
|
}), 500
|
|
|
|
# @openai_bp.route('/chat/completions', methods=['POST'])
|
|
# def generate_openai():
|
|
# print(request.data)
|
|
# return '', 200
|