This commit is contained in:
Timothy J. Baek 2024-01-22 01:37:54 -08:00
parent 60afd6ecdd
commit b26e0fb7e7

View file

@ -37,19 +37,16 @@ async def get_openai_url(user=Depends(get_current_user)):
if user and user.role == "admin": if user and user.role == "admin":
return {"OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL} return {"OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL}
else: else:
raise HTTPException(status_code=401, raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
@app.post("/url/update") @app.post("/url/update")
async def update_openai_url(form_data: UrlUpdateForm, async def update_openai_url(form_data: UrlUpdateForm, user=Depends(get_current_user)):
user=Depends(get_current_user)):
if user and user.role == "admin": if user and user.role == "admin":
app.state.OPENAI_API_BASE_URL = form_data.url app.state.OPENAI_API_BASE_URL = form_data.url
return {"OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL} return {"OPENAI_API_BASE_URL": app.state.OPENAI_API_BASE_URL}
else: else:
raise HTTPException(status_code=401, raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
@app.get("/key") @app.get("/key")
@ -57,19 +54,16 @@ async def get_openai_key(user=Depends(get_current_user)):
if user and user.role == "admin": if user and user.role == "admin":
return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY} return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY}
else: else:
raise HTTPException(status_code=401, raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
@app.post("/key/update") @app.post("/key/update")
async def update_openai_key(form_data: KeyUpdateForm, async def update_openai_key(form_data: KeyUpdateForm, user=Depends(get_current_user)):
user=Depends(get_current_user)):
if user and user.role == "admin": if user and user.role == "admin":
app.state.OPENAI_API_KEY = form_data.key app.state.OPENAI_API_KEY = form_data.key
return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY} return {"OPENAI_API_KEY": app.state.OPENAI_API_KEY}
else: else:
raise HTTPException(status_code=401, raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"]) @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
@ -78,43 +72,29 @@ async def proxy(path: str, request: Request, user=Depends(get_current_user)):
print(target_url, app.state.OPENAI_API_KEY) print(target_url, app.state.OPENAI_API_KEY)
if user.role not in ["user", "admin"]: if user.role not in ["user", "admin"]:
raise HTTPException(status_code=401, raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
if app.state.OPENAI_API_KEY == "": if app.state.OPENAI_API_KEY == "":
raise HTTPException(status_code=401, raise HTTPException(status_code=401, detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
detail=ERROR_MESSAGES.API_KEY_NOT_FOUND)
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 to decode the body of the request from bytes to a UTF-8 string (Require add max_token to fix gpt-4-vision)
try: try:
body_str = (await request.body()).decode('utf-8') body = body.decode("utf-8")
except UnicodeDecodeError as e: body = json.loads(body)
print("Error decoding request body:", e)
raise HTTPException(status_code=400, detail="Invalid request body")
# Check if the body is not empty
if body_str:
try:
body_dict = json.loads(body_str)
except json.JSONDecodeError as e:
print("Error loading request body into a dictionary:", e)
raise HTTPException(status_code=400, detail="Invalid JSON in request body")
# Check if the model is "gpt-4-vision-preview" and set "max_tokens" to 4000 # 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 # This is a workaround until OpenAI fixes the issue with this model
if body_dict.get("model") == "gpt-4-vision-preview": if body.get("model") == "gpt-4-vision-preview":
body_dict["max_tokens"] = 4000 body["max_tokens"] = 4000
print("Modified body_dict:", body_dict) print("Modified body_dict:", body)
# Try to convert the modified body back to JSON
try:
# Convert the modified body back to JSON # Convert the modified body back to JSON
body_json = json.dumps(body_dict) body = json.dumps(body)
except TypeError as e: except json.JSONDecodeError as e:
print("Error converting modified body to JSON:", e) print("Error loading request body into a dictionary:", e)
raise HTTPException(status_code=500, detail="Internal server error") raise HTTPException(status_code=400, detail="Invalid JSON in request body")
else:
body_json = body_str # If the body is empty, use it as is
headers = {} headers = {}
headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}" headers["Authorization"] = f"Bearer {app.state.OPENAI_API_KEY}"
@ -124,7 +104,7 @@ async def proxy(path: str, request: Request, user=Depends(get_current_user)):
r = requests.request( r = requests.request(
method=request.method, method=request.method,
url=target_url, url=target_url,
data=body_json, data=body,
headers=headers, headers=headers,
stream=True, stream=True,
) )
@ -153,8 +133,8 @@ async def proxy(path: str, request: Request, user=Depends(get_current_user)):
if "openai" in app.state.OPENAI_API_BASE_URL and path == "models": if "openai" in app.state.OPENAI_API_BASE_URL and path == "models":
response_data["data"] = list( response_data["data"] = list(
filter(lambda model: "gpt" in model["id"], filter(lambda model: "gpt" in model["id"], response_data["data"])
response_data["data"])) )
return response_data return response_data
except Exception as e: except Exception as e: