Skip to content

GitLab

"...composable_kernel_onnxruntime.git" did not exist on "70456328851dea8e7c6d908193e2d32c2d7c509c"
Commit bb297126 authored by Timothy J. Baek's avatar Timothy J. Baek
Browse files

fix: backend proxy

parent 439185be
Hide whitespace changes
Inline Side-by-side
Showing with 120 additions and 188 deletions
Dockerfile
View file @ bb297126
...@@ -12,10 +12,9 @@ RUN npm run build ...@@ -12,10 +12,9 @@ RUN npm run build
FROM python:3.11-slim-buster as base FROM python:3.11-slim-buster as base
ARG OLLAMA_API_BASE_URL='/ollama/api'
ENV ENV=prod ENV ENV=prod
ENV OLLAMA_API_BASE_URL $OLLAMA_API_BASE_URL
ENV OLLAMA_API_BASE_URL "/ollama/api"
ENV OPENAI_API_BASE_URL "" ENV OPENAI_API_BASE_URL ""
ENV OPENAI_API_KEY "" ENV OPENAI_API_KEY ""
......
backend/apps/ollama/main.py
View file @ bb297126
from fastapi import FastAPI, Request, Response, HTTPException, Depends from fastapi import FastAPI, Request, Response, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from fastapi.concurrency import run_in_threadpool
import requests import requests
import json import json
...@@ -11,8 +12,6 @@ from constants import ERROR_MESSAGES ...@@ -11,8 +12,6 @@ from constants import ERROR_MESSAGES
from utils.utils import decode_token, get_current_user from utils.utils import decode_token, get_current_user
from config import OLLAMA_API_BASE_URL, WEBUI_AUTH from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
import aiohttp
app = FastAPI() app = FastAPI()
app.add_middleware( app.add_middleware(
CORSMiddleware, CORSMiddleware,
...@@ -50,25 +49,9 @@ async def update_ollama_api_url( ...@@ -50,25 +49,9 @@ async def update_ollama_api_url(
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED) raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
# async def fetch_sse(method, target_url, body, headers):
# async with aiohttp.ClientSession() as session:
# try:
# async with session.request(
# method, target_url, data=body, headers=headers
# ) as response:
# print(response.status)
# async for line in response.content:
# yield line
# except Exception as e:
# print(e)
# error_detail = "Ollama WebUI: Server Connection Error"
# yield json.dumps({"error": error_detail, "message": str(e)}).encode()
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"]) @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy(path: str, request: Request, user=Depends(get_current_user)): async def proxy(path: str, request: Request, user=Depends(get_current_user)):
target_url = f"{app.state.OLLAMA_API_BASE_URL}/{path}" target_url = f"{app.state.OLLAMA_API_BASE_URL}/{path}"
print(target_url)
body = await request.body() body = await request.body()
headers = dict(request.headers) headers = dict(request.headers)
...@@ -87,41 +70,42 @@ async def proxy(path: str, request: Request, user=Depends(get_current_user)): ...@@ -87,41 +70,42 @@ async def proxy(path: str, request: Request, user=Depends(get_current_user)):
headers.pop("Origin", None) headers.pop("Origin", None)
headers.pop("Referer", None) headers.pop("Referer", None)
session = aiohttp.ClientSession() r = None
response = None
try: def get_request():
response = await session.request( nonlocal r
request.method, target_url, data=body, headers=headers try:
) r = requests.request(
method=request.method,
print(response) url=target_url,
if not response.ok: data=body,
data = await response.json() headers=headers,
print(data) stream=True,
response.raise_for_status() )
async def generate(): r.raise_for_status()
async for line in response.content:
print(line) return StreamingResponse(
yield line r.iter_content(chunk_size=8192),
await session.close() status_code=r.status_code,
headers=dict(r.headers),
return StreamingResponse(generate(), response.status) )
except Exception as e:
raise e
try:
return await run_in_threadpool(get_request)
except Exception as e: except Exception as e:
print(e)
error_detail = "Ollama WebUI: Server Connection Error" error_detail = "Ollama WebUI: Server Connection Error"
if r is not None:
if response is not None:
try: try:
res = await response.json() res = r.json()
if "error" in res: if "error" in res:
error_detail = f"Ollama: {res['error']}" error_detail = f"Ollama: {res['error']}"
except: except:
error_detail = f"Ollama: {e}" error_detail = f"Ollama: {e}"
await session.close()
raise HTTPException( raise HTTPException(
status_code=response.status if response else 500, status_code=r.status_code if r else 500,
detail=error_detail, detail=error_detail,
) )
backend/apps/ollama/old_main.py
View file @ bb297126
from flask import Flask, request, Response, jsonify from fastapi import FastAPI, Request, Response, HTTPException, Depends
from flask_cors import CORS from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
import requests import requests
import json import json
from pydantic import BaseModel
from apps.web.models.users import Users from apps.web.models.users import Users
from constants import ERROR_MESSAGES 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 from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
app = Flask(__name__) import aiohttp
CORS(
app
) # Enable Cross-Origin Resource Sharing (CORS) to allow requests from different domains
# Define the target server URL app = FastAPI()
TARGET_SERVER_URL = OLLAMA_API_BASE_URL app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.state.OLLAMA_API_BASE_URL = OLLAMA_API_BASE_URL
@app.route("/url", methods=["GET"]) # TARGET_SERVER_URL = OLLAMA_API_BASE_URL
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"]) @app.get("/url")
def update_ollama_api_url(): async def get_ollama_api_url(user=Depends(get_current_user)):
headers = dict(request.headers) if user and user.role == "admin":
data = request.get_json(force=True) return {"OLLAMA_API_BASE_URL": app.state.OLLAMA_API_BASE_URL}
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: else:
return ( raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}),
401,
) class UrlUpdateForm(BaseModel):
url: str
@app.route("/", @app.post("/url/update")
defaults={"path": ""}, async def update_ollama_api_url(
methods=["GET", "POST", "PUT", "DELETE"]) form_data: UrlUpdateForm, user=Depends(get_current_user)
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE"]) ):
def proxy(path): if user and user.role == "admin":
# Combine the base URL of the target server with the requested path app.state.OLLAMA_API_BASE_URL = form_data.url
target_url = f"{TARGET_SERVER_URL}/{path}" return {"OLLAMA_API_BASE_URL": app.state.OLLAMA_API_BASE_URL}
else:
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
# async def fetch_sse(method, target_url, body, headers):
# async with aiohttp.ClientSession() as session:
# try:
# async with session.request(
# method, target_url, data=body, headers=headers
# ) as response:
# print(response.status)
# async for line in response.content:
# yield line
# except Exception as e:
# print(e)
# error_detail = "Ollama WebUI: Server Connection Error"
# yield json.dumps({"error": error_detail, "message": str(e)}).encode()
@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}"
print(target_url) print(target_url)
# Get data from the original request body = await request.body()
data = request.get_data()
headers = dict(request.headers) headers = dict(request.headers)
# Basic RBAC support if user.role in ["user", "admin"]:
if WEBUI_AUTH: if path in ["pull", "delete", "push", "copy", "create"]:
if "Authorization" in headers: if user.role != "admin":
_, credentials = headers["Authorization"].split() raise HTTPException(
token_data = decode_token(credentials) status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
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: else:
pass raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
r = None
headers.pop("Host", None) headers.pop("Host", None)
headers.pop("Authorization", None) headers.pop("Authorization", None)
headers.pop("Origin", None) headers.pop("Origin", None)
headers.pop("Referer", None) headers.pop("Referer", None)
session = aiohttp.ClientSession()
response = None
try: try:
# Make a request to the target server response = await session.request(
r = requests.request( request.method, target_url, data=body, headers=headers
method=request.method,
url=target_url,
data=data,
headers=headers,
stream=True, # Enable streaming for server-sent events
) )
r.raise_for_status() print(response)
if not response.ok:
# Proxy the target server's response to the client data = await response.json()
def generate(): print(data)
for chunk in r.iter_content(chunk_size=8192): response.raise_for_status()
yield chunk
response = Response(generate(), status=r.status_code) async def generate():
async for line in response.content:
print(line)
yield line
await session.close()
# Copy headers from the target server's response to the client's response return StreamingResponse(generate(), response.status)
for key, value in r.headers.items():
response.headers[key] = value
return response
except Exception as e: except Exception as e:
print(e) print(e)
error_detail = "Ollama WebUI: Server Connection Error" 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__": if response is not None:
app.run(debug=True) try:
res = await response.json()
if "error" in res:
error_detail = f"Ollama: {res['error']}"
except:
error_detail = f"Ollama: {e}"
await session.close()
raise HTTPException(
status_code=response.status if response else 500,
detail=error_detail,
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment