main.py 9.39 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 requests
Timothy J. Baek's avatar
Timothy J. Baek committed
8

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
18
from apps.ollama.main import app as ollama_app
19
from apps.openai.main import app as openai_app
Timothy J. Baek's avatar
Timothy J. Baek committed
20
from apps.litellm.main import app as litellm_app, startup as litellm_app_startup
Timothy J. Baek's avatar
Timothy J. Baek committed
21
from apps.audio.main import app as audio_app
Timothy J. Baek's avatar
Timothy J. Baek committed
22
23
from apps.images.main import app as images_app
from apps.rag.main import app as rag_app
24
from apps.web.main import app as webui_app
Timothy J. Baek's avatar
Timothy J. Baek committed
25

Timothy J. Baek's avatar
Timothy J. Baek committed
26
27
from pydantic import BaseModel
from typing import List
Timothy J. Baek's avatar
Timothy J. Baek committed
28

Timothy J. Baek's avatar
Timothy J. Baek committed
29
30

from utils.utils import get_admin_user
Timothy J. Baek's avatar
Timothy J. Baek committed
31
32
from apps.rag.utils import query_doc, query_collection, rag_template

33
34
35
36
37
38
39
40
41
from config import (
    WEBUI_NAME,
    ENV,
    VERSION,
    CHANGELOG,
    FRONTEND_BUILD_DIR,
    MODEL_FILTER_ENABLED,
    MODEL_FILTER_LIST,
)
42
43
from constants import ERROR_MESSAGES

Timothy J. Baek's avatar
Timothy J. Baek committed
44
45
46
47
48
49
50
51
52
53
54
55

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


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

58
59
app.state.MODEL_FILTER_ENABLED = MODEL_FILTER_ENABLED
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
60

Timothy J. Baek's avatar
Timothy J. Baek committed
61
62
63
64
65
66
67
68
69
70
71
origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
@app.on_event("startup")
async def on_startup():
Timothy J. Baek's avatar
Timothy J. Baek committed
74
    await litellm_app_startup()
Timothy J. Baek's avatar
Timothy J. Baek committed
75
76


Timothy J. Baek's avatar
Timothy J. Baek committed
77
78
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
79
80
81
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
82
83
            print(request.url.path)

Timothy J. Baek's avatar
Timothy J. Baek committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
            # 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 {}

            # Example: Add a new key-value pair or modify existing ones
            # data["modified"] = True  # Example modification
            if "docs" in data:
                docs = data["docs"]
                print(docs)

                last_user_message_idx = None
                for i in range(len(data["messages"]) - 1, -1, -1):
                    if data["messages"][i]["role"] == "user":
                        last_user_message_idx = i
                        break

Timothy J. Baek's avatar
Timothy J. Baek committed
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
                user_message = data["messages"][last_user_message_idx]

                if isinstance(user_message["content"], list):
                    # Handle list content input
                    content_type = "list"
                    query = ""
                    for content_item in user_message["content"]:
                        if content_item["type"] == "text":
                            query = content_item["text"]
                            break
                elif isinstance(user_message["content"], str):
                    # Handle text content input
                    content_type = "text"
                    query = user_message["content"]
                else:
                    # Fallback in case the input does not match expected types
                    content_type = None
                    query = ""
Timothy J. Baek's avatar
Timothy J. Baek committed
121
122
123
124
125

                relevant_contexts = []

                for doc in docs:
                    context = None
Timothy J. Baek's avatar
Timothy J. Baek committed
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

                    try:
                        if doc["type"] == "collection":
                            context = query_collection(
                                collection_names=doc["collection_names"],
                                query=query,
                                k=rag_app.state.TOP_K,
                                embedding_function=rag_app.state.sentence_transformer_ef,
                            )
                        else:
                            context = query_doc(
                                collection_name=doc["collection_name"],
                                query=query,
                                k=rag_app.state.TOP_K,
                                embedding_function=rag_app.state.sentence_transformer_ef,
                            )
                    except Exception as e:
                        print(e)
                        context = None

Timothy J. Baek's avatar
Timothy J. Baek committed
146
147
148
149
150
151
152
                    relevant_contexts.append(context)

                context_string = ""
                for context in relevant_contexts:
                    if context:
                        context_string += " ".join(context["documents"][0]) + "\n"

