open-webui/backend/main.py

337 lines
9.6 KiB
Python
Raw Permalink Normal View History

2024-02-23 09:30:26 +01:00
from bs4 import BeautifulSoup
import json
import markdown
2024-01-07 11:48:21 +01:00
import time
2024-02-24 09:21:53 +01:00
import os
import sys
import logging
2024-04-10 08:03:05 +02:00
import aiohttp
2024-02-25 20:26:58 +01:00
import requests
2024-02-23 09:30:26 +01:00
2024-02-25 20:26:58 +01:00
from fastapi import FastAPI, Request, Depends, status
2023-11-15 01:28:51 +01:00
from fastapi.staticfiles import StaticFiles
from fastapi import HTTPException
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.middleware.cors import CORSMiddleware
2023-11-19 01:47:12 +01:00
from starlette.exceptions import HTTPException as StarletteHTTPException
2024-03-09 07:34:47 +01:00
from starlette.middleware.base import BaseHTTPMiddleware
2023-11-15 01:28:51 +01:00
2024-01-07 07:07:20 +01:00
2023-11-15 01:28:51 +01:00
from apps.ollama.main import app as ollama_app
2024-01-05 03:38:03 +01:00
from apps.openai.main import app as openai_app
2024-04-10 08:03:05 +02:00
2024-04-21 08:46:09 +02:00
from apps.litellm.main import (
app as litellm_app,
start_litellm_background,
shutdown_litellm_background,
)
2024-02-11 09:17:50 +01:00
from apps.audio.main import app as audio_app
2024-02-22 03:12:01 +01:00
from apps.images.main import app as images_app
from apps.rag.main import app as rag_app
2023-11-19 01:47:12 +01:00
from apps.web.main import app as webui_app
2024-01-07 07:07:20 +01:00
2024-04-21 08:22:02 +02:00
import asyncio
2024-03-10 06:19:20 +01:00
from pydantic import BaseModel
from typing import List
2024-02-24 07:44:56 +01:00
2024-03-10 06:19:20 +01:00
from utils.utils import get_admin_user
2024-03-11 02:40:50 +01:00
from apps.rag.utils import rag_messages
2024-03-09 07:34:47 +01:00
2024-03-10 06:47:01 +01:00
from config import (
2024-03-24 04:16:18 +01:00
CONFIG_DATA,
2024-03-10 06:47:01 +01:00
WEBUI_NAME,
ENV,
VERSION,
CHANGELOG,
FRONTEND_BUILD_DIR,
CACHE_DIR,
STATIC_DIR,
2024-04-26 23:19:50 +02:00
ENABLE_LITELLM,
ENABLE_MODEL_FILTER,
2024-03-10 06:47:01 +01:00
MODEL_FILTER_LIST,
GLOBAL_LOG_LEVEL,
SRC_LOG_LEVELS,
2024-03-21 02:35:02 +01:00
WEBHOOK_URL,
ENABLE_ADMIN_EXPORT,
2024-03-10 06:47:01 +01:00
)
2024-02-25 20:26:58 +01:00
from constants import ERROR_MESSAGES
logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MAIN"])
2023-11-15 01:28:51 +01:00
2024-03-28 10:45:56 +01:00
2023-11-15 01:28:51 +01:00
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
2024-04-02 12:03:55 +02:00
print(
f"""
___ __ __ _ _ _ ___
/ _ \ _ __ ___ _ __ \ \ / /__| |__ | | | |_ _|
| | | | '_ \ / _ \ '_ \ \ \ /\ / / _ \ '_ \| | | || |
| |_| | |_) | __/ | | | \ V V / __/ |_) | |_| || |
\___/| .__/ \___|_| |_| \_/\_/ \___|_.__/ \___/|___|
|_|
v{VERSION} - building the best open-source AI user interface.
https://github.com/open-webui/open-webui
"""
)
2024-01-07 11:48:21 +01:00
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
2023-11-15 01:28:51 +01:00
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
2024-03-10 06:47:01 +01:00
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
2024-03-10 06:19:20 +01:00
2024-03-21 02:35:02 +01:00
app.state.WEBHOOK_URL = WEBHOOK_URL
2023-11-15 01:28:51 +01:00
origins = ["*"]
2024-02-24 09:21:53 +01:00
2024-03-09 07:34:47 +01:00
class RAGMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
if request.method == "POST" and (
"/api/chat" in request.url.path or "/chat/completions" in request.url.path
):
log.debug(f"request.url.path: {request.url.path}")
2024-03-09 07:52:42 +01:00
2024-03-09 07:34:47 +01:00
# 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:
2024-03-11 02:40:50 +01:00
data = {**data}
data["messages"] = rag_messages(
2024-04-27 21:38:50 +02:00
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,
2024-03-11 02:40:50 +01:00
)
2024-03-09 07:34:47 +01:00
del data["docs"]
log.debug(f"data['messages']: {data['messages']}")
2024-03-10 05:12:32 +01:00
2024-03-09 07:34:47 +01:00
modified_body_bytes = json.dumps(data).encode("utf-8")
2024-03-11 02:40:50 +01:00
# 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"
],
]
2024-03-09 07:34:47 +01:00
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)
2024-03-11 02:40:50 +01:00
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2023-11-15 01:28:51 +01:00
@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
2024-03-11 02:40:50 +01:00
@app.on_event("startup")
async def on_startup():
2024-04-26 23:19:50 +02:00
if ENABLE_LITELLM:
asyncio.create_task(start_litellm_background())
2024-03-11 02:40:50 +01:00
2023-11-19 01:47:12 +01:00
app.mount("/api/v1", webui_app)
2024-02-22 12:22:23 +01:00
app.mount("/litellm/api", litellm_app)
app.mount("/ollama", ollama_app)
2024-01-05 03:38:03 +01:00
app.mount("/openai/api", openai_app)
2024-02-11 09:17:50 +01:00
2024-02-22 03:12:01 +01:00
app.mount("/images/api/v1", images_app)
2024-02-11 09:17:50 +01:00
app.mount("/audio/api/v1", audio_app)
2024-01-07 07:07:20 +01:00
app.mount("/rag/api/v1", rag_app)
2024-03-31 22:59:39 +02:00
2024-02-22 03:12:01 +01:00
@app.get("/api/config")
async def get_app_config():
# Checking and Handling the Absence of 'ui' in CONFIG_DATA
2024-03-31 22:59:39 +02:00
default_locale = "en-US"
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
2024-02-22 03:12:01 +01:00
return {
"status": True,
2024-02-24 02:12:19 +01:00
"name": WEBUI_NAME,
2024-02-23 09:30:26 +01:00
"version": VERSION,
"default_locale": default_locale,
2024-02-22 03:12:01 +01:00
"images": images_app.state.ENABLED,
"default_models": webui_app.state.DEFAULT_MODELS,
2024-03-31 22:59:39 +02:00
"default_prompt_suggestions": webui_app.state.DEFAULT_PROMPT_SUGGESTIONS,
"trusted_header_auth": bool(webui_app.state.AUTH_TRUSTED_EMAIL_HEADER),
"admin_export_enabled": ENABLE_ADMIN_EXPORT,
2024-02-22 03:12:01 +01:00
}
2024-03-10 06:19:20 +01:00
@app.get("/api/config/model/filter")
async def get_model_filter_config(user=Depends(get_admin_user)):
2024-03-10 06:47:01 +01:00
return {
"enabled": app.state.ENABLE_MODEL_FILTER,
2024-03-10 06:47:01 +01:00
"models": app.state.MODEL_FILTER_LIST,
}
2024-03-10 06:19:20 +01:00
class ModelFilterConfigForm(BaseModel):
enabled: bool
models: List[str]
@app.post("/api/config/model/filter")
2024-03-21 02:35:02 +01:00
async def update_model_filter_config(
2024-03-10 06:19:20 +01:00
form_data: ModelFilterConfigForm, user=Depends(get_admin_user)
):
app.state.ENABLE_MODEL_FILTER = form_data.enabled
2024-03-10 06:47:01 +01:00
app.state.MODEL_FILTER_LIST = form_data.models
2024-03-10 06:19:20 +01:00
ollama_app.state.ENABLE_MODEL_FILTER = app.state.ENABLE_MODEL_FILTER
2024-03-10 06:47:01 +01:00
ollama_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
2024-03-10 06:19:20 +01:00
openai_app.state.ENABLE_MODEL_FILTER = app.state.ENABLE_MODEL_FILTER
2024-03-10 06:47:01 +01:00
openai_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
2024-03-10 06:19:20 +01:00
litellm_app.state.ENABLE_MODEL_FILTER = app.state.ENABLE_MODEL_FILTER
2024-03-28 10:45:56 +01:00
litellm_app.state.MODEL_FILTER_LIST = app.state.MODEL_FILTER_LIST
2024-03-10 06:47:01 +01:00
return {
"enabled": app.state.ENABLE_MODEL_FILTER,
2024-03-10 06:47:01 +01:00
"models": app.state.MODEL_FILTER_LIST,
}
2024-03-10 06:19:20 +01:00
2024-03-21 02:35:02 +01:00
@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,
}
2024-03-05 09:59:35 +01:00
@app.get("/api/version")
async def get_app_config():
return {
"version": VERSION,
}
2024-02-23 09:30:26 +01:00
@app.get("/api/changelog")
async def get_app_changelog():
2024-03-31 10:10:57 +02:00
return {key: CHANGELOG[key] for idx, key in enumerate(CHANGELOG) if idx < 5}
2024-02-23 09:30:26 +01:00
2024-02-25 20:26:58 +01:00
@app.get("/api/version/updates")
async def get_app_latest_release_version():
try:
2024-04-10 08:03:05 +02:00
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:
2024-02-25 20:26:58 +01:00
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
2024-02-25 20:55:15 +01:00
detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED,
2024-02-25 20:26:58 +01:00
)
2024-04-10 10:27:19 +02:00
@app.get("/manifest.json")
async def get_manifest_json():
return {
2024-04-04 05:43:55 +02:00
"name": WEBUI_NAME,
"short_name": WEBUI_NAME,
"start_url": "/",
"display": "standalone",
"background_color": "#343541",
"theme_color": "#343541",
"orientation": "portrait-primary",
2024-05-02 04:32:36 +02:00
"icons": [{"src": "/static/logo.png", "type": "image/png", "sizes": "500x500"}],
}
2024-04-10 10:27:19 +02:00
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.mount("/cache", StaticFiles(directory=CACHE_DIR), name="cache")
2024-02-24 02:12:19 +01:00
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."
)
2024-04-21 08:46:09 +02:00
@app.on_event("shutdown")
async def shutdown_event():
2024-04-26 23:19:50 +02:00
if ENABLE_LITELLM:
await shutdown_litellm_background()