fix: replace deprecated on_event with lifespan

This commit is contained in:
Tang Ziya 2024-05-09 12:00:03 +08:00
parent 90503be2ed
commit 1d9c745b35
2 changed files with 25 additions and 21 deletions

View File

@ -1,4 +1,5 @@
import sys
from contextlib import asynccontextmanager
from fastapi import FastAPI, Depends, HTTPException
from fastapi.routing import APIRoute
@ -46,7 +47,16 @@ import asyncio
import subprocess
import yaml
app = FastAPI()
@asynccontextmanager
async def lifespan(app: FastAPI):
log.info("startup_event")
# TODO: Check config.yaml file and create one
asyncio.create_task(start_litellm_background())
yield
app = FastAPI(lifespan=lifespan)
origins = ["*"]
@ -141,13 +151,6 @@ async def shutdown_litellm_background():
background_process = None
@app.on_event("startup")
async def startup_event():
log.info("startup_event")
# TODO: Check config.yaml file and create one
asyncio.create_task(start_litellm_background())
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST

View File

@ -1,3 +1,4 @@
from contextlib import asynccontextmanager
from bs4 import BeautifulSoup
import json
import markdown
@ -92,7 +93,19 @@ https://github.com/open-webui/open-webui
"""
)
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
@asynccontextmanager
async def lifespan(app: FastAPI):
if ENABLE_LITELLM:
asyncio.create_task(start_litellm_background())
yield
if ENABLE_LITELLM:
await shutdown_litellm_background()
app = FastAPI(
docs_url="/docs" if ENV == "dev" else None, redoc_url=None, lifespan=lifespan
)
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
@ -211,12 +224,6 @@ async def check_url(request: Request, call_next):
return response
@app.on_event("startup")
async def on_startup():
if ENABLE_LITELLM:
asyncio.create_task(start_litellm_background())
app.mount("/api/v1", webui_app)
app.mount("/litellm/api", litellm_app)
@ -381,9 +388,3 @@ else:
log.warning(
f"Frontend build directory not found at '{FRONTEND_BUILD_DIR}'. Serving API only."
)
@app.on_event("shutdown")
async def shutdown_event():
if ENABLE_LITELLM:
await shutdown_litellm_background()