open-webui/backend/apps/openai/main.py

342 lines
10 KiB
Python
Raw Normal View History

2024-01-05 01:49:34 +01:00
from fastapi import FastAPI, Request, Response, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
2024-02-06 07:51:08 +01:00
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
2024-01-05 01:49:34 +01:00
import requests
2024-03-07 01:13:25 +01:00
import aiohttp
import asyncio
2024-01-05 01:49:34 +01:00
import json
2024-03-07 01:13:25 +01:00
2024-01-05 01:49:34 +01:00
from pydantic import BaseModel
2024-02-06 07:51:08 +01:00
2024-01-05 01:49:34 +01:00
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
2024-02-17 08:30:38 +01:00
from utils.utils import (
decode_token,
get_current_user,
get_verified_user,
get_admin_user,
)
2024-03-10 06:47:01 +01:00
from config import (
OPENAI_API_BASE_URLS,
OPENAI_API_KEYS,
CACHE_DIR,
MODEL_FILTER_ENABLED,
MODEL_FILTER_LIST,
)
2024-03-07 01:13:25 +01:00
from typing import List, Optional
2024-02-06 07:51:08 +01:00
import hashlib
from pathlib import Path
2024-01-05 01:49:34 +01:00
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2024-03-10 06:47:01 +01:00
app.state.MODEL_FILTER_ENABLED = MODEL_FILTER_ENABLED
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
2024-03-10 06:19:20 +01:00
2024-03-07 01:13:25 +01:00
app.state.OPENAI_API_BASE_URLS = OPENAI_API_BASE_URLS
app.state.OPENAI_API_KEYS = OPENAI_API_KEYS
app.state.MODELS = {}
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
@app.middleware("http")
async def check_url(request: Request, call_next):
if len(app.state.MODELS) == 0:
await get_all_models()
else:
pass
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
response = await call_next(request)
return response
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
class UrlsUpdateForm(BaseModel):
urls: List[str]
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
class KeysUpdateForm(BaseModel):
keys: List[str]
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
@app.get("/urls")
async def get_openai_urls(user=Depends(get_admin_user)):
return {"OPENAI_API_BASE_URLS": app.state.OPENAI_API_BASE_URLS}
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
@app.post("/urls/update")
async def update_openai_urls(form_data: UrlsUpdateForm, user=Depends(get_admin_user)):
app.state.OPENAI_API_BASE_URLS = form_data.urls
return {"OPENAI_API_BASE_URLS": app.state.OPENAI_API_BASE_URLS}
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
@app.get("/keys")
async def get_openai_keys(user=Depends(get_admin_user)):
return {"OPENAI_API_KEYS": app.state.OPENAI_API_KEYS}
@app.post("/keys/update")
async def update_openai_key(form_data: KeysUpdateForm, user=Depends(get_admin_user)):
app.state.OPENAI_API_KEYS = form_data.keys
return {"OPENAI_API_KEYS": app.state.OPENAI_API_KEYS}
2024-01-05 01:49:34 +01:00
2024-02-06 07:51:08 +01:00
@app.post("/audio/speech")
async def speech(request: Request, user=Depends(get_verified_user)):
2024-03-07 01:13:25 +01:00
idx = None
try:
idx = app.state.OPENAI_API_BASE_URLS.index("https://api.openai.com/v1")
body = await request.body()
name = hashlib.sha256(body).hexdigest()
SPEECH_CACHE_DIR = Path(CACHE_DIR).joinpath("./audio/speech/")
SPEECH_CACHE_DIR.mkdir(parents=True, exist_ok=True)
file_path = SPEECH_CACHE_DIR.joinpath(f"{name}.mp3")
file_body_path = SPEECH_CACHE_DIR.joinpath(f"{name}.json")
# Check if the file already exists in the cache
if file_path.is_file():
return FileResponse(file_path)
headers = {}
headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEYS[idx]}"
headers["Content-Type"] = "application/json"
2024-03-16 10:45:24 +01:00
r = None
2024-03-07 01:13:25 +01:00
try:
r = requests.post(
url=f"{app.state.OPENAI_API_BASE_URLS[idx]}/audio/speech",
data=body,
headers=headers,
stream=True,
)
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
r.raise_for_status()
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
# Save the streaming content to a file
with open(file_path, "wb") as f:
for chunk in r.iter_content(chunk_size=8192):
f.write(chunk)
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
with open(file_body_path, "w") as f:
json.dump(json.loads(body.decode("utf-8")), f)
2024-02-06 07:57:24 +01:00
2024-03-07 01:13:25 +01:00
# Return the saved file
return FileResponse(file_path)
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
except Exception as e:
print(e)
error_detail = "Open WebUI: Server Connection Error"
if r is not None:
try:
res = r.json()
if "error" in res:
error_detail = f"External: {res['error']}"
except:
error_detail = f"External: {e}"
2024-03-16 11:25:20 +01:00
raise HTTPException(
status_code=r.status_code if r else 500, detail=error_detail
)
2024-03-07 01:13:25 +01:00
except ValueError:
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.OPENAI_NOT_FOUND)
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
async def fetch_url(url, key):
2024-02-06 07:51:08 +01:00
try:
2024-03-07 01:13:25 +01:00
headers = {"Authorization": f"Bearer {key}"}
async with aiohttp.ClientSession() as session:
async with session.get(url, headers=headers) as response:
return await response.json()
except Exception as e:
# Handle connection error here
print(f"Connection error: {e}")
return None
def merge_models_lists(model_lists):
merged_list = []
for idx, models in enumerate(model_lists):
2024-03-18 09:11:48 +01:00
if models is not None and "error" not in models:
merged_list.extend(
[
{**model, "urlIdx": idx}
for model in models
if "api.openai.com" not in app.state.OPENAI_API_BASE_URLS[idx]
or "gpt" in model["id"]
]
)
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
return merged_list
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
async def get_all_models():
print("get_all_models")
2024-03-11 03:00:43 +01:00
if len(app.state.OPENAI_API_KEYS) == 1 and app.state.OPENAI_API_KEYS[0] == "":
models = {"data": []}
else:
tasks = [
fetch_url(f"{url}/models", app.state.OPENAI_API_KEYS[idx])
for idx, url in enumerate(app.state.OPENAI_API_BASE_URLS)
]
2024-03-18 09:11:48 +01:00
2024-03-11 03:00:43 +01:00
responses = await asyncio.gather(*tasks)
models = {
"data": merge_models_lists(
2024-03-18 09:11:48 +01:00
list(
map(
lambda response: response["data"] if response else None,
responses,
)
)
2024-03-11 03:00:43 +01:00
)
}
2024-03-18 09:11:48 +01:00
print(models)
2024-03-11 03:00:43 +01:00
app.state.MODELS = {model["id"]: model for model in models["data"]}
2024-02-06 07:51:08 +01:00
2024-03-11 03:00:43 +01:00
return models
2024-02-06 07:51:08 +01:00
2024-03-07 01:13:25 +01:00
@app.get("/models")
@app.get("/models/{url_idx}")
2024-03-10 06:19:20 +01:00
async def get_models(url_idx: Optional[int] = None, user=Depends(get_current_user)):
2024-03-07 01:13:25 +01:00
if url_idx == None:
2024-03-10 06:19:20 +01:00
models = await get_all_models()
if app.state.MODEL_FILTER_ENABLED:
if user.role == "user":
models["data"] = list(
filter(
2024-03-10 06:47:01 +01:00
lambda model: model["id"] in app.state.MODEL_FILTER_LIST,
models["data"],
)
2024-03-10 06:19:20 +01:00
)
return models
return models
2024-03-07 01:13:25 +01:00
else:
url = app.state.OPENAI_API_BASE_URLS[url_idx]
2024-03-16 10:45:24 +01:00
r = None
2024-03-07 01:13:25 +01:00
try:
r = requests.request(method="GET", url=f"{url}/models")
r.raise_for_status()
response_data = r.json()
if "api.openai.com" in url:
response_data["data"] = list(
filter(lambda model: "gpt" in model["id"], response_data["data"])
)
return response_data
except Exception as e:
print(e)
error_detail = "Open WebUI: Server Connection Error"
if r is not None:
try:
res = r.json()
if "error" in res:
error_detail = f"External: {res['error']}"
except:
error_detail = f"External: {e}"
raise HTTPException(
status_code=r.status_code if r else 500,
detail=error_detail,
)
2024-02-06 07:51:08 +01:00
2024-01-05 01:49:34 +01:00
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
2024-03-07 01:13:25 +01:00
idx = 0
2024-01-05 01:49:34 +01:00
2024-01-22 10:37:54 +01:00
body = await request.body()
# TODO: Remove below after gpt-4-vision fix from Open AI
# Try to decode the body of the request from bytes to a UTF-8 string (Require add max_token to fix gpt-4-vision)
try:
2024-01-22 10:37:54 +01:00
body = body.decode("utf-8")
body = json.loads(body)
2024-03-07 01:13:25 +01:00
idx = app.state.MODELS[body.get("model")]["urlIdx"]
# Check if the model is "gpt-4-vision-preview" and set "max_tokens" to 4000
# This is a workaround until OpenAI fixes the issue with this model
2024-01-22 10:37:54 +01:00
if body.get("model") == "gpt-4-vision-preview":
if "max_tokens" not in body:
body["max_tokens"] = 4000
2024-01-22 10:37:54 +01:00
print("Modified body_dict:", body)
# Fix for ChatGPT calls failing because the num_ctx key is in body
if "num_ctx" in body:
# If 'num_ctx' is in the dictionary, delete it
# Leaving it there generates an error with the
# OpenAI API (Feb 2024)
del body["num_ctx"]
2024-01-22 10:37:54 +01:00
# Convert the modified body back to JSON
body = json.dumps(body)
except json.JSONDecodeError as e:
print("Error loading request body into a dictionary:", e)
2024-01-05 01:49:34 +01:00
2024-03-07 01:13:25 +01:00
url = app.state.OPENAI_API_BASE_URLS[idx]
key = app.state.OPENAI_API_KEYS[idx]
target_url = f"{url}/{path}"
if key == "":
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
2024-01-05 03:38:03 +01:00
headers = {}
2024-03-07 01:13:25 +01:00
headers["Authorization"] = f"Bearer {key}"
2024-01-05 03:54:00 +01:00
headers["Content-Type"] = "application/json"
2024-01-05 01:49:34 +01:00
2024-03-16 10:45:24 +01:00
r = None
2024-01-05 01:49:34 +01:00
try:
r = requests.request(
method=request.method,
url=target_url,
2024-01-22 10:37:54 +01:00
data=body,
2024-01-05 01:49:34 +01:00
headers=headers,
stream=True,
)
r.raise_for_status()
2024-01-05 03:38:03 +01:00
# Check if response is SSE
if "text/event-stream" in r.headers.get("Content-Type", ""):
return StreamingResponse(
r.iter_content(chunk_size=8192),
status_code=r.status_code,
headers=dict(r.headers),
)
else:
response_data = r.json()
return response_data
2024-01-05 01:49:34 +01:00
except Exception as e:
print(e)
2024-02-17 08:30:38 +01:00
error_detail = "Open WebUI: Server Connection Error"
2024-01-05 01:49:34 +01:00
if r is not None:
try:
res = r.json()
if "error" in res:
error_detail = f"External: {res['error']}"
except:
error_detail = f"External: {e}"
2024-03-16 11:25:20 +01:00
raise HTTPException(
status_code=r.status_code if r else 500, detail=error_detail
)