main.py 7.66 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 logging
8
import requests
Timothy J. Baek's avatar
Timothy J. Baek committed
9

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

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
27
28
from pydantic import BaseModel
from typing import List
Timothy J. Baek's avatar
Timothy J. Baek committed
29

Timothy J. Baek's avatar
Timothy J. Baek committed
30
31

from utils.utils import get_admin_user
Timothy J. Baek's avatar
Timothy J. Baek committed
32
from apps.rag.utils import rag_messages
Timothy J. Baek's avatar
Timothy J. Baek committed
33

34
from config import (
35
    CONFIG_DATA,
36
37
38
39
40
41
42
    WEBUI_NAME,
    ENV,
    VERSION,
    CHANGELOG,
    FRONTEND_BUILD_DIR,
    MODEL_FILTER_ENABLED,
    MODEL_FILTER_LIST,
43
44
    GLOBAL_LOG_LEVEL,
    SRC_LOG_LEVELS,
Timothy J. Baek's avatar
Timothy J. Baek committed
45
    WEBHOOK_URL,
46
)
47
48
from constants import ERROR_MESSAGES

49
50
51
logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MAIN"])
Timothy J. Baek's avatar
Timothy J. Baek committed
52

53

Timothy J. Baek's avatar
Timothy J. Baek committed
54
55
56
57
58
59
60
61
62
63
64
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


65
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
Timothy J. Baek's avatar
Timothy J. Baek committed
66

67
68
app.state.MODEL_FILTER_ENABLED = MODEL_FILTER_ENABLED
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
69

Timothy J. Baek's avatar
Timothy J. Baek committed
70
71
72
app.state.WEBHOOK_URL = WEBHOOK_URL


Timothy J. Baek's avatar
Timothy J. Baek committed
73
74
origins = ["*"]

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

Timothy J. Baek's avatar
Timothy J. Baek committed
76
77
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
78
79
80
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
81
            log.debug(f"request.url.path: {request.url.path}")
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
82

Timothy J. Baek's avatar
Timothy J. Baek committed
83
84
85
86
87
88
89
90
91
92
            # Read the original request body
            body = await request.body()
            # Decode body to string
            body_str = body.decode("utf-8")
            # Parse string to JSON
            data = json.loads(body_str) if body_str else {}

            # Example: Add a new key-value pair or modify existing ones
            # data["modified"] = True  # Example modification
            if "docs" in data:
Timothy J. Baek's avatar
Timothy J. Baek committed
93
94
95
96
97
98
99
100
                data = {**data}
                data["messages"] = rag_messages(
                    data["docs"],
                    data["messages"],
                    rag_app.state.RAG_TEMPLATE,
                    rag_app.state.TOP_K,
                    rag_app.state.sentence_transformer_ef,
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
101
102
                del data["docs"]

103
                log.debug(f"data['messages']: {data['messages']}")
Timothy J. Baek's avatar
Timothy J. Baek committed
104

Timothy J. Baek's avatar
Timothy J. Baek committed
105
106
            modified_body_bytes = json.dumps(data).encode("utf-8")

Timothy J. Baek's avatar
Timothy J. Baek committed
107
108
109
110
111
112
113
114
115
116
117
118
            # Replace the request body with the modified one
            request._body = modified_body_bytes

            # Set custom header to ensure content-length matches new body length
            request.headers.__dict__["_list"] = [
                (b"content-length", str(len(modified_body_bytes)).encode("utf-8")),
                *[
                    (k, v)
                    for k, v in request.headers.raw
                    if k.lower() != b"content-length"
                ],
            ]
Timothy J. Baek's avatar
Timothy J. Baek committed
119
120
121
122
123
124
125
126
127
128
129

        response = await call_next(request)
        return response

    async def _receive(self, body: bytes):
        return {"type": "http.request", "body": body, "more_body": False}


app.add_middleware(RAGMiddleware)


Timothy J. Baek's avatar
Timothy J. Baek committed
130
131
132
133
134
135
136
137
138
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
139
140
141
142
143
144
145
146
147
148
@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


Timothy J. Baek's avatar
Timothy J. Baek committed
149
150
151
152
153
@app.on_event("startup")
async def on_startup():
    await litellm_app_startup()


154
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
155
156
app.mount("/litellm/api", litellm_app)

157
app.mount("/ollama", ollama_app)
158
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
159

Timothy J. Baek's avatar
Timothy J. Baek committed
160
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
161
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
162
163
app.mount("/rag/api/v1", rag_app)

164

Timothy J. Baek's avatar
Timothy J. Baek committed
165
166
167
168
@app.get("/api/config")
async def get_app_config():
    return {
        "status": True,
169
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
170
        "version": VERSION,
171
172
        "default_locale": (
            CONFIG_DATA["ui"]["default_locale"]
Timothy J. Baek's avatar
Timothy J. Baek committed
173
            if "ui" in CONFIG_DATA and "default_locale" in CONFIG_DATA["ui"]
174
175
            else "en-US"
        ),
Timothy J. Baek's avatar
Timothy J. Baek committed
176
177
178
179
180
181
        "images": images_app.state.ENABLED,
        "default_models": webui_app.state.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
182
183
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
184
185
186
187
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
188
189
190
191
192
193
194
195


class ModelFilterConfigForm(BaseModel):
    enabled: bool
    models: List[str]


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
196
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
197
198
199
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
    app.state.MODEL_FILTER_ENABLED = form_data.enabled
200
    app.state.MODEL_FILTER_LIST = form_data.models
Timothy J. Baek's avatar
Timothy J. Baek committed
201
202

    ollama_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
203
    ollama_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
204
205

    openai_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
206
    openai_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
207

208
209
210
    litellm_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
    litellm_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST

211
212
213
214
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
215
216


Timothy J. Baek's avatar
Timothy J. Baek committed
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
@app.get("/api/webhook")
async def get_webhook_url(user=Depends(get_admin_user)):
    return {
        "url": app.state.WEBHOOK_URL,
    }


class UrlForm(BaseModel):
    url: str


@app.post("/api/webhook")
async def update_webhook_url(form_data: UrlForm, user=Depends(get_admin_user)):
    app.state.WEBHOOK_URL = form_data.url

    webui_app.state.WEBHOOK_URL = app.state.WEBHOOK_URL

    return {
        "url": app.state.WEBHOOK_URL,
    }


239
240
241
242
243
244
245
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
246
247
@app.get("/api/changelog")
async def get_app_changelog():
248
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
249
250


251
252
253
254
255
256
257
258
259
@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
260
        return {"current": VERSION, "latest": latest_version[1:]}
261
262
263
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
264
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
265
266
267
        )


268
app.mount("/static", StaticFiles(directory="static"), name="static")
Timothy J. Baek's avatar
Timothy J. Baek committed
269
app.mount("/cache", StaticFiles(directory="data/cache"), name="cache")
270
271


272
273
app.mount(
    "/",
lucasew's avatar
lucasew committed
274
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
275
276
    name="spa-static-files",
)