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

216 lines
6.9 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
import json
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
from utils.utils import decode_token, get_current_user
2024-02-06 07:51:08 +01:00
from config import OPENAI_API_BASE_URL, OPENAI_API_KEY, CACHE_DIR
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=["*"],
)
app.state.OPENAI_API_BASE_URL = OPENAI_API_BASE_URL
app.state.OPENAI_API_KEY = OPENAI_API_KEY
class UrlUpdateForm(BaseModel):
url: str
class KeyUpdateForm(BaseModel):
key: str
@app.get("/url")
async def get_openai_url(user=Depends(get_current_user)):
if user and user.role == "admin":
return {"OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL}
else:
2024-01-22 10:37:54 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
2024-01-05 01:49:34 +01:00
@app.post("/url/update")
2024-01-22 10:37:54 +01:00
async def update_openai_url(form_data: UrlUpdateForm, user=Depends(get_current_user)):
2024-01-05 01:49:34 +01:00
if user and user.role == "admin":
app.state.OPENAI_API_BASE_URL = form_data.url
return {"OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL}
else:
2024-01-22 10:37:54 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
2024-01-05 01:49:34 +01:00
@app.get("/key")
async def get_openai_key(user=Depends(get_current_user)):
if user and user.role == "admin":
return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY}
else:
2024-01-22 10:37:54 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
2024-01-05 01:49:34 +01:00
@app.post("/key/update")
2024-01-22 10:37:54 +01:00
async def update_openai_key(form_data: KeyUpdateForm, user=Depends(get_current_user)):
2024-01-05 01:49:34 +01:00
if user and user.role == "admin":
app.state.OPENAI_API_KEY = form_data.key
return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY}
else:
2024-01-22 10:37:54 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
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_current_user)):
target_url = f"{app.state.OPENAI_API_BASE_URL}/audio/speech"
if user.role not in ["user", "admin"]:
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
if app.state.OPENAI_API_KEY == "":
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
body = await request.body()
2024-02-06 07:57:24 +01:00
name = hashlib.sha256(body).hexdigest()
2024-02-06 07:51:08 +01:00
SPEECH_CACHE_DIR = Path(CACHE_DIR).joinpath("./audio/speech/")
SPEECH_CACHE_DIR.mkdir(parents=True, exist_ok=True)
2024-02-06 07:57:24 +01:00
file_path = SPEECH_CACHE_DIR.joinpath(f"{name}.mp3")
file_body_path = SPEECH_CACHE_DIR.joinpath(f"{name}.json")
2024-02-06 07:51:08 +01:00
# 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_KEY}"
headers["Content-Type"] = "application/json"
try:
print("openai")
r = requests.post(
url=target_url,
data=body,
headers=headers,
stream=True,
)
r.raise_for_status()
# 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:57:24 +01:00
with open(file_body_path, "w") as f:
json.dump(json.loads(body.decode("utf-8")), f)
2024-02-06 07:51:08 +01:00
# Return the saved file
return FileResponse(file_path)
except Exception as e:
print(e)
error_detail = "Ollama 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, detail=error_detail)
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_current_user)):
target_url = f"{app.state.OPENAI_API_BASE_URL}/{path}"
2024-01-05 03:38:03 +01:00
print(target_url, app.state.OPENAI_API_KEY)
2024-01-05 01:49:34 +01:00
if user.role not in ["user", "admin"]:
2024-01-22 10:37:54 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
2024-01-05 03:38:03 +01:00
if app.state.OPENAI_API_KEY == "":
2024-01-22 10:37:54 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
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)
# 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)
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-01-05 03:38:03 +01:00
headers = {}
2024-01-05 01:49:34 +01:00
headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
2024-01-05 03:54:00 +01:00
headers["Content-Type"] = "application/json"
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:
# For non-SSE, read the response and return it
# response_data = (
# r.json()
# if r.headers.get("Content-Type", "")
# == "application/json"
# else r.text
# )
response_data = r.json()
if "openai" in app.state.OPENAI_API_BASE_URL and path == "models":
response_data["data"] = list(
2024-01-22 10:37:54 +01:00
filter(lambda model: "gpt" in model["id"], response_data["data"])
)
2024-01-05 03:38:03 +01:00
return response_data
2024-01-05 01:49:34 +01:00
except Exception as e:
print(e)
error_detail = "Ollama 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, detail=error_detail)