main.py 15.8 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

22
23
from apps.ollama.main import app as ollama_app, get_all_models as get_ollama_models
from apps.openai.main import app as openai_app, get_all_models as get_openai_models
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
42
from utils.utils import get_admin_user, get_verified_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
57
    ENABLE_OPENAI_API,
    ENABLE_OLLAMA_API,
58
    ENABLE_LITELLM,
Timothy J. Baek's avatar
Timothy J. Baek committed
59
    ENABLE_MODEL_FILTER,
60
    MODEL_FILTER_LIST,
61
62
    GLOBAL_LOG_LEVEL,
    SRC_LOG_LEVELS,
Timothy J. Baek's avatar
Timothy J. Baek committed
63
    WEBHOOK_URL,
64
    ENABLE_ADMIN_EXPORT,
65
    AppConfig,
66
)
67
68
from constants import ERROR_MESSAGES

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

73

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

101
102
103
104
105
106
107
108
109
110
111
112
@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
113

114
app.state.config = AppConfig()
115
116
117
118

app.state.config.ENABLE_OPENAI_API = ENABLE_OPENAI_API
app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API

119
120
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
121

122

123
app.state.config.WEBHOOK_URL = WEBHOOK_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
124

Timothy J. Baek's avatar
Timothy J. Baek committed
125
126
127

app.state.MODELS = {}

Timothy J. Baek's avatar
Timothy J. Baek committed
128
129
origins = ["*"]

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

Timothy J. Baek's avatar
Timothy J. Baek committed
131
# Custom middleware to add security headers
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
132
133
134
135
136
137
# 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
138
139


Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
140
# app.add_middleware(SecurityHeadersMiddleware)
Timothy J. Baek's avatar
Timothy J. Baek committed
141
142


Timothy J. Baek's avatar
Timothy J. Baek committed
143
144
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
145
146
        return_citations = False

147
148
149
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
150
            log.debug(f"request.url.path: {request.url.path}")
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
151

Timothy J. Baek's avatar
Timothy J. Baek committed
152
153
154
155
156
157
158
            # 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 {}

159
160
161
162
            return_citations = data.get("citations", False)
            if "citations" in data:
                del data["citations"]

Timothy J. Baek's avatar
Timothy J. Baek committed
163
164
165
            # 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
166
                data = {**data}
