From 25987fe3c901ef386d2d84c96f878396565dce10 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 13 Dec 2023 17:37:29 -0800
Subject: [PATCH] feat: better error handling for ollama reverse proxy

---
 backend/apps/ollama/main.py | 39 +++++++++++++++++++++----------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
index 53daefac0..98048ca53 100644
--- a/backend/apps/ollama/main.py
+++ b/backend/apps/ollama/main.py
@@ -59,27 +59,32 @@ def proxy(path):
     else:
         pass
 
-    # Make a request to the target server
-    target_response = requests.request(
-        method=request.method,
-        url=target_url,
-        data=data,
-        headers=headers,
-        stream=True,  # Enable streaming for server-sent events
-    )
+    try:
+        # Make a request to the target server
+        target_response = requests.request(
+            method=request.method,
+            url=target_url,
+            data=data,
+            headers=headers,
+            stream=True,  # Enable streaming for server-sent events
+        )
 
-    # Proxy the target server's response to the client
-    def generate():
-        for chunk in target_response.iter_content(chunk_size=8192):
-            yield chunk
+        target_response.raise_for_status()
 
-    response = Response(generate(), status=target_response.status_code)
+        # Proxy the target server's response to the client
+        def generate():
+            for chunk in target_response.iter_content(chunk_size=8192):
+                yield chunk
 
-    # Copy headers from the target server's response to the client's response
-    for key, value in target_response.headers.items():
-        response.headers[key] = value
+        response = Response(generate(), status=target_response.status_code)
 
-    return response
+        # Copy headers from the target server's response to the client's response
+        for key, value in target_response.headers.items():
+            response.headers[key] = value
+
+        return response
+    except Exception as e:
+        return jsonify({"detail": "Server Connection Error", "message": str(e)}), 400
 
 
 if __name__ == "__main__":