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

128 lines
3.8 KiB
Python
Raw Normal View History

2024-01-06 02:16:35 +01:00
from fastapi import FastAPI, Request, Response, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
2024-01-04 22:06:31 +01:00
import requests
import json
2024-01-06 02:16:35 +01:00
from pydantic import BaseModel
2024-01-04 22:06:31 +01:00
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
2024-01-06 02:16:35 +01:00
from utils.utils import decode_token, get_current_user
2024-01-04 22:06:31 +01:00
from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
2024-01-06 02:16:35 +01:00
import aiohttp
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
app.state.OLLAMA_API_BASE_URL = OLLAMA_API_BASE_URL
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
# TARGET_SERVER_URL = OLLAMA_API_BASE_URL
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
@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}
2024-01-04 22:06:31 +01:00
else:
2024-01-06 02:16:35 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
class UrlUpdateForm(BaseModel):
url: str
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
@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)
# 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}"
2024-01-04 22:06:31 +01:00
print(target_url)
2024-01-06 02:16:35 +01:00
body = await request.body()
2024-01-04 22:06:31 +01:00
headers = dict(request.headers)
2024-01-06 02:16:35 +01:00
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
)
2024-01-04 22:06:31 +01:00
else:
2024-01-06 02:16:35 +01:00
raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
2024-01-04 22:06:31 +01:00
headers.pop("Host", None)
headers.pop("Authorization", None)
headers.pop("Origin", None)
headers.pop("Referer", None)
2024-01-06 02:16:35 +01:00
session = aiohttp.ClientSession()
response = None
2024-01-04 22:06:31 +01:00
try:
2024-01-06 02:16:35 +01:00
response = await session.request(
request.method, target_url, data=body, headers=headers
2024-01-04 22:06:31 +01:00
)
2024-01-06 02:16:35 +01:00
print(response)
if not response.ok:
data = await response.json()
print(data)
response.raise_for_status()
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
async def generate():
async for line in response.content:
print(line)
yield line
await session.close()
2024-01-04 22:06:31 +01:00
2024-01-06 02:16:35 +01:00
return StreamingResponse(generate(), response.status)
2024-01-04 22:06:31 +01:00
except Exception as e:
print(e)
error_detail = "Ollama WebUI: Server Connection Error"
2024-01-06 02:16:35 +01:00
if response is not None:
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,
)