main.py 13.7 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
Timothy J. Baek's avatar
Timothy J. Baek committed
11

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

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

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


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

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

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

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

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

70

Timothy J. Baek's avatar
Timothy J. Baek committed
71
72
73
74
75
76
77
78
79
80
81
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
82
print(
Timothy J. Baek's avatar
Timothy J. Baek committed
83
    rf"""
Timothy J. Baek's avatar
Timothy J. Baek committed
84
85
86
87
88
89
90
91
92
93
94
95
96
  ___                    __        __   _     _   _ ___ 
 / _ \ _ __   ___ _ __   \ \      / /__| |__ | | | |_ _|
| | | | '_ \ / _ \ '_ \   \ \ /\ / / _ \ '_ \| | | || | 
| |_| | |_) |  __/ | | |   \ 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
97

98
99
100
101
102
103
104
105
106
107
108
109
@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
110

111
112
113
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
114

115
app.state.MODEL_CONFIG = [model.to_form() for model in Models.get_all_models()]
116

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

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
122
# Custom middleware to add security headers
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
123
124
125
126
127
128
# 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
129
130


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


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

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

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

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

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

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

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

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

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

        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
205
206
207
208
209
        return response

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

210
211
212
213
214
215
216
217
218
219
    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
220
221
222
223

app.add_middleware(RAGMiddleware)


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


Timothy J. Baek's avatar
Timothy J. Baek committed
233
234
235
236
237
238
239
240
241
242
@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


243
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
244
245
app.mount("/litellm/api", litellm_app)

246
app.mount("/ollama", ollama_app)
247
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
248

Timothy J. Baek's avatar
Timothy J. Baek committed
249
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
250
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
251
252
app.mount("/rag/api/v1", rag_app)

253

Timothy J. Baek's avatar
Timothy J. Baek committed
254
255
@app.get("/api/config")
async def get_app_config():
256
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA
257
258

    default_locale = "en-US"
259
260
261
262
    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
263
264
    return {
        "status": True,
265
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
266
        "version": VERSION,
267
        "auth": WEBUI_AUTH,
268
        "default_locale": default_locale,
269
270
271
        "images": images_app.state.config.ENABLED,
        "default_models": webui_app.state.config.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
272
        "trusted_header_auth": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
273
        "admin_export_enabled": ENABLE_ADMIN_EXPORT,
Timothy J. Baek's avatar
Timothy J. Baek committed
274
275
276
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
277
278
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
279
    return {
280
281
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
282
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
283
284
285
286
287
288
289
290


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
291
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
292
293
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
294
295
    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
296

Timothy J. Baek's avatar
Timothy J. Baek committed
297
298
    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
299

Timothy J. Baek's avatar
Timothy J. Baek committed
300
301
    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
302

303
304
    litellm_app.state.ENABLE_MODEL_FILTER = app.state.config.ENABLE_MODEL_FILTER
    litellm_app.state.MODEL_FILTER_LIST = app.state.config.MODEL_FILTER_LIST
305

306
    return {
307
308
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
309
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
310
311


312
class SetModelConfigForm(BaseModel):
313
    models: List[ModelForm]
314
315
316
317
318
319


@app.post("/api/config/models")
async def update_model_config(
    form_data: SetModelConfigForm, user=Depends(get_admin_user)
):
320
321
322
323
324
    if not Models.update_all_models(form_data.models):
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=ERROR_MESSAGES.DEFAULT("Failed to update model config"),
        )
325

326
327
328
    ollama_app.state.MODEL_CONFIG = [
        model for model in form_data.models if model.source == "ollama"
    ]
329

330
331
332
    openai_app.state.MODEL_CONFIG = [
        model for model in form_data.models if model.source == "openai"
    ]
333

334
335
336
    litellm_app.state.MODEL_CONFIG = [
        model for model in form_data.models if model.source == "litellm"
    ]
337

338
    app.state.MODEL_CONFIG = [model for model in form_data.models]
339

340
    return {"models": app.state.MODEL_CONFIG}
341
342
343
344


@app.get("/api/config/models")
async def get_model_config(user=Depends(get_admin_user)):
345
    return {"models": app.state.MODEL_CONFIG}
346

347

Timothy J. Baek's avatar
Timothy J. Baek committed
348
349
350
@app.get("/api/webhook")
async def get_webhook_url(user=Depends(get_admin_user)):
    return {
351
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
352
353
354
355
356
357
358
359
360
    }


class UrlForm(BaseModel):
    url: str


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

363
    webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
364
365

    return {
366
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
367
368
369
    }


370
371
372
373
374
375
376
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
377
378
@app.get("/api/changelog")
async def get_app_changelog():
379
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
380
381


382
383
384
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
385
386
387
388
389
390
391
392
393
394
        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:
395
396
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
397
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
398
399
        )

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

401
402
403
@app.get("/manifest.json")
async def get_manifest_json():
    return {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
404
405
        "name": WEBUI_NAME,
        "short_name": WEBUI_NAME,
406
407
408
409
410
        "start_url": "/",
        "display": "standalone",
        "background_color": "#343541",
        "theme_color": "#343541",
        "orientation": "portrait-primary",
Timothy J. Baek's avatar
Timothy J. Baek committed
411
        "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
412
413
    }

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

415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
@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
430
431
432
433
434
@app.get("/health")
async def healthcheck():
    return {"status": True}


435
436
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
437

438
439
440
441
442
443
444
445
446
447
if os.path.exists(FRONTEND_BUILD_DIR):
    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."
    )