Explorar o código

fix: replace deprecated on_event with lifespan

Tang Ziya hai 11 meses
pai
achega
1d9c745b35
Modificáronse 2 ficheiros con 25 adicións e 21 borrados
  1. 11 8
      backend/apps/litellm/main.py
  2. 14 13
      backend/main.py

+ 11 - 8
backend/apps/litellm/main.py

@@ -1,4 +1,5 @@
 import sys
 import sys
+from contextlib import asynccontextmanager
 
 
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.routing import APIRoute
 from fastapi.routing import APIRoute
@@ -46,7 +47,16 @@ import asyncio
 import subprocess
 import subprocess
 import yaml
 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 = ["*"]
 origins = ["*"]
 
 
@@ -141,13 +151,6 @@ async def shutdown_litellm_background():
         background_process = None
         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.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
 app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
 app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
 
 

+ 14 - 13
backend/main.py

@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 from bs4 import BeautifulSoup
 from bs4 import BeautifulSoup
 import json
 import json
 import markdown
 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.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
 app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
 app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
@@ -211,12 +224,6 @@ async def check_url(request: Request, call_next):
     return response
     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("/api/v1", webui_app)
 app.mount("/litellm/api", litellm_app)
 app.mount("/litellm/api", litellm_app)
 
 
@@ -381,9 +388,3 @@ else:
     log.warning(
     log.warning(
         f"Frontend build directory not found at '{FRONTEND_BUILD_DIR}'. Serving API only."
         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()