diff --git a/backend/apps/litellm/main.py b/backend/apps/litellm/main.py index 95f442067..6db426439 100644 --- a/backend/apps/litellm/main.py +++ b/backend/apps/litellm/main.py @@ -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 diff --git a/backend/main.py b/backend/main.py index 139819f7c..2d8d9ed68 100644 --- a/backend/main.py +++ b/backend/main.py @@ -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()