Spaces:

iblfe
/

test

Runtime error

App Files Files Community

test / openai_server /server.py

iblfe

Upload folder using huggingface_hub

b585c7f verified 5 months ago

raw history blame contribute delete

No virus

8.31 kB

	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()