main.py 3.69 KB
Newer Older
1
from flask import Flask, request, Response, jsonify
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
4
5
6
from flask_cors import CORS

import requests
import json

7
8
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
9
from utils.utils import decode_token
10
from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
Timothy J. Baek's avatar
Timothy J. Baek committed
11
12
13
14
15
16
17
18
19
20

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


21
22
23
@app.route("/",
           defaults={"path": ""},
           methods=["GET", "POST", "PUT", "DELETE"])
Timothy J. Baek's avatar
Timothy J. Baek committed
24
25
26
27
@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}"
Timothy J. Baek's avatar
Timothy J. Baek committed
28
    print(target_url)
Timothy J. Baek's avatar
Timothy J. Baek committed
29
30
31
32
33

    # Get data from the original request
    data = request.get_data()
    headers = dict(request.headers)

Timothy J. Baek's avatar
Timothy J. Baek committed
34
    # Basic RBAC support
35
    if WEBUI_AUTH:
36
        if "Authorization" in headers:
37
38
39
40
41
42
            _, 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"])
43
            if user:
Timothy J. Baek's avatar
Timothy J. Baek committed
44
45
46
47
48
49
50
51
                # 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 (
52
53
54
55
                                jsonify({
                                    "detail":
                                    ERROR_MESSAGES.ACCESS_PROHIBITED
                                }),
Timothy J. Baek's avatar
Timothy J. Baek committed
56
57
58
59
60
                                401,
                            )
                    else:
                        pass
                else:
61
62
                    return jsonify(
                        {"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}), 401
63
64
65
66
67
68
69
            else:
                return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
        else:
            return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
    else:
        pass

70
71
    r = None

Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
    headers.pop("Host", None)
    headers.pop("Authorization", None)
    headers.pop("Origin", None)
    headers.pop("Referer", None)
Timothy J. Baek's avatar
Timothy J. Baek committed
76

77
78
    try:
        # Make a request to the target server
79
        r = requests.request(
80
81
82
83
84
85
86
            method=request.method,
            url=target_url,
            data=data,
            headers=headers,
            stream=True,  # Enable streaming for server-sent events
        )

87
        r.raise_for_status()
88
89
90

        # Proxy the target server's response to the client
        def generate():
91
            for chunk in r.iter_content(chunk_size=8192):
92
93
                yield chunk

94
        response = Response(generate(), status=r.status_code)
95
96

        # Copy headers from the target server's response to the client's response
97
        for key, value in r.headers.items():
98
99
100
101
            response.headers[key] = value

        return response
    except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
102
        print(e)
103
104
        error_detail = "Ollama WebUI: Server Connection Error"
        if r != None:
Timothy J. Baek's avatar
Timothy J. Baek committed
105
            print(r.text)
106
107
108
109
110
111
            res = r.json()
            if "error" in res:
                error_detail = f"Ollama: {res['error']}"
            print(res)

        return (
112
113
114
115
            jsonify({
                "detail": error_detail,
                "message": str(e),
            }),
116
117
            400,
        )
Timothy J. Baek's avatar
Timothy J. Baek committed
118
119
120
121


if __name__ == "__main__":
    app.run(debug=True)