From 2b84af878a2bd0deab5423761a48705dcd8cb984 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 8 Mar 2024 13:33:56 -0800
Subject: [PATCH] refac: litellm

---
 backend/apps/litellm/main.py | 41 ++++++++++++++++++++++++++++++++++++
 backend/main.py              | 39 ++--------------------------------
 2 files changed, 43 insertions(+), 37 deletions(-)
 create mode 100644 backend/apps/litellm/main.py

diff --git a/backend/apps/litellm/main.py b/backend/apps/litellm/main.py
new file mode 100644
index 00000000..21b9e58a
--- /dev/null
+++ b/backend/apps/litellm/main.py
@@ -0,0 +1,41 @@
+from litellm.proxy.proxy_server import ProxyConfig, initialize
+from litellm.proxy.proxy_server import app
+
+from fastapi import FastAPI, Request, Depends, status
+from fastapi.responses import JSONResponse
+from utils.utils import get_http_authorization_cred, get_current_user
+from config import ENV
+
+proxy_config = ProxyConfig()
+
+
+async def config():
+    router, model_list, general_settings = await proxy_config.load_config(
+        router=None, config_file_path="./data/litellm/config.yaml"
+    )
+
+    await initialize(config="./data/litellm/config.yaml", telemetry=False)
+
+
+async def startup():
+    await config()
+
+
+@app.on_event("startup")
+async def on_startup():
+    await startup()
+
+
+@app.middleware("http")
+async def auth_middleware(request: Request, call_next):
+    auth_header = request.headers.get("Authorization", "")
+
+    if ENV != "dev":
+        try:
+            user = get_current_user(get_http_authorization_cred(auth_header))
+            print(user)
+        except Exception as e:
+            return JSONResponse(status_code=400, content={"detail": str(e)})
+
+    response = await call_next(request)
+    return response
diff --git a/backend/main.py b/backend/main.py
index 5f6b4441..9e04ee48 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -9,17 +9,14 @@ import requests
 from fastapi import FastAPI, Request, Depends, status
 from fastapi.staticfiles import StaticFiles
 from fastapi import HTTPException
-from fastapi.responses import JSONResponse
 from fastapi.middleware.wsgi import WSGIMiddleware
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
 
-from litellm.proxy.proxy_server import ProxyConfig, initialize
-from litellm.proxy.proxy_server import app as litellm_app
-
 from apps.ollama.main import app as ollama_app
 from apps.openai.main import app as openai_app
+from apps.litellm.main import app as litellm_app, startup as litellm_app_startup
 from apps.audio.main import app as audio_app
 from apps.images.main import app as images_app
 from apps.rag.main import app as rag_app
@@ -29,8 +26,6 @@ from apps.web.main import app as webui_app
 from config import WEBUI_NAME, ENV, VERSION, CHANGELOG, FRONTEND_BUILD_DIR
 from constants import ERROR_MESSAGES
 
-from utils.utils import get_http_authorization_cred, get_current_user
-
 
 class SPAStaticFiles(StaticFiles):
     async def get_response(self, path: str, scope):
@@ -43,21 +38,6 @@ class SPAStaticFiles(StaticFiles):
                 raise ex
 
 
-proxy_config = ProxyConfig()
-
-
-async def config():
-    router, model_list, general_settings = await proxy_config.load_config(
-        router=None, config_file_path="./data/litellm/config.yaml"
-    )
-
-    await initialize(config="./data/litellm/config.yaml", telemetry=False)
-
-
-async def startup():
-    await config()
-
-
 app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
 
 origins = ["*"]
@@ -73,7 +53,7 @@ app.add_middleware(
 
 @app.on_event("startup")
 async def on_startup():
-    await startup()
+    await litellm_app_startup()
 
 
 @app.middleware("http")
@@ -86,21 +66,6 @@ async def check_url(request: Request, call_next):
     return response
 
 
-@litellm_app.middleware("http")
-async def auth_middleware(request: Request, call_next):
-    auth_header = request.headers.get("Authorization", "")
-
-    if ENV != "dev":
-        try:
-            user = get_current_user(get_http_authorization_cred(auth_header))
-            print(user)
-        except Exception as e:
-            return JSONResponse(status_code=400, content={"detail": str(e)})
-
-    response = await call_next(request)
-    return response
-
-
 app.mount("/api/v1", webui_app)
 app.mount("/litellm/api", litellm_app)