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

120 lines
3.6 KiB
Python
Raw Normal View History

2023-11-19 01:47:12 +01:00
from flask import Flask, request, Response, jsonify
2023-11-15 01:28:51 +01:00
from flask_cors import CORS
import requests
import json
2023-11-19 01:47:12 +01:00
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
2023-11-15 01:28:51 +01:00
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("/", 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}"
2023-12-26 06:44:28 +01:00
print(target_url)
2023-11-15 01:28:51 +01:00
# Get data from the original request
data = request.get_data()
headers = dict(request.headers)
2023-11-19 06:41:43 +01:00
# Basic RBAC support
if WEBUI_AUTH:
2023-11-19 01:47:12 +01:00
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"])
2023-11-19 01:47:12 +01:00
if user:
2023-11-19 06:41:43 +01:00
# 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
2023-11-19 01:47:12 +01:00
else:
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
else:
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
else:
pass
2023-12-15 02:05:46 +01:00
r = None
2023-12-26 22:43:32 +01:00
headers.pop("Host", None)
headers.pop("Authorization", None)
headers.pop("Origin", None)
headers.pop("Referer", None)
2023-12-26 22:40:03 +01:00
try:
# Make a request to the target server
2023-12-15 02:05:46 +01:00
r = requests.request(
method=request.method,
url=target_url,
data=data,
headers=headers,
stream=True, # Enable streaming for server-sent events
)
2023-12-15 02:05:46 +01:00
r.raise_for_status()
# Proxy the target server's response to the client
def generate():
2023-12-15 02:05:46 +01:00
for chunk in r.iter_content(chunk_size=8192):
yield chunk
2023-12-15 02:05:46 +01:00
response = Response(generate(), status=r.status_code)
# Copy headers from the target server's response to the client's response
2023-12-15 02:05:46 +01:00
for key, value in r.headers.items():
response.headers[key] = value
return response
except Exception as e:
2023-12-26 22:40:03 +01:00
print(e)
2023-12-15 02:05:46 +01:00
error_detail = "Ollama WebUI: Server Connection Error"
if r != None:
2023-12-26 22:40:03 +01:00
print(r.text)
2023-12-15 02:05:46 +01:00
res = r.json()
if "error" in res:
error_detail = f"Ollama: {res['error']}"
print(res)
return (
jsonify(
{
"detail": error_detail,
"message": str(e),
}
),
400,
)
2023-11-15 01:28:51 +01:00
if __name__ == "__main__":
app.run(debug=True)