main.py 3.37 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
3
from bs4 import BeautifulSoup
import json
import markdown
4
import time
Timothy J. Baek's avatar
Timothy J. Baek committed
5
6
import os
import sys
7
import requests
Timothy J. Baek's avatar
Timothy J. Baek committed
8

9
from fastapi import FastAPI, Request, Depends, status
Timothy J. Baek's avatar
Timothy J. Baek committed
10
11
12
13
from fastapi.staticfiles import StaticFiles
from fastapi import HTTPException
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.middleware.cors import CORSMiddleware
14
from starlette.exceptions import HTTPException as StarletteHTTPException
Timothy J. Baek's avatar
Timothy J. Baek committed
15

Timothy J. Baek's avatar
Timothy J. Baek committed
16

Timothy J. Baek's avatar
Timothy J. Baek committed
17
from apps.ollama.main import app as ollama_app
18
from apps.openai.main import app as openai_app
Timothy J. Baek's avatar
Timothy J. Baek committed
19
from apps.litellm.main import app as litellm_app, startup as litellm_app_startup
Timothy J. Baek's avatar
Timothy J. Baek committed
20
from apps.audio.main import app as audio_app
Timothy J. Baek's avatar
Timothy J. Baek committed
21
22
from apps.images.main import app as images_app
from apps.rag.main import app as rag_app
23
from apps.web.main import app as webui_app
Timothy J. Baek's avatar
Timothy J. Baek committed
24

Timothy J. Baek's avatar
Timothy J. Baek committed
25

26
from config import WEBUI_NAME, ENV, VERSION, CHANGELOG, FRONTEND_BUILD_DIR
27
28
from constants import ERROR_MESSAGES

Timothy J. Baek's avatar
Timothy J. Baek committed
29
30
31
32
33
34
35
36
37
38
39
40

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


41
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
Timothy J. Baek's avatar
Timothy J. Baek committed
42
43
44
45
46
47
48
49
50
51
52
53

origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
54
55
@app.on_event("startup")
async def on_startup():
Timothy J. Baek's avatar
Timothy J. Baek committed
56
    await litellm_app_startup()
Timothy J. Baek's avatar
Timothy J. Baek committed
57
58


Timothy J. Baek's avatar
Timothy J. Baek committed
59
60
61
62
63
64
65
66
67
68
@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


69
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
70
71
app.mount("/litellm/api", litellm_app)

72
app.mount("/ollama", ollama_app)
73
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
74

Timothy J. Baek's avatar
Timothy J. Baek committed
75
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
76
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
app.mount("/rag/api/v1", rag_app)

79

Timothy J. Baek's avatar
Timothy J. Baek committed
80
81
@app.get("/api/config")
async def get_app_config():
Timothy J. Baek's avatar
Timothy J. Baek committed
82

Timothy J. Baek's avatar
Timothy J. Baek committed
83
84
    return {
        "status": True,
85
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
86
        "version": VERSION,
Timothy J. Baek's avatar
Timothy J. Baek committed
87
88
89
90
91
92
        "images": images_app.state.ENABLED,
        "default_models": webui_app.state.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
    }


93
94
95
96
97
98
99
100
@app.get("/api/version")
async def get_app_config():

    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
101
102
@app.get("/api/changelog")
async def get_app_changelog():
Timothy J. Baek's avatar
Timothy J. Baek committed
103
    return CHANGELOG
Timothy J. Baek's avatar
Timothy J. Baek committed
104
105


106
107
108
109
110
111
112
113
114
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
        response = requests.get(
            f"https://api.github.com/repos/open-webui/open-webui/releases/latest"
        )
        response.raise_for_status()
        latest_version = response.json()["tag_name"]

Timothy J. Baek's avatar
Timothy J. Baek committed
115
        return {"current": VERSION, "latest": latest_version[1:]}
116
117
118
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
119
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
120
121
122
        )


123
app.mount("/static", StaticFiles(directory="static"), name="static")
Timothy J. Baek's avatar
Timothy J. Baek committed
124
app.mount("/cache", StaticFiles(directory="data/cache"), name="cache")
125
126


127
128
app.mount(
    "/",
lucasew's avatar
lucasew committed
129
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
130
131
    name="spa-static-files",
)