main.py 11.8 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
Timothy J. Baek's avatar
Timothy J. Baek committed
8
import aiohttp
9
import requests
Timothy J. Baek's avatar
Timothy J. Baek committed
10

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

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

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


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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
39
40

from utils.utils import get_admin_user
Timothy J. Baek's avatar
Timothy J. Baek committed
41
from apps.rag.utils import rag_messages
Timothy J. Baek's avatar
Timothy J. Baek committed
42

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

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

69

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

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
98
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
99
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
100

Timothy J. Baek's avatar
Timothy J. Baek committed
101
102
app.state.WEBHOOK_URL = WEBHOOK_URL

Timothy J. Baek's avatar
Timothy J. Baek committed
103
104
origins = ["*"]

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

Timothy J. Baek's avatar
Timothy J. Baek committed
106
107
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
108
109
        return_citations = False

110
111
112
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
113
            log.debug(f"request.url.path: {request.url.path}")
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
114

Timothy J. Baek's avatar
Timothy J. Baek committed
115
116
117
118
119
120
121
            # 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 {}

122
123
124
125
            return_citations = data.get("citations", False)
            if "citations" in data:
                del data["citations"]

Timothy J. Baek's avatar
Timothy J. Baek committed
126
127
128
            # 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
129
                data = {**data}
130
                data["messages"], citations = rag_messages(
Timothy J. Baek's avatar
Timothy J. Baek committed
131
132
133
134
135
136
137
138
                    docs=data["docs"],
                    messages=data["messages"],
                    template=rag_app.state.RAG_TEMPLATE,
                    embedding_function=rag_app.state.EMBEDDING_FUNCTION,
                    k=rag_app.state.TOP_K,
                    reranking_function=rag_app.state.sentence_transformer_rf,
                    r=rag_app.state.RELEVANCE_THRESHOLD,
                    hybrid_search=rag_app.state.ENABLE_RAG_HYBRID_SEARCH,
Timothy J. Baek's avatar
Timothy J. Baek committed
139
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
140
141
                del data["docs"]

142
143
144
                log.debug(
                    f"data['messages']: {data['messages']}, citations: {citations}"
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
145

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

Timothy J. Baek's avatar
Timothy J. Baek committed
148
149
150
151
152
153
154
155
156
157
158
159
            # 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
160
161

        response = await call_next(request)
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

        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
177
178
179
180
181
        return response

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

182
183
184
185
186
187
188
189
190
191
    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
192
193
194
195

app.add_middleware(RAGMiddleware)


Timothy J. Baek's avatar
Timothy J. Baek committed
196
197
198
199
200
201
202
203
204
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
207
208
209
210
211
212
213
214
@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
215
216
@app.on_event("startup")
async def on_startup():
217
218
    if ENABLE_LITELLM:
        asyncio.create_task(start_litellm_background())
Timothy J. Baek's avatar
Timothy J. Baek committed
219
220


221
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
222
223
app.mount("/litellm/api", litellm_app)

224
app.mount("/ollama", ollama_app)
225
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
226

Timothy J. Baek's avatar
Timothy J. Baek committed
227
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
228
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
229
230
app.mount("/rag/api/v1", rag_app)

231

Timothy J. Baek's avatar
Timothy J. Baek committed
232
233
@app.get("/api/config")
async def get_app_config():
234
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA
235
236

    default_locale = "en-US"
237
238
239
240
    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
241
242
    return {
        "status": True,
243
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
244
        "version": VERSION,
245
        "auth": WEBUI_AUTH,
246
        "default_locale": default_locale,
Timothy J. Baek's avatar
Timothy J. Baek committed
247
248
        "images": images_app.state.ENABLED,
        "default_models": webui_app.state.DEFAULT_MODELS,
249
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
250
        "trusted_header_auth": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
251
        "admin_export_enabled": ENABLE_ADMIN_EXPORT,
252
        "websearch": RAG_WEB_SEARCH_ENABLED,
Timothy J. Baek's avatar
Timothy J. Baek committed
253
254
255
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
256
257
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
258
    return {
Timothy J. Baek's avatar
Timothy J. Baek committed
259
        "enabled": app.state.ENABLE_MODEL_FILTER,
260
261
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
262
263
264
265
266
267
268
269


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
270
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
271
272
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
Timothy J. Baek committed
273
    app.state.ENABLE_MODEL_FILTER = form_data.enabled
274
    app.state.MODEL_FILTER_LIST = form_data.models
Timothy J. Baek's avatar
Timothy J. Baek committed
275

Timothy J. Baek's avatar
Timothy J. Baek committed
276
    ollama_app.state.ENABLE_MODEL_FILTER = app.state.ENABLE_MODEL_FILTER
277
    ollama_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
278

Timothy J. Baek's avatar
Timothy J. Baek committed
279
    openai_app.state.ENABLE_MODEL_FILTER = app.state.ENABLE_MODEL_FILTER
280
    openai_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
281

Timothy J. Baek's avatar
Timothy J. Baek committed
282
    litellm_app.state.ENABLE_MODEL_FILTER = app.state.ENABLE_MODEL_FILTER
283
284
    litellm_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST

285
    return {
Timothy J. Baek's avatar
Timothy J. Baek committed
286
        "enabled": app.state.ENABLE_MODEL_FILTER,
287
288
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
289
290


Timothy J. Baek's avatar
Timothy J. Baek committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
@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,
    }


313
314
315
316
317
318
319
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
320
321
@app.get("/api/changelog")
async def get_app_changelog():
322
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
323
324


325
326
327
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
328
329
330
331
332
333
334
335
336
337
        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:
338
339
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
340
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
341
342
        )

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

344
345
346
@app.get("/manifest.json")
async def get_manifest_json():
    return {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
347
348
        "name": WEBUI_NAME,
        "short_name": WEBUI_NAME,
349
350
351
352
353
        "start_url": "/",
        "display": "standalone",
        "background_color": "#343541",
        "theme_color": "#343541",
        "orientation": "portrait-primary",
Timothy J. Baek's avatar
Timothy J. Baek committed
354
        "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
355
356
    }

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

358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
@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")


373
374
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
375

376
377
378
379
380
381
382
383
384
385
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."
    )
386
387
388
389


@app.on_event("shutdown")
async def shutdown_event():
390
391
    if ENABLE_LITELLM:
        await shutdown_litellm_background()