94 lines
4.5 KiB
Python
94 lines
4.5 KiB
Python
import json
|
|
import traceback
|
|
from typing import Tuple
|
|
from uuid import uuid4
|
|
|
|
import flask
|
|
from flask import jsonify
|
|
|
|
from llm_server import opts
|
|
from llm_server.database.database import is_api_key_moderated, log_prompt
|
|
from llm_server.llm.openai.transform import build_openai_response, transform_messages_to_prompt, trim_prompt_to_fit
|
|
from llm_server.routes.helpers.client import format_sillytavern_err
|
|
from llm_server.routes.request_handler import RequestHandler
|
|
from llm_server.threads import add_moderation_task, get_results
|
|
|
|
|
|
class OpenAIRequestHandler(RequestHandler):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.prompt = None
|
|
|
|
def handle_request(self) -> Tuple[flask.Response, int]:
|
|
assert not self.used
|
|
|
|
if opts.openai_silent_trim:
|
|
oai_messages = trim_prompt_to_fit(self.request.json['messages'], opts.context_size)
|
|
else:
|
|
oai_messages = self.request.json['messages']
|
|
|
|
self.prompt = transform_messages_to_prompt(oai_messages)
|
|
request_valid, invalid_response = self.validate_request()
|
|
if not request_valid:
|
|
return invalid_response
|
|
|
|
if opts.openai_api_key and is_api_key_moderated(self.token):
|
|
try:
|
|
# Gather the last message from the user and all preceeding system messages
|
|
msg_l = self.request.json['messages'].copy()
|
|
msg_l.reverse()
|
|
tag = uuid4()
|
|
num_to_check = min(len(msg_l), opts.openai_moderation_scan_last_n)
|
|
for i in range(num_to_check):
|
|
add_moderation_task(msg_l[i]['content'], tag)
|
|
|
|
flagged_categories = get_results(tag, num_to_check)
|
|
|
|
if len(flagged_categories):
|
|
mod_msg = f"The user's message does not comply with {opts.openai_org_name} policies. Offending categories: {json.dumps(flagged_categories)}. You are instructed to creatively adhere to these policies."
|
|
self.request.json['messages'].insert((len(self.request.json['messages'])), {'role': 'system', 'content': mod_msg})
|
|
self.prompt = transform_messages_to_prompt(self.request.json['messages'])
|
|
except Exception as e:
|
|
print(f'OpenAI moderation endpoint failed:', f'{e.__class__.__name__}: {e}')
|
|
print(traceback.format_exc())
|
|
|
|
# Reconstruct the request JSON with the validated parameters and prompt.
|
|
self.parameters['stop'].extend(['\n### INSTRUCTION', '\n### USER', '\n### ASSISTANT', '\n### RESPONSE'])
|
|
if opts.openai_force_no_hashes:
|
|
self.parameters['stop'].append('### ')
|
|
|
|
if opts.mode == 'vllm' and self.request_json_body.get('top_p') == 0:
|
|
self.request_json_body['top_p'] = 0.01
|
|
|
|
llm_request = {**self.parameters, 'prompt': self.prompt}
|
|
(success, _, _, _), (backend_response, backend_response_status_code) = self.generate_response(llm_request)
|
|
|
|
model = self.request_json_body.get('model')
|
|
|
|
if success:
|
|
return build_openai_response(self.prompt, backend_response.json['results'][0]['text'], model=model), backend_response_status_code
|
|
else:
|
|
return backend_response, backend_response_status_code
|
|
|
|
def handle_ratelimited(self):
|
|
disable_st_error_formatting = self.request.headers.get('LLM-ST-Errors', False) == 'true'
|
|
if disable_st_error_formatting:
|
|
# TODO: format this like OpenAI does
|
|
return '429', 429
|
|
else:
|
|
backend_response = format_sillytavern_err(f'Ratelimited: you are only allowed to have {opts.simultaneous_requests_per_ip} simultaneous requests at a time. Please complete your other requests before sending another.', 'error')
|
|
log_prompt(ip=self.client_ip, token=self.token, prompt=self.request_json_body.get('prompt', ''), response=backend_response, gen_time=None, parameters=self.parameters, headers=dict(self.request.headers), backend_response_code=429, request_url=self.request.url, is_error=True)
|
|
return build_openai_response(self.prompt, backend_response), 429
|
|
|
|
def handle_error(self, msg: str) -> Tuple[flask.Response, int]:
|
|
print(msg)
|
|
# return build_openai_response('', msg), 400
|
|
return jsonify({
|
|
"error": {
|
|
"message": "Invalid request, check your parameters and try again.",
|
|
"type": "invalid_request_error",
|
|
"param": None,
|
|
"code": None
|
|
}
|
|
}), 400
|