forked from open-webui/open-webui
feat: enable backend ollama url update
This commit is contained in:
Timothy J. Baek
parent
bdb3f50771
commit
30aff2db53
13 changed files with 554 additions and 409 deletions
|
|
@ -1,69 +1,68 @@
|
|||
from flask import Flask, request, Response, jsonify
|
||||
from flask_cors import CORS
|
||||
|
||||
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
|
||||
from utils.utils import decode_token, get_current_user
|
||||
from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
|
||||
|
||||
app = Flask(__name__)
|
||||
CORS(
|
||||
app
|
||||
) # Enable Cross-Origin Resource Sharing (CORS) to allow requests from different domains
|
||||
app = FastAPI()
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Define the target server URL
|
||||
TARGET_SERVER_URL = OLLAMA_API_BASE_URL
|
||||
app.state.OLLAMA_API_BASE_URL = OLLAMA_API_BASE_URL
|
||||
|
||||
# TARGET_SERVER_URL = OLLAMA_API_BASE_URL
|
||||
|
||||
|
||||
@app.route("/", defaults={"path": ""}, methods=["GET", "POST", "PUT", "DELETE"])
|
||||
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE"])
|
||||
def proxy(path):
|
||||
# Combine the base URL of the target server with the requested path
|
||||
target_url = f"{TARGET_SERVER_URL}/{path}"
|
||||
print(target_url)
|
||||
@app.get("/url")
|
||||
async def get_ollama_api_url(user=Depends(get_current_user)):
|
||||
if user and user.role == "admin":
|
||||
return {"OLLAMA_API_BASE_URL": app.state.OLLAMA_API_BASE_URL}
|
||||
else:
|
||||
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
|
||||
|
||||
# Get data from the original request
|
||||
data = request.get_data()
|
||||
|
||||
class UrlUpdateForm(BaseModel):
|
||||
url: str
|
||||
|
||||
|
||||
@app.post("/url/update")
|
||||
async def update_ollama_api_url(
|
||||
form_data: UrlUpdateForm, user=Depends(get_current_user)
|
||||
):
|
||||
if user and user.role == "admin":
|
||||
app.state.OLLAMA_API_BASE_URL = form_data.url
|
||||
return {"OLLAMA_API_BASE_URL": app.state.OLLAMA_API_BASE_URL}
|
||||
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.OLLAMA_API_BASE_URL}/{path}"
|
||||
|
||||
body = await request.body()
|
||||
headers = dict(request.headers)
|
||||
|
||||
# Basic RBAC support
|
||||
if WEBUI_AUTH:
|
||||
if "Authorization" in headers:
|
||||
_, credentials = headers["Authorization"].split()
|
||||
token_data = decode_token(credentials)
|
||||
if token_data is None or "email" not in token_data:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
|
||||
user = Users.get_user_by_email(token_data["email"])
|
||||
if user:
|
||||
# Only user and admin roles can access
|
||||
if user.role in ["user", "admin"]:
|
||||
if path in ["pull", "delete", "push", "copy", "create"]:
|
||||
# Only admin role can perform actions above
|
||||
if user.role == "admin":
|
||||
pass
|
||||
else:
|
||||
return (
|
||||
jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}),
|
||||
401,
|
||||
)
|
||||
else:
|
||||
pass
|
||||
else:
|
||||
return jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}), 401
|
||||
else:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
else:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
if user.role in ["user", "admin"]:
|
||||
if path in ["pull", "delete", "push", "copy", "create"]:
|
||||
if user.role != "admin":
|
||||
raise HTTPException(
|
||||
status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
|
||||
)
|
||||
else:
|
||||
pass
|
||||
|
||||
r = None
|
||||
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
|
||||
|
||||
headers.pop("Host", None)
|
||||
headers.pop("Authorization", None)
|
||||
|
|
@ -71,49 +70,30 @@ def proxy(path):
|
|||
headers.pop("Referer", None)
|
||||
|
||||
try:
|
||||
# Make a request to the target server
|
||||
r = requests.request(
|
||||
method=request.method,
|
||||
url=target_url,
|
||||
data=data,
|
||||
data=body,
|
||||
headers=headers,
|
||||
stream=True, # Enable streaming for server-sent events
|
||||
stream=True,
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
# Proxy the target server's response to the client
|
||||
def generate():
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
yield chunk
|
||||
|
||||
response = Response(generate(), status=r.status_code)
|
||||
|
||||
# Copy headers from the target server's response to the client's response
|
||||
for key, value in r.headers.items():
|
||||
response.headers[key] = value
|
||||
|
||||
return response
|
||||
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 != None:
|
||||
print(r.text)
|
||||
res = r.json()
|
||||
if "error" in res:
|
||||
error_detail = f"Ollama: {res['error']}"
|
||||
print(res)
|
||||
if r is not None:
|
||||
try:
|
||||
res = r.json()
|
||||
if "error" in res:
|
||||
error_detail = f"Ollama: {res['error']}"
|
||||
except:
|
||||
error_detail = f"Ollama: {e}"
|
||||
|
||||
return (
|
||||
jsonify(
|
||||
{
|
||||
"detail": error_detail,
|
||||
"message": str(e),
|
||||
}
|
||||
),
|
||||
400,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(debug=True)
|
||||
raise HTTPException(status_code=r.status_code, detail=error_detail)
|
||||
|
|
|
|||
176
backend/apps/ollama/old_main.py
Normal file
176
backend/apps/ollama/old_main.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
from flask import Flask, request, Response, jsonify
|
||||
from flask_cors import CORS
|
||||
|
||||
|
||||
import requests
|
||||
import json
|
||||
|
||||
|
||||
from apps.web.models.users import Users
|
||||
from constants import ERROR_MESSAGES
|
||||
from utils.utils import decode_token
|
||||
from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
|
||||
|
||||
app = Flask(__name__)
|
||||
CORS(
|
||||
app
|
||||
) # Enable Cross-Origin Resource Sharing (CORS) to allow requests from different domains
|
||||
|
||||
# Define the target server URL
|
||||
TARGET_SERVER_URL = OLLAMA_API_BASE_URL
|
||||
|
||||
|
||||
@app.route("/url", methods=["GET"])
|
||||
def get_ollama_api_url():
|
||||
headers = dict(request.headers)
|
||||
if "Authorization" in headers:
|
||||
_, credentials = headers["Authorization"].split()
|
||||
token_data = decode_token(credentials)
|
||||
if token_data is None or "email" not in token_data:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
|
||||
user = Users.get_user_by_email(token_data["email"])
|
||||
if user and user.role == "admin":
|
||||
return (
|
||||
jsonify({"OLLAMA_API_BASE_URL": TARGET_SERVER_URL}),
|
||||
200,
|
||||
)
|
||||
else:
|
||||
return (
|
||||
jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}),
|
||||
401,
|
||||
)
|
||||
else:
|
||||
return (
|
||||
jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}),
|
||||
401,
|
||||
)
|
||||
|
||||
|
||||
@app.route("/url/update", methods=["POST"])
|
||||
def update_ollama_api_url():
|
||||
headers = dict(request.headers)
|
||||
data = request.get_json(force=True)
|
||||
|
||||
if "Authorization" in headers:
|
||||
_, credentials = headers["Authorization"].split()
|
||||
token_data = decode_token(credentials)
|
||||
if token_data is None or "email" not in token_data:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
|
||||
user = Users.get_user_by_email(token_data["email"])
|
||||
if user and user.role == "admin":
|
||||
TARGET_SERVER_URL = data["url"]
|
||||
return (
|
||||
jsonify({"OLLAMA_API_BASE_URL": TARGET_SERVER_URL}),
|
||||
200,
|
||||
)
|
||||
else:
|
||||
return (
|
||||
jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}),
|
||||
401,
|
||||
)
|
||||
else:
|
||||
return (
|
||||
jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}),
|
||||
401,
|
||||
)
|
||||
|
||||
|
||||
@app.route("/", defaults={"path": ""}, methods=["GET", "POST", "PUT", "DELETE"])
|
||||
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE"])
|
||||
def proxy(path):
|
||||
# Combine the base URL of the target server with the requested path
|
||||
target_url = f"{TARGET_SERVER_URL}/{path}"
|
||||
print(target_url)
|
||||
|
||||
# Get data from the original request
|
||||
data = request.get_data()
|
||||
headers = dict(request.headers)
|
||||
|
||||
# Basic RBAC support
|
||||
if WEBUI_AUTH:
|
||||
if "Authorization" in headers:
|
||||
_, credentials = headers["Authorization"].split()
|
||||
token_data = decode_token(credentials)
|
||||
if token_data is None or "email" not in token_data:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
|
||||
user = Users.get_user_by_email(token_data["email"])
|
||||
if user:
|
||||
# Only user and admin roles can access
|
||||
if user.role in ["user", "admin"]:
|
||||
if path in ["pull", "delete", "push", "copy", "create"]:
|
||||
# Only admin role can perform actions above
|
||||
if user.role == "admin":
|
||||
pass
|
||||
else:
|
||||
return (
|
||||
jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}),
|
||||
401,
|
||||
)
|
||||
else:
|
||||
pass
|
||||
else:
|
||||
return jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}), 401
|
||||
else:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
else:
|
||||
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
|
||||
else:
|
||||
pass
|
||||
|
||||
r = None
|
||||
|
||||
headers.pop("Host", None)
|
||||
headers.pop("Authorization", None)
|
||||
headers.pop("Origin", None)
|
||||
headers.pop("Referer", None)
|
||||
|
||||
try:
|
||||
# Make a request to the target server
|
||||
r = requests.request(
|
||||
method=request.method,
|
||||
url=target_url,
|
||||
data=data,
|
||||
headers=headers,
|
||||
stream=True, # Enable streaming for server-sent events
|
||||
)
|
||||
|
||||
r.raise_for_status()
|
||||
|
||||
# Proxy the target server's response to the client
|
||||
def generate():
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
yield chunk
|
||||
|
||||
response = Response(generate(), status=r.status_code)
|
||||
|
||||
# Copy headers from the target server's response to the client's response
|
||||
for key, value in r.headers.items():
|
||||
response.headers[key] = value
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
print(e)
|
||||
error_detail = "Ollama WebUI: Server Connection Error"
|
||||
if r != None:
|
||||
print(r.text)
|
||||
res = r.json()
|
||||
if "error" in res:
|
||||
error_detail = f"Ollama: {res['error']}"
|
||||
print(res)
|
||||
|
||||
return (
|
||||
jsonify(
|
||||
{
|
||||
"detail": error_detail,
|
||||
"message": str(e),
|
||||
}
|
||||
),
|
||||
400,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(debug=True)
|
||||
Loading…
Add table
Add a link
Reference in a new issue