main.py 8.41 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
    MANIFEST_NAME,
    MANIFEST_SHORT_NAME
48
)
49
50
from constants import ERROR_MESSAGES

51
52
53
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
54

55

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


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

69
70
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
71

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

74
75
app.state.MANIFEST_NAME = MANIFEST_NAME
app.state.MANIFEST_SHORT_NAME = MANIFEST_SHORT_NAME
Timothy J. Baek's avatar
Timothy J. Baek committed
76

Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
origins = ["*"]

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
87
88
89
90
91
92
93
94
95
96
            # 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
97
98
99
100
101
102
103
104
                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
105
106
                del data["docs"]

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
111
112
113
114
115
116
117
118
119
120
121
122
            # 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
123
124
125
126
127
128
129
130
131
132
133

        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
134
135
136
137
138
139
140
141
142
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
143
144
145
146
147
148
149
150
151
152
@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
153
154
155
156
157
@app.on_event("startup")
async def on_startup():
    await litellm_app_startup()


158
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
159
160
app.mount("/litellm/api", litellm_app)

161
app.mount("/ollama", ollama_app)
162
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
163

Timothy J. Baek's avatar
Timothy J. Baek committed
164
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
165
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
166
167
app.mount("/rag/api/v1", rag_app)

168

Timothy J. Baek's avatar
Timothy J. Baek committed
169
170
@app.get("/api/config")
async def get_app_config():
171
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA
172
173

    default_locale = "en-US"
174
175
176
177
    if "ui" in CONFIG_DATA:
        default_locale = CONFIG_DATA["ui"].get("default_locale", "en-US")

    # The Rest of the Function Now Uses the Variables Defined Above
Timothy J. Baek's avatar
Timothy J. Baek committed
178
179
    return {
        "status": True,
180
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
181
        "version": VERSION,
182
        "default_locale": default_locale,
Timothy J. Baek's avatar
Timothy J. Baek committed
183
184
        "images": images_app.state.ENABLED,
        "default_models": webui_app.state.DEFAULT_MODELS,
185
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
Timothy J. Baek's avatar
Timothy J. Baek committed
186
187
188
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
191
192
193
194
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
195
196
197
198
199
200
201
202


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
203
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
204
205
206
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
    app.state.MODEL_FILTER_ENABLED = form_data.enabled
207
    app.state.MODEL_FILTER_LIST = form_data.models
Timothy J. Baek's avatar
Timothy J. Baek committed
208
209

    ollama_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
210
    ollama_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
211
212

    openai_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
213
    openai_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
214

215
216
217
    litellm_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
    litellm_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST

218
219
220
221
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
222
223


Timothy J. Baek's avatar
Timothy J. Baek committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
@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,
    }


246
247
248
249
250
251
252
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
253
254
@app.get("/api/changelog")
async def get_app_changelog():
255
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
256
257


258
259
260
261
262
263
264
265
266
@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
267
        return {"current": VERSION, "latest": latest_version[1:]}
268
269
270
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
271
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
272
273
274
        )


275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
@app.get("/manifest.json")
async def get_manifest_json():
    return {
        "name": app.state.MANIFEST_NAME,
        "short_name": app.state.MANIFEST_SHORT_NAME,
        "start_url": "/",
        "display": "standalone",
        "background_color": "#343541",
        "theme_color": "#343541",
        "orientation": "portrait-primary",
        "icons": [
            {
                "src": "/favicon.png",
                "type": "image/png",
                "sizes": "844x884"
            }
        ]
    }


295
app.mount("/static", StaticFiles(directory="static"), name="static")
Timothy J. Baek's avatar
Timothy J. Baek committed
296
app.mount("/cache", StaticFiles(directory="data/cache"), name="cache")
297
298


299
300
app.mount(
    "/",
lucasew's avatar
lucasew committed
301
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
302
303
    name="spa-static-files",
)