open-webui/backend/main.py

138 lines
3.5 KiB
Python
Raw Normal View History

2024-02-23 09:30:26 +01:00
from bs4 import BeautifulSoup
import json
import markdown
2024-01-07 11:48:21 +01:00
import time
2024-02-24 09:21:53 +01:00
import os
import sys
2024-02-23 09:30:26 +01:00
2024-02-24 07:44:56 +01:00
from fastapi import FastAPI, Request, Depends
2023-11-15 01:28:51 +01:00
from fastapi.staticfiles import StaticFiles
from fastapi import HTTPException
2024-02-24 07:44:56 +01:00
from fastapi.responses import JSONResponse
2023-11-15 01:28:51 +01:00
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.middleware.cors import CORSMiddleware
2023-11-19 01:47:12 +01:00
from starlette.exceptions import HTTPException as StarletteHTTPException
2023-11-15 01:28:51 +01:00
2024-01-07 07:07:20 +01:00
2024-02-24 09:21:53 +01:00
from litellm.proxy.proxy_server import ProxyConfig, initialize
2024-02-22 12:22:23 +01:00
from litellm.proxy.proxy_server import app as litellm_app
2023-11-15 01:28:51 +01:00
from apps.ollama.main import app as ollama_app
2024-01-05 03:38:03 +01:00
from apps.openai.main import app as openai_app
2024-02-11 09:17:50 +01:00
from apps.audio.main import app as audio_app
2024-02-22 03:12:01 +01:00
from apps.images.main import app as images_app
from apps.rag.main import app as rag_app
2023-11-19 01:47:12 +01:00
from apps.web.main import app as webui_app
2024-01-07 07:07:20 +01:00
2024-02-24 07:44:56 +01:00
2024-02-24 02:12:19 +01:00
from config import WEBUI_NAME, ENV, VERSION, CHANGELOG, FRONTEND_BUILD_DIR
2024-02-24 07:44:56 +01:00
from utils.utils import get_http_authorization_cred, get_current_user
2023-11-15 01:28:51 +01:00
class SPAStaticFiles(StaticFiles):
async def get_response(self, path: str, scope):
try:
return await super().get_response(path, scope)
except (HTTPException, StarletteHTTPException) as ex:
if ex.status_code == 404:
return await super().get_response("index.html", scope)
else:
raise ex
2024-02-24 09:21:53 +01:00
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()
2024-01-07 11:48:21 +01:00
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
2023-11-15 01:28:51 +01:00
origins = ["*"]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2024-02-24 09:21:53 +01:00
@app.on_event("startup")
async def on_startup():
await startup()
2023-11-15 01:28:51 +01:00
@app.middleware("http")
async def check_url(request: Request, call_next):
start_time = int(time.time())
response = await call_next(request)
process_time = int(time.time()) - start_time
response.headers["X-Process-Time"] = str(process_time)
return response
2024-02-24 07:44:56 +01:00
@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
2023-11-19 01:47:12 +01:00
app.mount("/api/v1", webui_app)
2024-02-22 12:22:23 +01:00
app.mount("/litellm/api", litellm_app)
2024-01-04 22:06:31 +01:00
app.mount("/ollama/api", ollama_app)
2024-01-05 03:38:03 +01:00
app.mount("/openai/api", openai_app)
2024-02-11 09:17:50 +01:00
2024-02-22 03:12:01 +01:00
app.mount("/images/api/v1", images_app)
2024-02-11 09:17:50 +01:00
app.mount("/audio/api/v1", audio_app)
2024-01-07 07:07:20 +01:00
app.mount("/rag/api/v1", rag_app)
2024-01-04 22:06:31 +01:00
2024-02-22 03:12:01 +01:00
@app.get("/api/config")
async def get_app_config():
2024-02-23 09:30:26 +01:00
2024-02-22 03:12:01 +01:00
return {
"status": True,
2024-02-24 02:12:19 +01:00
"name": WEBUI_NAME,
2024-02-23 09:30:26 +01:00
"version": VERSION,
2024-02-22 03:12:01 +01:00
"images": images_app.state.ENABLED,
"default_models": webui_app.state.DEFAULT_MODELS,
"default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
}
2024-02-23 09:30:26 +01:00
@app.get("/api/changelog")
async def get_app_changelog():
2024-02-23 09:36:53 +01:00
return CHANGELOG
2024-02-23 09:30:26 +01:00
2024-02-24 02:12:19 +01:00
app.mount("/static", StaticFiles(directory="static"), name="static")
2024-01-22 10:47:07 +01:00
app.mount(
"/",
SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
2024-01-22 10:47:07 +01:00
name="spa-static-files",
)