Timothy J. Baek's avatar
Timothy J. Baek committed
153
                ra_content = rag_template(
Timothy J. Baek's avatar
Timothy J. Baek committed
154
155
156
157
158
                    template=rag_app.state.RAG_TEMPLATE,
                    context=context_string,
                    query=query,
                )

Timothy J. Baek's avatar
Timothy J. Baek committed
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
                if content_type == "list":
                    new_content = []
                    for content_item in user_message["content"]:
                        if content_item["type"] == "text":
                            # Update the text item's content with ra_content
                            new_content.append({"type": "text", "text": ra_content})
                        else:
                            # Keep other types of content as they are
                            new_content.append(content_item)
                    new_user_message = {**user_message, "content": new_content}
                else:
                    new_user_message = {
                        **user_message,
                        "content": ra_content,
                    }

Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
177
                data["messages"][last_user_message_idx] = new_user_message
                del data["docs"]

Timothy J. Baek's avatar
Timothy J. Baek committed
178
179
                print(data["messages"])

Timothy J. Baek's avatar
Timothy J. Baek committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
            modified_body_bytes = json.dumps(data).encode("utf-8")

            # Create a new request with the modified body
            scope = request.scope
            scope["body"] = modified_body_bytes
            request = Request(scope, receive=lambda: self._receive(modified_body_bytes))

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


Timothy J. Baek's avatar
Timothy J. Baek committed
197
198
199
200
201
202
203
204
205
206
@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


207
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
208
209
app.mount("/litellm/api", litellm_app)

210
app.mount("/ollama", ollama_app)
211
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
212

Timothy J. Baek's avatar
Timothy J. Baek committed
213
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
214
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
215
216
app.mount("/rag/api/v1", rag_app)

217

Timothy J. Baek's avatar
Timothy J. Baek committed
218
219
@app.get("/api/config")
async def get_app_config():
Timothy J. Baek's avatar
Timothy J. Baek committed
220

Timothy J. Baek's avatar
Timothy J. Baek committed
221
222
    return {
        "status": True,
223
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
224
        "version": VERSION,
Timothy J. Baek's avatar
Timothy J. Baek committed
225
226
227
228
229
230
        "images": images_app.state.ENABLED,
        "default_models": webui_app.state.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
231
232
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
233
234
235
236
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
237
238
239
240
241
242
243
244
245
246
247
248
249


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


@app.post("/api/config/model/filter")
async def get_model_filter_config(
    form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):

    app.state.MODEL_FILTER_ENABLED = form_data.enabled
250
    app.state.MODEL_FILTER_LIST = form_data.models
Timothy J. Baek's avatar
Timothy J. Baek committed
251
252

    ollama_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
253
    ollama_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
254
255

    openai_app.state.MODEL_FILTER_ENABLED = app.state.MODEL_FILTER_ENABLED
256
    openai_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
Timothy J. Baek's avatar
Timothy J. Baek committed
257

258
259
260
261
    return {
        "enabled": app.state.MODEL_FILTER_ENABLED,
        "models": app.state.MODEL_FILTER_LIST,
    }
Timothy J. Baek's avatar
Timothy J. Baek committed
262
263


264
265
266
267
268
269
270
271
@app.get("/api/version")
async def get_app_config():

    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
272
273
@app.get("/api/changelog")
async def get_app_changelog():
Timothy J. Baek's avatar
Timothy J. Baek committed
274
    return CHANGELOG
Timothy J. Baek's avatar
Timothy J. Baek committed
275
276


277
278
279
280
281
282
283
284
285
@app.get("/api/version/updates")
async def get_app_latest_release_version():
    try:
        response = requests.get(
            f"https://api.github.com/repos/open-webui/open-webui/releases/latest"
        )
        response.raise_for_status()
        latest_version = response.json()["tag_name"]

Timothy J. Baek's avatar
Timothy J. Baek committed
286
        return {"current": VERSION, "latest": latest_version[1:]}
287
288
289
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
290
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
291
292
293
        )


294
app.mount("/static", StaticFiles(directory="static"), name="static")
Timothy J. Baek's avatar
Timothy J. Baek committed
295
app.mount("/cache", StaticFiles(directory="data/cache"), name="cache")
296
297


298
299
app.mount(
    "/",
lucasew's avatar
lucasew committed
300
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
301
302
    name="spa-static-files",
)