refac: concurrency

This commit is contained in:
Timothy J. Baek 2024-05-21 17:29:31 -07:00
parent eaa4112f46
commit 6b4fba3309

22
main.py
View File

@ -1,5 +1,7 @@
from fastapi import FastAPI, Request, Depends, status, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.concurrency import run_in_threadpool
from starlette.responses import StreamingResponse, Response
from pydantic import BaseModel, ConfigDict
@ -116,12 +118,13 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
detail=f"Model {form_data.model} not found",
)
def job():
get_response = PIPELINES[form_data.model]["module"].get_response
if form_data.stream:
def stream_content():
res = get_response(user_message, messages=form_data.messages)
if isinstance(res, str):
@ -130,6 +133,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
elif isinstance(res, Generator):
for message in res:
print(message)
message = stream_message_template(message)
yield f"data: {json.dumps(message)}\n\n"
@ -139,7 +143,12 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
"created": int(time.time()),
"model": MODEL_ID,
"choices": [
{"index": 0, "delta": {}, "logprobs": None, "finish_reason": "stop"}
{
"index": 0,
"delta": {},
"logprobs": None,
"finish_reason": "stop",
}
],
}
@ -176,6 +185,15 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
],
}
try:
return await run_in_threadpool(job)
except Exception as e:
print(e)
raise HTTPException(
status_code=500,
detail="{e}",
)
@app.get("/")
async def get_status():