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

Timothy J. Baek's avatar
Timothy J. Baek committed
27
28
from apps.rag.utils import query_doc, query_collection, rag_template

29
from config import WEBUI_NAME, ENV, VERSION, CHANGELOG, FRONTEND_BUILD_DIR
30
31
from constants import ERROR_MESSAGES

Timothy J. Baek's avatar
Timothy J. Baek committed
32
33
34
35
36
37
38
39
40
41
42
43

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


44
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
Timothy J. Baek's avatar
Timothy J. Baek committed
45
46
47
48
49
50
51
52
53
54
55
56

origins = ["*"]

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


Timothy J. Baek's avatar
Timothy J. Baek committed
57
58
@app.on_event("startup")
async def on_startup():
Timothy J. Baek's avatar
Timothy J. Baek committed
59
    await litellm_app_startup()
Timothy J. Baek's avatar
Timothy J. Baek committed
60
61


Timothy J. Baek's avatar
Timothy J. Baek committed
62
63
64
class RAGMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):

65
66
        "chat/completions" in request.url.path

Timothy J. Baek's avatar
Timothy J. Baek committed
67
        print(request.url.path)
68
69
70
        if request.method == "POST" and (
            "/api/chat" in request.url.path or "/chat/completions" in request.url.path
        ):
Timothy J. Baek's avatar
Timothy J. Baek committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
            # 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

                query = data["messages"][last_user_message_idx]["content"]

                relevant_contexts = []

                for doc in docs:
                    context = None
                    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,
                        )
                    relevant_contexts.append(context)

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

                content = rag_template(
                    template=rag_app.state.RAG_TEMPLATE,
                    context=context_string,
                    query=query,
                )

                new_user_message = {
                    **data["messages"][last_user_message_idx],
                    "content": content,
                }
                data["messages"][last_user_message_idx] = new_user_message
                del data["docs"]

            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
147
148
149
150
151
152
153
154
155
156
@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


157
app.mount("/api/v1", webui_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
158
159
app.mount("/litellm/api", litellm_app)

160
app.mount("/ollama", ollama_app)
161
app.mount("/openai/api", openai_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
162

Timothy J. Baek's avatar
Timothy J. Baek committed
163
app.mount("/images/api/v1", images_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
164
app.mount("/audio/api/v1", audio_app)
Timothy J. Baek's avatar
Timothy J. Baek committed
165
166
app.mount("/rag/api/v1", rag_app)

167

Timothy J. Baek's avatar
Timothy J. Baek committed
168
169
@app.get("/api/config")
async def get_app_config():
Timothy J. Baek's avatar
Timothy J. Baek committed
170

Timothy J. Baek's avatar
Timothy J. Baek committed
171
172
    return {
        "status": True,
173
        "name": WEBUI_NAME,
Timothy J. Baek's avatar
Timothy J. Baek committed
174
        "version": VERSION,
Timothy J. Baek's avatar
Timothy J. Baek committed
175
176
177
178
179
180
        "images": images_app.state.ENABLED,
        "default_models": webui_app.state.DEFAULT_MODELS,
        "default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
    }


181
182
183
184
185
186
187
188
@app.get("/api/version")
async def get_app_config():

    return {
        "version": VERSION,
    }


Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
@app.get("/api/changelog")
async def get_app_changelog():
Timothy J. Baek's avatar
Timothy J. Baek committed
191
    return CHANGELOG
Timothy J. Baek's avatar
Timothy J. Baek committed
192
193


194
195
196
197
198
199
200
201
202
@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
203
        return {"current": VERSION, "latest": latest_version[1:]}
204
205
206
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
Timothy J. Baek's avatar
Timothy J. Baek committed
207
            detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
208
209
210
        )


211
app.mount("/static", StaticFiles(directory="static"), name="static")
Timothy J. Baek's avatar
Timothy J. Baek committed
212
app.mount("/cache", StaticFiles(directory="data/cache"), name="cache")
213
214


215
216
app.mount(
    "/",
lucasew's avatar
lucasew committed
217
    SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
218
219
    name="spa-static-files",
)