Merge branch 'dev' into feature/img-gen-steps-setting

backend/.gitignore

@@ -6,6 +6,11 @@ uploads
 *.db
 _test
 Pipfile
-data/*
+!/data
+/data/*
+!/data/litellm
+/data/litellm/*
+!data/litellm/config.yaml
+
 !data/config.json
 .webui_secret_key

backend/apps/images/main.py

@@ -136,6 +136,7 @@ def get_models(user=Depends(get_current_user)):
         models = r.json()
         return models
     except Exception as e:
+        app.state.ENABLED = False
         raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
 
 
@@ -147,6 +148,7 @@ async def get_default_model(user=Depends(get_admin_user)):
 
         return {"model": options["sd_model_checkpoint"]}
     except Exception as e:
+        app.state.ENABLED = False
         raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
 
 

backend/config.py

@@ -83,8 +83,6 @@ for version in soup.find_all("h2"):
     # Find the next sibling that is a h3 tag (section title)
     current = version.find_next_sibling()
 
-    print(current)
-
     while current and current.name != "h2":
         if current.name == "h3":
             section_title = current.get_text().lower()  # e.g., "added", "fixed"

backend/data/litellm/config.yaml

@@ -0,0 +1,4 @@
+general_settings: {}
+litellm_settings: {}
+model_list: []
+router_settings: {}

backend/main.py

@@ -2,25 +2,31 @@ from bs4 import BeautifulSoup
 import json
 import markdown
 import time
+import os
+import sys
 
-
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Depends
 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.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
-
 from apps.web.main import app as webui_app
 
+
 from config import WEBUI_NAME, ENV, VERSION, CHANGELOG, FRONTEND_BUILD_DIR
+from utils.utils import get_http_authorization_cred, get_current_user
 
 
 class SPAStaticFiles(StaticFiles):
@@ -34,6 +40,21 @@ 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 = ["*"]
@@ -47,6 +68,11 @@ app.add_middleware(
 )
 
 
+@app.on_event("startup")
+async def on_startup():
+    await startup()
+
+
 @app.middleware("http")
 async def check_url(request: Request, call_next):
     start_time = int(time.time())
@@ -57,7 +83,23 @@ 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)
 
 app.mount("/ollama/api", ollama_app)
 app.mount("/openai/api", openai_app)

backend/requirements.txt

@@ -16,6 +16,9 @@ aiohttp
 peewee
 bcrypt
 
+litellm
+apscheduler
+
 langchain
 langchain-community
 chromadb

backend/utils/utils.py

@@ -58,6 +58,17 @@ def extract_token_from_auth_header(auth_header: str):
     return auth_header[len("Bearer ") :]
 
 
+def get_http_authorization_cred(auth_header: str):
+    try:
+        scheme, credentials = auth_header.split(" ")
+        return {
+            "scheme": scheme,
+            "credentials": credentials,
+        }
+    except:
+        raise ValueError(ERROR_MESSAGES.INVALID_TOKEN)
+
+
 def get_current_user(
     auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
 ):