main.py 7.02 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
from starlette.middleware.base import BaseHTTPMiddleware
Timothy J. Baek's avatar
Timothy J. Baek committed
16

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

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

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

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

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
45
46
47
48
49
50
51
52
53
54
55
56

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


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

59
60
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
61

Timothy J. Baek's avatar
Timothy J. Baek committed
62
63
64
app.state.WEBHOOK_URL = WEBHOOK_URL


Timothy J. Baek's avatar
Timothy J. Baek committed
65
66
origins = ["*"]

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

Timothy J. Baek's avatar
Timothy J. Baek committed
68
69
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
70
71
72
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
73
74
            print(request.url.path)

Timothy J. Baek's avatar
Timothy J. Baek committed
75
76
77
78
79
80
81
82
83
84
85
            # 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
86
87
88
89
90
91
92
93
                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
94
95
                del data["docs"]

Timothy J. Baek's avatar
Timothy J. Baek committed
96
97
                print(data["messages"])

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

Timothy J. Baek's avatar
Timothy J. Baek committed
100
101
102
103
104
105
106
107
108
109
110
111
            # 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
112
113
114
115
116
117
118
119
120
121
122

        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
123
124
125
126
127
128
129
130
131
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
132
133
134
135
136
137
138
139
140
141
@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
142
143
144
145
146
@app.on_event("startup")
async def on_startup():
    await litellm_app_startup()


147
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
148
149
app.mount("/litellm/api", litellm_app)

150
app.mount("/ollama", ollama_app)
151
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
152

Timothy J. Baek's avatar
Timothy J. Baek committed
153
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
154
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
155
156
app.mount("/rag/api/v1", rag_app)

157

Timothy J. Baek's avatar
Timothy J. Baek committed
158
159
@app.get("/api/config")
async def get_app_config():
Timothy J. Baek's avatar
Timothy J. Baek committed
160

Timothy J. Baek's avatar
Timothy J. Baek committed
161
162
    return {
        "status": True,
163
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
164
        "version": VERSION,
Timothy J. Baek's avatar
Timothy J. Baek committed
165
166
167
168
169
170
        "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
171
172
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
173
174
175
176
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
177
178
179
180
181
182
183
184


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
185
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
186
187
188
189
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):

    app.state.MODEL_FILTER_ENABLED = form_data.enabled
190
    app.state.MODEL_FILTER_LIST = form_data.models
Timothy J. Baek's avatar
Timothy J. Baek committed
191
192

    ollama_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
193
    ollama_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
194
195

    openai_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
196
    openai_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
197

198
199
200
201
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
202
203


Timothy J. Baek's avatar
Timothy J. Baek committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
@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,
    }


226
227
228
229
230
231
232
233
@app.get("/api/version")
async def get_app_config():

    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
234
235
@app.get("/api/changelog")
async def get_app_changelog():
Timothy J. Baek's avatar
Timothy J. Baek committed
236
    return CHANGELOG
Timothy J. Baek's avatar
Timothy J. Baek committed
237
238


239
240
241
242
243
244
245
246
247
@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
248
        return {"current": VERSION, "latest": latest_version[1:]}
249
250
251
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
252
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
253
254
255
        )


256
app.mount("/static", StaticFiles(directory="static"), name="static")
Timothy J. Baek's avatar
Timothy J. Baek committed
257
app.mount("/cache", StaticFiles(directory="data/cache"), name="cache")
258
259


260
261
app.mount(
    "/",
lucasew's avatar
lucasew committed
262
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
263
264
    name="spa-static-files",
)