main.py 13 KB
Newer Older
1
from contextlib import asynccontextmanager
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
4
from bs4 import BeautifulSoup
import json
import markdown
5
import time
Timothy J. Baek's avatar
Timothy J. Baek committed
6
7
import os
import sys
8
import logging
Timothy J. Baek's avatar
Timothy J. Baek committed
9
import aiohttp
10
import requests
11
import mimetypes
Timothy J. Baek's avatar
Timothy J. Baek committed
12

13
from fastapi import FastAPI, Request, Depends, status
Timothy J. Baek's avatar
Timothy J. Baek committed
14
15
16
17
from fastapi.staticfiles import StaticFiles
from fastapi import HTTPException
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.middleware.cors import CORSMiddleware
18
from starlette.exceptions import HTTPException as StarletteHTTPException
Timothy J. Baek's avatar
Timothy J. Baek committed
19
from starlette.middleware.base import BaseHTTPMiddleware
20
from starlette.responses import StreamingResponse, Response
Timothy J. Baek's avatar
Timothy J. Baek committed
21

Timothy J. Baek's avatar
Timothy J. Baek committed
22
from apps.ollama.main import app as ollama_app
23
from apps.openai.main import app as openai_app
Timothy J. Baek's avatar
Timothy J. Baek committed
24

