From 5e4dc98f4490808df0e234ac23137f44c75c546a Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Thu, 4 Jan 2024 16:49:34 -0800
Subject: [PATCH] feat: openai backend support

---
 backend/apps/openai/main.py | 113 ++++++++++++++++++++++++++++++++++++
 backend/config.py           |   8 +++
 2 files changed, 121 insertions(+)
 create mode 100644 backend/apps/openai/main.py

diff --git a/backend/apps/openai/main.py b/backend/apps/openai/main.py
new file mode 100644
index 00000000..0a12137a
--- /dev/null
+++ b/backend/apps/openai/main.py
@@ -0,0 +1,113 @@
+from fastapi import FastAPI, Request, Response, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+
+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:
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+
+
+@app.post("/url/update")
+async def update_openai_url(form_data: UrlUpdateForm, user=Depends(get_current_user)):
+    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:
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+
+
+@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:
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+
+
+@app.post("/key/update")
+async def update_openai_key(form_data: KeyUpdateForm, user=Depends(get_current_user)):
+    if user and user.role == "admin":
+        app.state.OPENAI_API_KEY = form_data.key
+        return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY}
+    else:
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+
+
+@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}"
+
+    body = await request.body()
+    headers = dict(request.headers)
+
+    if user.role not in ["user", "admin"]:
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+
+    headers.pop("Host", None)
+    headers.pop("Authorization", None)
+    headers.pop("Origin", None)
+    headers.pop("Referer", None)
+
+    headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
+
+    try:
+        r = requests.request(
+            method=request.method,
+            url=target_url,
+            data=body,
+            headers=headers,
+            stream=True,
+        )
+
+        r.raise_for_status()
+
+        return StreamingResponse(
+            r.iter_content(chunk_size=8192),
+            status_code=r.status_code,
+            headers=dict(r.headers),
+        )
+    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)
diff --git a/backend/config.py b/backend/config.py
index 8e100fe5..90900b9d 100644
--- a/backend/config.py
+++ b/backend/config.py
@@ -27,6 +27,14 @@ if ENV == "prod":
     if OLLAMA_API_BASE_URL == "/ollama/api":
         OLLAMA_API_BASE_URL = "http://host.docker.internal:11434/api"
 
+
+####################################
+# OPENAI_API
+####################################
+
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
+OPENAI_API_BASE_URL = os.environ.get("OPENAI_API_BASE_URL", "https://api.openai.com/v1")
+
 ####################################
 # WEBUI_VERSION
 ####################################