main.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. from fastapi import FastAPI, Request, Response, HTTPException, Depends
  2. from fastapi.middleware.cors import CORSMiddleware
  3. from fastapi.responses import StreamingResponse
  4. import requests
  5. import json
  6. from pydantic import BaseModel
  7. from apps.web.models.users import Users
  8. from constants import ERROR_MESSAGES
  9. from utils.utils import decode_token, get_current_user
  10. from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
  11. app = FastAPI()
  12. app.add_middleware(
  13. CORSMiddleware,
  14. allow_origins=["*"],
  15. allow_credentials=True,
  16. allow_methods=["*"],
  17. allow_headers=["*"],
  18. )
  19. app.state.OLLAMA_API_BASE_URL = OLLAMA_API_BASE_URL
  20. # TARGET_SERVER_URL = OLLAMA_API_BASE_URL
  21. @app.get("/url")
  22. async def get_ollama_api_url(user=Depends(get_current_user)):
  23. if user and user.role == "admin":
  24. return {"OLLAMA_API_BASE_URL": app.state.OLLAMA_API_BASE_URL}
  25. else:
  26. raise HTTPException(status_code=401,
  27. detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
  28. class UrlUpdateForm(BaseModel):
  29. url: str
  30. @app.post("/url/update")
  31. async def update_ollama_api_url(form_data: UrlUpdateForm,
  32. user=Depends(get_current_user)):
  33. if user and user.role == "admin":
  34. app.state.OLLAMA_API_BASE_URL = form_data.url
  35. return {"OLLAMA_API_BASE_URL": app.state.OLLAMA_API_BASE_URL}
  36. else:
  37. raise HTTPException(status_code=401,
  38. detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
  39. @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
  40. async def proxy(path: str, request: Request, user=Depends(get_current_user)):
  41. target_url = f"{app.state.OLLAMA_API_BASE_URL}/{path}"
  42. body = await request.body()
  43. headers = dict(request.headers)
  44. if user.role in ["user", "admin"]:
  45. if path in ["pull", "delete", "push", "copy", "create"]:
  46. if user.role != "admin":
  47. raise HTTPException(status_code=401,
  48. detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
  49. else:
  50. raise HTTPException(status_code=401,
  51. detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
  52. headers.pop("Host", None)
  53. headers.pop("Authorization", None)
  54. headers.pop("Origin", None)
  55. headers.pop("Referer", None)
  56. try:
  57. r = requests.request(
  58. method=request.method,
  59. url=target_url,
  60. data=body,
  61. headers=headers,
  62. stream=True,
  63. )
  64. r.raise_for_status()
  65. return StreamingResponse(
  66. r.iter_content(chunk_size=8192),
  67. status_code=r.status_code,
  68. headers=dict(r.headers),
  69. )
  70. except Exception as e:
  71. print(e)
  72. error_detail = "Ollama WebUI: Server Connection Error"
  73. if r is not None:
  74. try:
  75. res = r.json()
  76. if "error" in res:
  77. error_detail = f"Ollama: {res['error']}"
  78. except:
  79. error_detail = f"Ollama: {e}"
  80. raise HTTPException(status_code=r.status_code, detail=error_detail)