167
                data["messages"], citations = rag_messages(
Timothy J. Baek's avatar
Timothy J. Baek committed
168
169
                    docs=data["docs"],
                    messages=data["messages"],
170
                    template=rag_app.state.config.RAG_TEMPLATE,
Timothy J. Baek's avatar
Timothy J. Baek committed
171
                    embedding_function=rag_app.state.EMBEDDING_FUNCTION,
172
                    k=rag_app.state.config.TOP_K,
Timothy J. Baek's avatar
Timothy J. Baek committed
173
                    reranking_function=rag_app.state.sentence_transformer_rf,
174
175
                    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
176
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
177
178
                del data["docs"]

179
180
181
                log.debug(
                    f"data['messages']: {data['messages']}, citations: {citations}"
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
182

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

Timothy J. Baek's avatar
Timothy J. Baek committed
185
186
187
188
189
190
191
192
193
194
195
196
            # 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
197
198

        response = await call_next(request)
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213

        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
214
215
216
217
218
        return response

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

219
220
221
222
223
224
225
226
227
228
    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
229
230
231
232

app.add_middleware(RAGMiddleware)


Timothy J. Baek's avatar
Timothy J. Baek committed
233
234
235
236
237
238
239
240
241
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
242
243
@app.middleware("http")
async def check_url(request: Request, call_next):
Timothy J. Baek's avatar
Timothy J. Baek committed
244
245
246
247
248
    if len(app.state.MODELS) == 0:
        await get_all_models()
    else:
        pass

Timothy J. Baek's avatar
Timothy J. Baek committed
249
250
251
252
253
254
255
256
    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
257
258
259
260
261
262
@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
263

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

265
# TODO: Deprecate LiteLLM
Timothy J. Baek's avatar
Timothy J. Baek committed
266
app.mount("/litellm/api", litellm_app)
267

268
app.mount("/ollama", ollama_app)
269
app.mount("/openai", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
270

Timothy J. Baek's avatar
Timothy J. Baek committed
271
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
272
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
273
274
app.mount("/rag/api/v1", rag_app)

Timothy J. Baek's avatar
Timothy J. Baek committed
275
276
277
278
app.mount("/api/v1", webui_app)

webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION

279

Timothy J. Baek's avatar
Timothy J. Baek committed
280
async def get_all_models():
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
    openai_models = []
    ollama_models = []

    if app.state.config.ENABLE_OPENAI_API:
        openai_models = await get_openai_models()

        openai_models = openai_models["data"]

    if app.state.config.ENABLE_OLLAMA_API:
        ollama_models = await get_ollama_models()

        ollama_models = [
            {
                "id": model["model"],
                "name": model["name"],
                "object": "model",
                "created": int(time.time()),
                "owned_by": "ollama",
                "ollama": model,
            }
            for model in ollama_models["models"]
        ]

    models = openai_models + ollama_models
    custom_models = Models.get_all_models()

    for custom_model in custom_models:
        if custom_model.base_model_id == None:
            for model in models:
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
310
311
312
313
                if (
                    custom_model.id == model["id"]
                    or custom_model.id == model["id"].split(":")[0]
                ):
314
315
316
                    model["name"] = custom_model.name
                    model["info"] = custom_model.model_dump()
        else:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
317
318
            owned_by = "openai"
            for model in models:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
319
320
321
322
                if (
                    custom_model.base_model_id == model["id"]
                    or custom_model.base_model_id == model["id"].split(":")[0]
                ):
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
323
324
325
                    owned_by = model["owned_by"]
                    break

326
327
328
329
330
331
            models.append(
                {
                    "id": custom_model.id,
                    "name": custom_model.name,
                    "object": "model",
                    "created": custom_model.created_at,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
332
                    "owned_by": owned_by,
333
                    "info": custom_model.model_dump(),
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
334
                    "preset": True,
335
336
337
                }
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
338
339
340
341
342
343
344
345
346
347
    app.state.MODELS = {model["id"]: model for model in models}

    webui_app.state.MODELS = app.state.MODELS

    return models


@app.get("/api/models")
async def get_models(user=Depends(get_verified_user)):
    models = await get_all_models()
348
349
350
351
352
353
354
355
356
357
358
359
360
    if app.state.config.ENABLE_MODEL_FILTER:
        if user.role == "user":
            models = list(
                filter(
                    lambda model: model["id"] in app.state.config.MODEL_FILTER_LIST,
                    models,
                )
            )
            return {"data": models}

    return {"data": models}


Timothy J. Baek's avatar
Timothy J. Baek committed
361
362
@app.get("/api/config")
async def get_app_config():
363
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA
364
365

    default_locale = "en-US"
366
367
368
369
    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
370
371
    return {
        "status": True,
372
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
373
        "version": VERSION,
374
        "auth": WEBUI_AUTH,
375
        "default_locale": default_locale,
376
377
378
        "images": images_app.state.config.ENABLED,
        "default_models": webui_app.state.config.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
379
        "trusted_header_auth": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
380
        "admin_export_enabled": ENABLE_ADMIN_EXPORT,
Timothy J. Baek's avatar
Timothy J. Baek committed
381
382
383
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
384
385
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
386
    return {
387
388
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
389
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
390
391
392
393
394
395
396
397


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
398
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
399
400
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
401
402
    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
403

Timothy J. Baek's avatar
Timothy J. Baek committed
404
405
    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
406

Timothy J. Baek's avatar
Timothy J. Baek committed
407
408
    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
409

410
411
    litellm_app.state.ENABLE_MODEL_FILTER = app.state.config.ENABLE_MODEL_FILTER
    litellm_app.state.MODEL_FILTER_LIST = app.state.config.MODEL_FILTER_LIST
412

413
    return {
414
415
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
416
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
417
418


Timothy J. Baek's avatar
Timothy J. Baek committed
419
420
421
@app.get("/api/webhook")
async def get_webhook_url(user=Depends(get_admin_user)):
    return {
422
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
423
424
425
426
427
428
429
430
431
    }


class UrlForm(BaseModel):
    url: str


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

434
    webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
Timothy J. Baek's avatar
Timothy J. Baek committed
435
436

    return {
437
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
438
439
440
    }


441
442
443
444
445
446
447
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
448
449
@app.get("/api/changelog")
async def get_app_changelog():
450
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
451
452


453
454
455
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
456
457
458
459
460
461
462
463
464
465
        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:
466
467
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
468
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
469
470
        )

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

472
473
474
@app.get("/manifest.json")
async def get_manifest_json():
    return {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
475
476
        "name": WEBUI_NAME,
        "short_name": WEBUI_NAME,
477
478
479
480
481
        "start_url": "/",
        "display": "standalone",
        "background_color": "#343541",
        "theme_color": "#343541",
        "orientation": "portrait-primary",
Timothy J. Baek's avatar
Timothy J. Baek committed
482
        "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
483
484
    }

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

486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
@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
501
502
503
504
505
@app.get("/health")
async def healthcheck():
    return {"status": True}


506
507
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
508

509
if os.path.exists(FRONTEND_BUILD_DIR):
Timothy J. Baek's avatar
Timothy J. Baek committed
510
    mimetypes.add_type("text/javascript", ".js")
511
512
513
514
515
516
517
518
519
    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."
    )