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

144 lines
4.5 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-01-05 03:38:03 +01:00
from fastapi.responses import StreamingResponse, JSONResponse
2024-01-05 01:49:34 +01:00
import requests
import json
from pydantic import BaseModel
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
from utils.utils import decode_token, get_current_user
from config import OPENAI_API_BASE_URL, OPENAI_API_KEY
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-05 05:05:19 +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-05 05:05:19 +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-05 05:05:19 +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-05 05:05:19 +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-05 05:05:19 +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-05 05:05:19 +01:00
raise HTTPException(status_code=401,
detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
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-05 05:05:19 +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-05 05:05:19 +01:00
raise HTTPException(status_code=401,
detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
2024-01-05 01:49:34 +01:00
2024-01-05 03:38:03 +01:00
body = await request.body()
# headers = dict(request.headers)
# print(headers)
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,
data=body,
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()
print(type(response_data))
if "openai" in app.state.OPENAI_API_BASE_URL and path == "models":
response_data["data"] = list(
2024-01-05 05:05:19 +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)