test / openai_server /server.py
iblfe's picture
Upload folder using huggingface_hub
b585c7f verified
import contextlib
import logging
import os
import sys
import ast
import json
from threading import Thread
import time
from traceback import print_exception
from typing import List
from pydantic import BaseModel, Field
import uvicorn
from fastapi import Depends, FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.requests import Request
from fastapi.responses import JSONResponse
from sse_starlette import EventSourceResponse
from starlette.responses import PlainTextResponse
from openai_server.log import logger
sys.path.append('openai_server')
# https://github.com/h2oai/h2ogpt/issues/1132
# https://github.com/jquesnelle/transformers-openai-api
# https://community.openai.com/t/trying-to-turn-this-into-an-automatic-web-search-engine/306383
class Generation(BaseModel):
# put here things not supported by OpenAI but are by torch or vLLM
# https://github.com/vllm-project/vllm/blob/main/vllm/sampling_params.py
top_k: int | None = 1
repetition_penalty: float | None = 1
min_p: float | None = 0.0
max_time: float | None = 360
class Params(BaseModel):
# https://platform.openai.com/docs/api-reference/completions/create
user: str | None = Field(default=None, description="Track user")
model: str | None = Field(default=None, description="Choose model")
best_of: int | None = Field(default=1, description="Unused")
frequency_penalty: float | None = 0.0
max_tokens: int | None = 256
n: int | None = Field(default=1, description="Unused")
presence_penalty: float | None = 0.0
stop: str | List[str] | None = None
stop_token_ids: List[int] | None = None
stream: bool | None = False
temperature: float | None = 0.3
top_p: float | None = 1.0
seed: int | None = 1234
class CompletionParams(Params):
prompt: str | List[str]
logit_bias: dict | None = None
logprobs: int | None = None
class TextRequest(Generation, CompletionParams):
pass
class TextResponse(BaseModel):
id: str
choices: List[dict]
created: int = int(time.time())
model: str
object: str = "text_completion"
usage: dict
class ChatParams(Params):
messages: List[dict]
tools: list | None = Field(default=None, description="WIP")
tool_choice: str | None = Field(default=None, description="WIP")
class ChatRequest(Generation, ChatParams):
# https://platform.openai.com/docs/api-reference/chat/create
pass
class ChatResponse(BaseModel):
id: str
choices: List[dict]
created: int = int(time.time())
model: str
object: str = "chat.completion"
usage: dict
class ModelInfoResponse(BaseModel):
model_name: str
class ModelListResponse(BaseModel):
model_names: List[str]
def verify_api_key(authorization: str = Header(None)) -> None:
server_api_key = os.getenv('H2OGPT_OPENAI_API_KEY', 'EMPTY')
if server_api_key == 'EMPTY':
# dummy case since '' cannot be handled
return
if server_api_key and (authorization is None or authorization != f"Bearer {server_api_key}"):
raise HTTPException(status_code=401, detail="Unauthorized")
app = FastAPI()
check_key = [Depends(verify_api_key)]
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"]
)
# https://platform.openai.com/docs/models/how-we-use-your-data
class InvalidRequestError(Exception):
pass
@app.exception_handler(Exception)
async def validation_exception_handler(request, exc):
print_exception(exc)
exc2 = InvalidRequestError(str(exc))
return PlainTextResponse(str(exc2), status_code=400)
@app.options("/", dependencies=check_key)
async def options_route():
return JSONResponse(content="OK")
@app.post('/v1/completions', response_model=TextResponse, dependencies=check_key)
async def openai_completions(request: Request, request_data: TextRequest):
if request_data.stream:
async def generator():
from openai_server.backend import stream_completions
response = stream_completions(dict(request_data))
for resp in response:
disconnected = await request.is_disconnected()
if disconnected:
break
yield {"data": json.dumps(resp)}
return EventSourceResponse(generator())
else:
from openai_server.backend import completions
response = completions(dict(request_data))
return JSONResponse(response)
@app.post('/v1/chat/completions', response_model=ChatResponse, dependencies=check_key)
async def openai_chat_completions(request: Request, request_data: ChatRequest):
if request_data.stream:
from openai_server.backend import stream_chat_completions
async def generator():
response = stream_chat_completions(dict(request_data))
for resp in response:
disconnected = await request.is_disconnected()
if disconnected:
break
yield {"data": json.dumps(resp)}
return EventSourceResponse(generator())
else:
from openai_server.backend import chat_completions
response = chat_completions(dict(request_data))
return JSONResponse(response)
# https://platform.openai.com/docs/api-reference/models/list
@app.get("/v1/models", dependencies=check_key)
@app.get("/v1/models/{model}", dependencies=check_key)
@app.get("/v1/models/{repo}/{model}", dependencies=check_key)
async def handle_models(request: Request):
path = request.url.path
model_name = path[len('/v1/models/'):]
from openai_server.backend import gradio_client
model_dict = ast.literal_eval(gradio_client.predict(api_name='/model_names'))
base_models = [x['base_model'] for x in model_dict]
if not model_name:
response = {
"object": "list",
"data": base_models,
}
else:
model_index = base_models.index(model_name)
if model_index >= 0:
response = model_dict[model_index]
else:
response = dict(model_name='INVALID')
return JSONResponse(response)
@app.get("/v1/internal/model/info", response_model=ModelInfoResponse, dependencies=check_key)
async def handle_model_info():
from openai_server.backend import get_model_info
return JSONResponse(content=get_model_info())
@app.get("/v1/internal/model/list", response_model=ModelListResponse, dependencies=check_key)
async def handle_list_models():
from openai_server.backend import get_model_list
return JSONResponse(content=get_model_list())
def run_server(host='0.0.0.0',
port=5000,
ssl_certfile=None,
ssl_keyfile=None,
gradio_prefix=None,
gradio_host=None,
gradio_port=None,
h2ogpt_key=None,
):
os.environ['GRADIO_PREFIX'] = gradio_prefix or 'http'
os.environ['GRADIO_SERVER_HOST'] = gradio_host or 'localhost'
os.environ['GRADIO_SERVER_PORT'] = gradio_port or '7860'
os.environ['GRADIO_H2OGPT_H2OGPT_KEY'] = h2ogpt_key or '' # don't use H2OGPT_H2OGPT_KEY, mixes things up
# use h2ogpt_key if no server api key, so OpenAI inherits key by default if any keys set and enforced via API for h2oGPT
# but OpenAI key cannot be '', so dummy value is EMPTY and if EMPTY we ignore the key in authorization
server_api_key = os.getenv('H2OGPT_OPENAI_API_KEY', os.environ['GRADIO_H2OGPT_H2OGPT_KEY']) or 'EMPTY'
os.environ['H2OGPT_OPENAI_API_KEY'] = server_api_key
port = int(os.getenv('H2OGPT_OPENAI_PORT', port))
ssl_certfile = os.getenv('H2OGPT_OPENAI_CERT_PATH', ssl_certfile)
ssl_keyfile = os.getenv('H2OGPT_OPENAI_KEY_PATH', ssl_keyfile)
prefix = 'https' if ssl_keyfile and ssl_certfile else 'http'
logger.info(f'OpenAI API URL: {prefix}://{host}:{port}')
logger.info(f'OpenAI API key: {server_api_key}')
logging.getLogger("uvicorn.error").propagate = False
uvicorn.run(app, host=host, port=port, ssl_certfile=ssl_certfile, ssl_keyfile=ssl_keyfile)
def run(wait=True, **kwargs):
if wait:
run_server(**kwargs)
else:
Thread(target=run_server, kwargs=kwargs, daemon=True).start()