main.py 30.6 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
14
import shutil
import os
import asyncio
Timothy J. Baek's avatar
Timothy J. Baek committed
15

Timothy J. Baek's avatar
Timothy J. Baek committed
16
from fastapi import FastAPI, Request, Depends, status, UploadFile, File, Form
Timothy J. Baek's avatar
Timothy J. Baek committed
17
from fastapi.staticfiles import StaticFiles
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
18
from fastapi.responses import JSONResponse
Timothy J. Baek's avatar
Timothy J. Baek committed
19
20
21
from fastapi import HTTPException
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.middleware.cors import CORSMiddleware
22
from starlette.exceptions import HTTPException as StarletteHTTPException
Timothy J. Baek's avatar
Timothy J. Baek committed
23
from starlette.middleware.base import BaseHTTPMiddleware
24
from starlette.responses import StreamingResponse, Response
Timothy J. Baek's avatar
Timothy J. Baek committed
25

Timothy J. Baek's avatar
Timothy J. Baek committed
26
27

from apps.socket.main import app as socket_app
28
29
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
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.webui.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

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.webui.models.models import Models, ModelModel
41
42
43
44
45
46
from utils.utils import (
    get_admin_user,
    get_verified_user,
    get_current_user,
    get_http_authorization_cred,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
47
from apps.rag.utils import rag_messages
Timothy J. Baek's avatar
Timothy J. Baek committed
48

49
from config import (
50
    CONFIG_DATA,
51
    WEBUI_NAME,
52
    WEBUI_URL,
53
    WEBUI_AUTH,
54
55
56
57
    ENV,
    VERSION,
    CHANGELOG,
    FRONTEND_BUILD_DIR,
58
59
    CACHE_DIR,
    STATIC_DIR,
60
61
    ENABLE_OPENAI_API,
    ENABLE_OLLAMA_API,
Timothy J. Baek's avatar
Timothy J. Baek committed
62
    ENABLE_MODEL_FILTER,
63
    MODEL_FILTER_LIST,
64
65
    GLOBAL_LOG_LEVEL,
    SRC_LOG_LEVELS,
Timothy J. Baek's avatar
Timothy J. Baek committed
66
    WEBHOOK_URL,
67
    ENABLE_ADMIN_EXPORT,
68
    AppConfig,
69
    WEBUI_BUILD_HASH,
70
)
71
72
from constants import ERROR_MESSAGES

73
74
75
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
76

77

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

      
99
v{VERSION} - building the best open-source AI user interface.
100
{f"Commit: {WEBUI_BUILD_HASH}" if WEBUI_BUILD_HASH != "dev-build" else ""}
Timothy J. Baek's avatar
Timothy J. Baek committed
101
102
103
104
https://github.com/open-webui/open-webui
"""
)

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

106
107
108
109
110
111
112
113
@asynccontextmanager
async def lifespan(app: FastAPI):
    yield


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

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

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

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

123

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

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

app.state.MODELS = {}

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

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
        if request.method == "POST" and (
148
149
            "/ollama/api/chat" in request.url.path
            or "/chat/completions" in request.url.path
150
        ):
151
            log.debug(f"request.url.path: {request.url.path}")
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
152

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

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

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

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

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

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

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

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

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

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

app.add_middleware(RAGMiddleware)


Timothy J. Baek's avatar
Timothy J. Baek committed
234
235
236
class PipelineMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if request.method == "POST" and (
237
238
            "/ollama/api/chat" in request.url.path
            or "/chat/completions" in request.url.path
Timothy J. Baek's avatar
Timothy J. Baek committed
239
240
241
242
243
244
245
246
247
248
249
        ):
            log.debug(f"request.url.path: {request.url.path}")

            # 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 {}

            model_id = data["model"]
250
            filters = [
Timothy J. Baek's avatar
Timothy J. Baek committed
251
252
253
                model
                for model in app.state.MODELS.values()
                if "pipeline" in model
Timothy J. Baek's avatar
Timothy J. Baek committed
254
                and "type" in model["pipeline"]
255
                and model["pipeline"]["type"] == "filter"
256
257
258
                and (
                    model["pipeline"]["pipelines"] == ["*"]
                    or any(
259
260
                        model_id == target_model_id
                        for target_model_id in model["pipeline"]["pipelines"]
261
262
                    )
                )
Timothy J. Baek's avatar
Timothy J. Baek committed
263
            ]
264
            sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])
Timothy J. Baek's avatar
Timothy J. Baek committed
265

266
            user = None
267
            if len(sorted_filters) > 0:
268
269
270
271
272
273
274
275
276
277
                try:
                    user = get_current_user(
                        get_http_authorization_cred(
                            request.headers.get("Authorization")
                        )
                    )
                    user = {"id": user.id, "name": user.name, "role": user.role}
                except:
                    pass

Timothy J. Baek's avatar
Timothy J. Baek committed
278
279
280
281
282
            model = app.state.MODELS[model_id]

            if "pipeline" in model:
                sorted_filters.append(model)

283
            for filter in sorted_filters:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
284
                r = None
Timothy J. Baek's avatar
Timothy J. Baek committed
285
                try:
286
                    urlIdx = filter["urlIdx"]
Timothy J. Baek's avatar
Timothy J. Baek committed
287
288
289
290
291
292
293

                    url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
                    key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]

                    if key != "":
                        headers = {"Authorization": f"Bearer {key}"}
                        r = requests.post(
Timothy J. Baek's avatar
Timothy J. Baek committed
294
                            f"{url}/{filter['id']}/filter/inlet",
Timothy J. Baek's avatar
Timothy J. Baek committed
295
296
                            headers=headers,
                            json={
297
                                "user": user,
Timothy J. Baek's avatar
Timothy J. Baek committed
298
299
300
301
302
303
304
305
                                "body": data,
                            },
                        )

                        r.raise_for_status()
                        data = r.json()
                except Exception as e:
                    # Handle connection error here
306
                    print(f"Connection error: {e}")
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
307
308
309
310
311
312
313
314
315
316
317
318
319
320

                    if r is not None:
                        try:
                            res = r.json()
                            if "detail" in res:
                                return JSONResponse(
                                    status_code=r.status_code,
                                    content=res,
                                )
                        except:
                            pass

                    else:
                        pass
Timothy J. Baek's avatar
Timothy J. Baek committed
321

Timothy J. Baek's avatar
Timothy J. Baek committed
322
323
324
325
326
327
            if "pipeline" not in app.state.MODELS[model_id]:
                if "chat_id" in data:
                    del data["chat_id"]

                if "title" in data:
                    del data["title"]
328

Timothy J. Baek's avatar
Timothy J. Baek committed
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
            modified_body_bytes = json.dumps(data).encode("utf-8")
            # 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"
                ],
            ]

        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(PipelineMiddleware)


Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
352
353
354
355
356
357
358
359
360
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
361
362
@app.middleware("http")
async def check_url(request: Request, call_next):
Timothy J. Baek's avatar
Timothy J. Baek committed
363
364
365
366
367
    if len(app.state.MODELS) == 0:
        await get_all_models()
    else:
        pass

Timothy J. Baek's avatar
Timothy J. Baek committed
368
369
370
371
372
373
374
375
    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
376
377
378
379
380
381
@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
382

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

Timothy J. Baek's avatar
Timothy J. Baek committed
384
385
386
app.mount("/ws", socket_app)


387
app.mount("/ollama", ollama_app)
388
app.mount("/openai", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
389

Timothy J. Baek's avatar
Timothy J. Baek committed
390
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
391
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
392
393
app.mount("/rag/api/v1", rag_app)

Timothy J. Baek's avatar
Timothy J. Baek committed
394
395
396
397
app.mount("/api/v1", webui_app)

webui_app.state.EMBEDDING_FUNCTION = rag_app.state.EMBEDDING_FUNCTION

398

Timothy J. Baek's avatar
Timothy J. Baek committed
399
async def get_all_models():
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
    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
429
430
431
432
                if (
                    custom_model.id == model["id"]
                    or custom_model.id == model["id"].split(":")[0]
                ):
433
434
435
                    model["name"] = custom_model.name
                    model["info"] = custom_model.model_dump()
        else:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
436
437
            owned_by = "openai"
            for model in models:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
438
439
440
441
                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
442
443
444
                    owned_by = model["owned_by"]
                    break

445
446
447
448
449
450
            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
451
                    "owned_by": owned_by,
452
                    "info": custom_model.model_dump(),
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
453
                    "preset": True,
454
455
456
                }
            )

Timothy J. Baek's avatar
Timothy J. Baek committed
457
458
459
460
461
462
463
464
465
466
    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()
Timothy J. Baek's avatar
Timothy J. Baek committed
467

468
    # Filter out filter pipelines
Timothy J. Baek's avatar
Timothy J. Baek committed
469
470
471
    models = [
        model
        for model in models
Timothy J. Baek's avatar
Timothy J. Baek committed
472
        if "pipeline" not in model or model["pipeline"].get("type", None) != "filter"
Timothy J. Baek's avatar
Timothy J. Baek committed
473
474
    ]

475
476
477
478
479
480
481
482
483
484
485
486
487
    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}


488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
@app.post("/api/chat/completed")
async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
    data = form_data
    model_id = data["model"]

    filters = [
        model
        for model in app.state.MODELS.values()
        if "pipeline" in model
        and "type" in model["pipeline"]
        and model["pipeline"]["type"] == "filter"
        and (
            model["pipeline"]["pipelines"] == ["*"]
            or any(
                model_id == target_model_id
                for target_model_id in model["pipeline"]["pipelines"]
            )
        )
    ]
    sorted_filters = sorted(filters, key=lambda x: x["pipeline"]["priority"])

509
510
    print(model_id)

Timothy J. Baek's avatar
Timothy J. Baek committed
511
512
513
514
    if model_id in app.state.MODELS:
        model = app.state.MODELS[model_id]
        if "pipeline" in model:
            sorted_filters = [model] + sorted_filters
Timothy J. Baek's avatar
Timothy J. Baek committed
515

516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
    for filter in sorted_filters:
        r = None
        try:
            urlIdx = filter["urlIdx"]

            url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
            key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]

            if key != "":
                headers = {"Authorization": f"Bearer {key}"}
                r = requests.post(
                    f"{url}/{filter['id']}/filter/outlet",
                    headers=headers,
                    json={
                        "user": {"id": user.id, "name": user.name, "role": user.role},
                        "body": data,
                    },
                )

                r.raise_for_status()
                data = r.json()
        except Exception as e:
            # Handle connection error here
            print(f"Connection error: {e}")

            if r is not None:
                try:
                    res = r.json()
                    if "detail" in res:
                        return JSONResponse(
                            status_code=r.status_code,
                            content=res,
                        )
                except:
                    pass

            else:
                pass

    return data


558
559
@app.get("/api/pipelines/list")
async def get_pipelines_list(user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
560
561
562
    responses = await get_openai_models(raw=True)

    print(responses)
Timothy J. Baek's avatar
Timothy J. Baek committed
563
564
565
566
567
    urlIdxs = [
        idx
        for idx, response in enumerate(responses)
        if response != None and "pipelines" in response
    ]
568
569
570
571
572
573
574
575
576
577
578
579

    return {
        "data": [
            {
                "url": openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx],
                "idx": urlIdx,
            }
            for urlIdx in urlIdxs
        ]
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
@app.post("/api/pipelines/upload")
async def upload_pipeline(
    urlIdx: int = Form(...), file: UploadFile = File(...), user=Depends(get_admin_user)
):
    print("upload_pipeline", urlIdx, file.filename)
    # Check if the uploaded file is a python file
    if not file.filename.endswith(".py"):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Only Python (.py) files are allowed.",
        )

    upload_folder = f"{CACHE_DIR}/pipelines"
    os.makedirs(upload_folder, exist_ok=True)
    file_path = os.path.join(upload_folder, file.filename)

    try:
        # Save the uploaded file
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]

        headers = {"Authorization": f"Bearer {key}"}

        with open(file_path, "rb") as f:
            files = {"file": f}
            r = requests.post(f"{url}/pipelines/upload", headers=headers, files=files)

        r.raise_for_status()
        data = r.json()

        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")

        detail = "Pipeline not found"
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass

        raise HTTPException(
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
            detail=detail,
        )
    finally:
        # Ensure the file is deleted after the upload is completed or on failure
        if os.path.exists(file_path):
            os.remove(file_path)


Timothy J. Baek's avatar
Timothy J. Baek committed
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
class AddPipelineForm(BaseModel):
    url: str
    urlIdx: int


@app.post("/api/pipelines/add")
async def add_pipeline(form_data: AddPipelineForm, user=Depends(get_admin_user)):

    r = None
    try:
        urlIdx = form_data.urlIdx

        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]

        headers = {"Authorization": f"Bearer {key}"}
        r = requests.post(
            f"{url}/pipelines/add", headers=headers, json={"url": form_data.url}
        )

        r.raise_for_status()
        data = r.json()

        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")

        detail = "Pipeline not found"
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass

        raise HTTPException(
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
            detail=detail,
        )


class DeletePipelineForm(BaseModel):
    id: str
    urlIdx: int


@app.delete("/api/pipelines/delete")
async def delete_pipeline(form_data: DeletePipelineForm, user=Depends(get_admin_user)):

    r = None
    try:
        urlIdx = form_data.urlIdx

        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]

        headers = {"Authorization": f"Bearer {key}"}
        r = requests.delete(
            f"{url}/pipelines/delete", headers=headers, json={"id": form_data.id}
        )

        r.raise_for_status()
        data = r.json()

        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")

        detail = "Pipeline not found"
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass

        raise HTTPException(
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
            detail=detail,
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
723
@app.get("/api/pipelines")
724
async def get_pipelines(urlIdx: Optional[int] = None, user=Depends(get_admin_user)):
Timothy J. Baek's avatar
Timothy J. Baek committed
725
726
727
    r = None
    try:
        urlIdx
728

Timothy J. Baek's avatar
Timothy J. Baek committed
729
730
        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
731

Timothy J. Baek's avatar
Timothy J. Baek committed
732
733
        headers = {"Authorization": f"Bearer {key}"}
        r = requests.get(f"{url}/pipelines", headers=headers)
734

Timothy J. Baek's avatar
Timothy J. Baek committed
735
736
        r.raise_for_status()
        data = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
737

Timothy J. Baek's avatar
Timothy J. Baek committed
738
739
740
741
        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")
Timothy J. Baek's avatar
Timothy J. Baek committed
742

Timothy J. Baek's avatar
Timothy J. Baek committed
743
744
745
746
747
748
749
750
751
752
753
754
755
        detail = "Pipeline not found"
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass

        raise HTTPException(
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
            detail=detail,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
756
757


Timothy J. Baek's avatar
Timothy J. Baek committed
758
759
760
761
762
763
764
@app.get("/api/pipelines/{pipeline_id}/valves")
async def get_pipeline_valves(
    urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
):
    models = await get_all_models()
    r = None
    try:
Timothy J. Baek's avatar
Timothy J. Baek committed
765

Timothy J. Baek's avatar
Timothy J. Baek committed
766
767
        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
Timothy J. Baek's avatar
Timothy J. Baek committed
768

Timothy J. Baek's avatar
Timothy J. Baek committed
769
770
        headers = {"Authorization": f"Bearer {key}"}
        r = requests.get(f"{url}/{pipeline_id}/valves", headers=headers)
Timothy J. Baek's avatar
Timothy J. Baek committed
771

Timothy J. Baek's avatar
Timothy J. Baek committed
772
773
        r.raise_for_status()
        data = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
774

Timothy J. Baek's avatar
Timothy J. Baek committed
775
776
777
778
        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")
779

Timothy J. Baek's avatar
Timothy J. Baek committed
780
        detail = "Pipeline not found"
781

Timothy J. Baek's avatar
Timothy J. Baek committed
782
783
784
785
786
787
788
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass
Timothy J. Baek's avatar
Timothy J. Baek committed
789
790

        raise HTTPException(
Timothy J. Baek's avatar
Timothy J. Baek committed
791
792
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
            detail=detail,
Timothy J. Baek's avatar
Timothy J. Baek committed
793
794
795
796
        )


@app.get("/api/pipelines/{pipeline_id}/valves/spec")
Timothy J. Baek's avatar
Timothy J. Baek committed
797
798
799
async def get_pipeline_valves_spec(
    urlIdx: Optional[int], pipeline_id: str, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
Timothy J. Baek committed
800
801
    models = await get_all_models()

Timothy J. Baek's avatar
Timothy J. Baek committed
802
803
804
805
    r = None
    try:
        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
Timothy J. Baek's avatar
Timothy J. Baek committed
806

Timothy J. Baek's avatar
Timothy J. Baek committed
807
808
        headers = {"Authorization": f"Bearer {key}"}
        r = requests.get(f"{url}/{pipeline_id}/valves/spec", headers=headers)
Timothy J. Baek's avatar
Timothy J. Baek committed
809

Timothy J. Baek's avatar
Timothy J. Baek committed
810
811
        r.raise_for_status()
        data = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
812

Timothy J. Baek's avatar
Timothy J. Baek committed
813
814
815
816
        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")
Timothy J. Baek's avatar
Timothy J. Baek committed
817

Timothy J. Baek's avatar
Timothy J. Baek committed
818
819
820
821
822
823
824
825
        detail = "Pipeline not found"
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass
826

Timothy J. Baek's avatar
Timothy J. Baek committed
827
        raise HTTPException(
Timothy J. Baek's avatar
Timothy J. Baek committed
828
829
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
            detail=detail,
Timothy J. Baek's avatar
Timothy J. Baek committed
830
831
832
833
834
        )


@app.post("/api/pipelines/{pipeline_id}/valves/update")
async def update_pipeline_valves(
Timothy J. Baek's avatar
Timothy J. Baek committed
835
836
837
838
    urlIdx: Optional[int],
    pipeline_id: str,
    form_data: dict,
    user=Depends(get_admin_user),
Timothy J. Baek's avatar
Timothy J. Baek committed
839
840
841
):
    models = await get_all_models()

Timothy J. Baek's avatar
Timothy J. Baek committed
842
843
844
845
    r = None
    try:
        url = openai_app.state.config.OPENAI_API_BASE_URLS[urlIdx]
        key = openai_app.state.config.OPENAI_API_KEYS[urlIdx]
Timothy J. Baek's avatar
Timothy J. Baek committed
846

Timothy J. Baek's avatar
Timothy J. Baek committed
847
848
849
850
851
852
        headers = {"Authorization": f"Bearer {key}"}
        r = requests.post(
            f"{url}/{pipeline_id}/valves/update",
            headers=headers,
            json={**form_data},
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
853

Timothy J. Baek's avatar
Timothy J. Baek committed
854
855
        r.raise_for_status()
        data = r.json()
Timothy J. Baek's avatar
Timothy J. Baek committed
856

Timothy J. Baek's avatar
Timothy J. Baek committed
857
858
859
860
        return {**data}
    except Exception as e:
        # Handle connection error here
        print(f"Connection error: {e}")
Timothy J. Baek's avatar
Timothy J. Baek committed
861

Timothy J. Baek's avatar
Timothy J. Baek committed
862
        detail = "Pipeline not found"
863

Timothy J. Baek's avatar
Timothy J. Baek committed
864
865
866
867
868
869
870
        if r is not None:
            try:
                res = r.json()
                if "detail" in res:
                    detail = res["detail"]
            except:
                pass
871

Timothy J. Baek's avatar
Timothy J. Baek committed
872
        raise HTTPException(
Timothy J. Baek's avatar
Timothy J. Baek committed
873
            status_code=(r.status_code if r is not None else status.HTTP_404_NOT_FOUND),
Timothy J. Baek's avatar
Timothy J. Baek committed
874
            detail=detail,
Timothy J. Baek's avatar
Timothy J. Baek committed
875
876
877
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
878
879
@app.get("/api/config")
async def get_app_config():
880
    # Checking and Handling the Absence of 'ui' in CONFIG_DATA
881
882

    default_locale = "en-US"
883
884
885
886
    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
887
888
    return {
        "status": True,
889
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
890
        "version": VERSION,
891
        "default_locale": default_locale,
892
893
        "default_models": webui_app.state.config.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.config.DEFAULT_PROMPT_SUGGESTIONS,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
894
        "features": {
895
896
            "auth": WEBUI_AUTH,
            "auth_trusted_header": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
897
            "enable_signup": webui_app.state.config.ENABLE_SIGNUP,
Timothy J. Baek's avatar
Timothy J. Baek committed
898
            "enable_web_search": rag_app.state.config.ENABLE_RAG_WEB_SEARCH,
899
            "enable_image_generation": images_app.state.config.ENABLED,
900
            "enable_community_sharing": webui_app.state.config.ENABLE_COMMUNITY_SHARING,
901
            "enable_admin_export": ENABLE_ADMIN_EXPORT,
902
        },
Timothy J. Baek's avatar
Timothy J. Baek committed
903
904
905
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
906
907
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
908
    return {
909
910
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
911
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
912
913
914
915
916
917
918
919


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


@app.post("/api/config/model/filter")
Timothy J. Baek's avatar
Timothy J. Baek committed
920
async def update_model_filter_config(
Timothy J. Baek's avatar
Timothy J. Baek committed
921
922
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
923
924
    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
925

926
    return {
927
928
        "enabled": app.state.config.ENABLE_MODEL_FILTER,
        "models": app.state.config.MODEL_FILTER_LIST,
929
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
930
931


Timothy J. Baek's avatar
Timothy J. Baek committed
932
933
934
@app.get("/api/webhook")
async def get_webhook_url(user=Depends(get_admin_user)):
    return {
935
        "url": app.state.config.WEBHOOK_URL,
Timothy J. Baek's avatar
Timothy J. Baek committed
936
937
938
939
940
941
942
943
944
    }


class UrlForm(BaseModel):
    url: str


@app.post("/api/webhook")
async def update_webhook_url(form_data: UrlForm, user=Depends(get_admin_user)):
945
946
    app.state.config.WEBHOOK_URL = form_data.url
    webui_app.state.WEBHOOK_URL = app.state.config.WEBHOOK_URL
947
    return {"url": app.state.config.WEBHOOK_URL}
948
949


950
951
952
953
954
955
956
@app.get("/api/version")
async def get_app_config():
    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
957
958
@app.get("/api/changelog")
async def get_app_changelog():
959
    return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
Timothy J. Baek's avatar
Timothy J. Baek committed
960
961


962
963
964
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
965
        async with aiohttp.ClientSession(trust_env=True) as session:
Timothy J. Baek's avatar
Timothy J. Baek committed
966
967
968
969
970
971
972
973
974
            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:
975
976
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
977
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
978
979
        )

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

981
982
983
@app.get("/manifest.json")
async def get_manifest_json():
    return {
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
984
985
        "name": WEBUI_NAME,
        "short_name": WEBUI_NAME,
986
987
988
989
990
        "start_url": "/",
        "display": "standalone",
        "background_color": "#343541",
        "theme_color": "#343541",
        "orientation": "portrait-primary",
Timothy J. Baek's avatar
Timothy J. Baek committed
991
        "icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
992
993
    }

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

995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
@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
1010
1011
1012
1013
1014
@app.get("/health")
async def healthcheck():
    return {"status": True}


1015
1016
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
1017

1018
if os.path.exists(FRONTEND_BUILD_DIR):
Timothy J. Baek's avatar
Timothy J. Baek committed
1019
    mimetypes.add_type("text/javascript", ".js")
1020
1021
1022
1023
1024
1025
1026
1027
1028
    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."
    )