25
26
27
28
29
from apps.litellm.main import (
    app as litellm_app,
    start_litellm_background,
    shutdown_litellm_background,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
30
31


Timothy J. Baek's avatar
Timothy J. Baek committed
32
from apps.audio.main import app as audio_app
Timothy J. Baek's avatar
Timothy J. Baek committed
33
34
from apps.images.main import app as images_app
from apps.rag.main import app as rag_app
35
from apps.web.main import app as webui_app
Timothy J. Baek's avatar
Timothy J. Baek committed
36

Timothy J. Baek's avatar
Timothy J. Baek committed
37
import asyncio
Timothy J. Baek's avatar
Timothy J. Baek committed
38
from pydantic import BaseModel
39
from typing import List, Optional
Timothy J. Baek's avatar
Timothy J. Baek committed
40

41
from apps.web.models.models import Models, ModelModel
Timothy J. Baek's avatar
Timothy J. Baek committed
42
from utils.utils import get_admin_user
Timothy J. Baek's avatar
Timothy J. Baek committed
43
from apps.rag.utils import rag_messages
Timothy J. Baek's avatar
Timothy J. Baek committed
44

45
from config import (
46
    CONFIG_DATA,
47
    WEBUI_NAME,
48
    WEBUI_URL,
49
    WEBUI_AUTH,
50
51
52
53
    ENV,
    VERSION,
    CHANGELOG,
    FRONTEND_BUILD_DIR,
54
55
    CACHE_DIR,
    STATIC_DIR,
56
    ENABLE_LITELLM,
Timothy J. Baek's avatar
Timothy J. Baek committed
57
    ENABLE_MODEL_FILTER,
58
    MODEL_FILTER_LIST,
59
60
    GLOBAL_LOG_LEVEL,
    SRC_LOG_LEVELS,
Timothy J. Baek's avatar
Timothy J. Baek committed
61
    WEBHOOK_URL,
62
    ENABLE_ADMIN_EXPORT,
63
    AppConfig,
64
)
65
66
from constants import ERROR_MESSAGES

67
68
69
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
70

71

Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
76
77
78
79
80
81
82
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


Timothy J. Baek's avatar
Timothy J. Baek committed
83
print(
Timothy J. Baek's avatar
Timothy J. Baek committed
84
    rf"""
Timothy J. Baek's avatar
Timothy J. Baek committed
85
86
87
88
89
90
91
92
93
94
95
96
97
  ___                    __        __   _     _   _ ___ 
 / _ \ _ __   ___ _ __   \ \      / /__| |__ | | | |_ _|
| | | | '_ \ / _ \ '_ \   \ \ /\ / / _ \ '_ \| | | || | 
| |_| | |_) |  __/ | | |   \ V  V /  __/ |_) | |_| || | 
 \___/| .__/ \___|_| |_|    \_/\_/ \___|_.__/ \___/|___|
      |_|                                               

      
v{VERSION} - building the best open-source AI user interface.      
https://github.com/open-webui/open-webui
"""
)

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

99
100
101
102
103
104
105
106
107
108
109
110
@asynccontextmanager
async def lifespan(app: FastAPI):
    if ENABLE_LITELLM:
        asyncio.create_task(start_litellm_background())
    yield
    if ENABLE_LITELLM:
        await shutdown_litellm_background()


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

112
113
114
app.state.config = AppConfig()
app.state.config.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.config.MODEL_FILTER_LIST = MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
115

116
app.state.MODEL_CONFIG = Models.get_all_models()
117

118
app.state.config.WEBHOOK_URL = WEBHOOK_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
119

Timothy J. Baek's avatar
Timothy J. Baek committed
120
121
origins = ["*"]

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

Timothy J. Baek's avatar
Timothy J. Baek committed
123
# Custom middleware to add security headers
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
124
125
126
127
128
129
# class SecurityHeadersMiddleware(BaseHTTPMiddleware):
#     async def dispatch(self, request: Request, call_next):
#         response: Response = await call_next(request)
#         response.headers["Cross-Origin-Opener-Policy"] = "same-origin"
#         response.headers["Cross-Origin-Embedder-Policy"] = "require-corp"
#         return response
Timothy J. Baek's avatar
Timothy J. Baek committed
130
131


Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
132
# app.add_middleware(SecurityHeadersMiddleware)
Timothy J. Baek's avatar
Timothy J. Baek committed
133
134


Timothy J. Baek's avatar
Timothy J. Baek committed
135
136
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
137
138
        return_citations = False

139
140
141
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
142
            log.debug(f"request.url.path: {request.url.path}")
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
143

Timothy J. Baek's avatar
Timothy J. Baek committed
144
145
146
147
148
149
150
            # 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 {}

151
152
153
154
            return_citations = data.get("citations", False)
            if "citations" in data:
                del data["citations"]

Timothy J. Baek's avatar
Timothy J. Baek committed
155
156
157
            # 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
158
                data = {**data}
159
                data["messages"], citations = rag_messages(
Timothy J. Baek's avatar
Timothy J. Baek committed
160
161
                    docs=data["docs"],
                    messages=data["messages"],
162
                    template=rag_app.state.config.RAG_TEMPLATE,
Timothy J. Baek's avatar
Timothy J. Baek committed
163
                    embedding_function=rag_app.state.EMBEDDING_FUNCTION,
164
                    k=rag_app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
165
                    reranking_function=rag_app.state.sentence_transformer_rf,
166
167
                    r=rag_app.state.config.RELEVANCE_THRESHOLD,
                    hybrid_search=rag_app.state.config.ENABLE_RAG_HYBRID_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
168
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
169
170
                del data["docs"]

171
172
173
                log.debug(
                    f"data['messages']: {data['messages']}, citations: {citations}"
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
174

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

Timothy J. Baek's avatar
Timothy J. Baek committed
177
178
179
180
181
182
183
184
185
186
187
188
            # 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
189
190

        response = await call_next(request)
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205

        if return_citations:
            # Inject the citations into the response
            if isinstance(response, StreamingResponse):
                # If it's a streaming response, inject it as SSE event or NDJSON line
                content_type = response.headers.get("Content-Type")
                if "text/event-stream" in content_type:
                    return StreamingResponse(
                        self.openai_stream_wrapper(response.body_iterator, citations),
                    )
                if "application/x-ndjson" in content_type:
                    return StreamingResponse(
                        self.ollama_stream_wrapper(response.body_iterator, citations),
                    )

Timothy J. Baek's avatar
Timothy J. Baek committed
206
207
208
209
210
        return response

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

211
212
213
214
215
216
217
218
219
220
    async def openai_stream_wrapper(self, original_generator, citations):
        yield f"data: {json.dumps({'citations': citations})}\n\n"
        async for data in original_generator:
            yield data

    async def ollama_stream_wrapper(self, original_generator, citations):
        yield f"{json.dumps({'citations': citations})}\n"
        async for data in original_generator:
            yield data

Timothy J. Baek's avatar
Timothy J. Baek committed
221
222
223
224

app.add_middleware(RAGMiddleware)


Timothy J. Baek's avatar
Timothy J. Baek committed
225
226
227
228
229
230
231
232
233
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
234
235
236
237
238
239
240
241
242
243
@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
244
245
246
247
248
249
@app.middleware("http")
async def update_embedding_function(request: Request, call_next):
    response = await call_next(request)
    if "/embedding/update" in request.url.path:
        webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION
    return response
Timothy J. Baek's avatar
Timothy J. Baek committed
250

Timothy J. Baek's avatar
Timothy J. Baek committed
251
252

app.mount("/litellm/api", litellm_app)
253
app.mount("/ollama", ollama_app)
254
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
255

Timothy J. Baek's avatar
Timothy J. Baek committed
256
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
257
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
258
259
app.mount("/rag/api/v1", rag_app)

Timothy J. Baek's avatar
Timothy J. Baek committed
260
261
262
263
app.mount("/api/v1", webui_app)

webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION

264

Timothy J. Baek's avatar
Timothy J. Baek committed
265
266
@app.get("/api/config")
async def get_app_config():
267
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA
268
269

    default_locale = "en-US"
270
271
272
273
    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
274
275
    return {
        "status": True,
276
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
277
        "version": VERSION,
278
        "auth": WEBUI_AUTH,
279
        "default_locale": default_locale,
280
281
282
        "images": images_app.state.config.ENABLED,
        "default_models": webui_app.state.config.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
283
        "trusted_header_auth": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
284
        "admin_export_enabled": ENABLE_ADMIN_EXPORT,
Timothy J. Baek's avatar
Timothy J. Baek committed
285
286
287
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
288
289
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
290
    return {
291
292
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
293
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
294
295
296
297
298
299
300
301


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
302
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
303
304
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
305
306
    app.state.config.ENABLE_MODEL_FILTER = form_data.enabled
    app.state.config.MODEL_FILTER_LIST = form_data.models
Timothy J. Baek's avatar
Timothy J. Baek committed
307

Timothy J. Baek's avatar
Timothy J. Baek committed
308
309
    ollama_app.state.config.ENABLE_MODEL_FILTER = app.state.config.ENABLE_MODEL_FILTER
    ollama_app.state.config.MODEL_FILTER_LIST = app.state.config.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
310

Timothy J. Baek's avatar
Timothy J. Baek committed
311
312
    openai_app.state.config.ENABLE_MODEL_FILTER = app.state.config.ENABLE_MODEL_FILTER
    openai_app.state.config.MODEL_FILTER_LIST = app.state.config.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
313

314
315
    litellm_app.state.ENABLE_MODEL_FILTER = app.state.config.ENABLE_MODEL_FILTER
    litellm_app.state.MODEL_FILTER_LIST = app.state.config.MODEL_FILTER_LIST
316

317
    return {
318
319
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
320
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
321
322


Timothy J. Baek's avatar
Timothy J. Baek committed
323
324
325
@app.get("/api/webhook")
async def get_webhook_url(user=Depends(get_admin_user)):
    return {
326
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
327
328
329
330
331
332
333
334
335
    }


class UrlForm(BaseModel):
    url: str


@app.post("/api/webhook")
async def update_webhook_url(form_data: UrlForm, user=Depends(get_admin_user)):
336
    app.state.config.WEBHOOK_URL = form_data.url
Timothy J. Baek's avatar
Timothy J. Baek committed
337

338
    webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
339
340

    return {
341
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
342
343
344
    }


345
346
347
348
349
350
351
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
352
353
@app.get("/api/changelog")
async def get_app_changelog():
354
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
355
356


357
358
359
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
360
361
362
363
364
365
366
367
368
369
        async with aiohttp.ClientSession() as session:
            async with session.get(
                "https://api.github.com/repos/open-webui/open-webui/releases/latest"
            ) as response:
                response.raise_for_status()
                data = await response.json()
                latest_version = data["tag_name"]

                return {"current": VERSION, "latest": latest_version[1:]}
    except aiohttp.ClientError as e:
370
371
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
372
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
373
374
        )

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

376
377
378
@app.get("/manifest.json")
async def get_manifest_json():
    return {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
379
380
        "name": WEBUI_NAME,
        "short_name": WEBUI_NAME,
381
382
383
384
385
        "start_url": "/",
        "display": "standalone",
        "background_color": "#343541",
        "theme_color": "#343541",
        "orientation": "portrait-primary",
Timothy J. Baek's avatar
Timothy J. Baek committed
386
        "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
387
388
    }

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

390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
@app.get("/opensearch.xml")
async def get_opensearch_xml():
    xml_content = rf"""
    <OpenSearchDescription xmlns="http://a9.com/-/spec/opensearch/1.1/" xmlns:moz="http://www.mozilla.org/2006/browser/search/">
    <ShortName>{WEBUI_NAME}</ShortName>
    <Description>Search {WEBUI_NAME}</Description>
    <InputEncoding>UTF-8</InputEncoding>
    <Image width="16" height="16" type="image/x-icon">{WEBUI_URL}/favicon.png</Image>
    <Url type="text/html" method="get" template="{WEBUI_URL}/?q={"{searchTerms}"}"/>
    <moz:SearchForm>{WEBUI_URL}</moz:SearchForm>
    </OpenSearchDescription>
    """
    return Response(content=xml_content, media_type="application/xml")


Timothy J. Baek's avatar
Timothy J. Baek committed
405
406
407
408
409
@app.get("/health")
async def healthcheck():
    return {"status": True}


410
411
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
412

413
if os.path.exists(FRONTEND_BUILD_DIR):
Timothy J. Baek's avatar
Timothy J. Baek committed
414
    mimetypes.add_type("text/javascript", ".js")
415
416
417
418
419
420
421
422
423
    app.mount(
        "/",
        SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
        name="spa-static-files",
    )
else:
    log.warning(
        f"Frontend build directory not found at '{FRONTEND_BUILD_DIR}'. Serving API only."
    )