From 47dc3b5fb25fc10b9c89625e3a0d3b7ef2580fed Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 5 Jan 2024 01:25:34 -0800
Subject: [PATCH] feat: async reverse proxy

---
 backend/apps/ollama/main.py | 77 +++++++++++++++++++++++++------------
 1 file changed, 52 insertions(+), 25 deletions(-)

diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
index 9310ea85..11fa707d 100644
--- a/backend/apps/ollama/main.py
+++ b/backend/apps/ollama/main.py
@@ -11,6 +11,8 @@ from constants import ERROR_MESSAGES
 from utils.utils import decode_token, get_current_user
 from config import OLLAMA_API_BASE_URL, WEBUI_AUTH
 
+import aiohttp
+
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -30,8 +32,7 @@ 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}
     else:
-        raise HTTPException(status_code=401,
-                            detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
 
 
 class UrlUpdateForm(BaseModel):
@@ -39,14 +40,29 @@ class UrlUpdateForm(BaseModel):
 
 
 @app.post("/url/update")
-async def update_ollama_api_url(form_data: UrlUpdateForm,
-                                user=Depends(get_current_user)):
+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)
+        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"])
@@ -59,42 +75,53 @@ async def proxy(path: str, request: Request, user=Depends(get_current_user)):
     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)
+                raise HTTPException(
+                    status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED
+                )
     else:
-        raise HTTPException(status_code=401,
-                            detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
+        raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
 
     headers.pop("Host", None)
     headers.pop("Authorization", None)
     headers.pop("Origin", None)
     headers.pop("Referer", None)
 
+    session = aiohttp.ClientSession()
+    response = None
     try:
-        r = requests.request(
-            method=request.method,
-            url=target_url,
-            data=body,
-            headers=headers,
-            stream=True,
+        response = await session.request(
+            request.method, target_url, data=body, headers=headers
         )
 
-        r.raise_for_status()
+        if not response.ok:
+            data = await response.json()
+            print(data)
+            response.raise_for_status()
+
+        async def gen():
+            async for line in response.content:
+                yield line
+            await session.close()
+
+        return StreamingResponse(gen(), response.status)
 
-        return StreamingResponse(
-            r.iter_content(chunk_size=8192),
-            status_code=r.status_code,
-            headers=dict(r.headers),
-        )
     except Exception as e:
         print(e)
         error_detail = "Ollama WebUI: Server Connection Error"
-        if r is not None:
+        if response is not None:
             try:
-                res = r.json()
+                res = await response.json()
                 if "error" in res:
                     error_detail = f"Ollama: {res['error']}"
             except:
                 error_detail = f"Ollama: {e}"
 
-        raise HTTPException(status_code=r.status_code, detail=error_detail)
+        await session.close()
+        raise HTTPException(
+            status_code=response.status if response else 500,
+            detail=error_detail,
+        )
+
+        # print(e)
+        # error_detail = "Ollama WebUI: Server Connection Error"
+        # return {"error": error_detail, "message": str(e)}