From be3ab88c88b120bdf144b7ab56c95799ad4e6867 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 12 Mar 2024 00:26:14 -0700
Subject: [PATCH] fix: error handling

---
 backend/apps/ollama/main.py | 17 +++++++++++++----
 backend/constants.py        |  1 +
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
index 5ecbaa297..039c570b5 100644
--- a/backend/apps/ollama/main.py
+++ b/backend/apps/ollama/main.py
@@ -123,6 +123,7 @@ async def get_all_models():
             map(lambda response: response["models"], responses)
         )
     }
+
     app.state.MODELS = {model["model"]: model for model in models["models"]}
 
     return models
@@ -181,11 +182,19 @@ async def get_ollama_versions(url_idx: Optional[int] = None):
         responses = await asyncio.gather(*tasks)
         responses = list(filter(lambda x: x is not None, responses))
 
-        lowest_version = min(
-            responses, key=lambda x: tuple(map(int, x["version"].split(".")))
-        )
+        print(responses)
 
-        return {"version": lowest_version["version"]}
+        if len(responses) > 0:
+            lowest_version = min(
+                responses, key=lambda x: tuple(map(int, x["version"].split(".")))
+            )
+
+            return {"version": lowest_version["version"]}
+        else:
+            raise HTTPException(
+                status_code=500,
+                detail=ERROR_MESSAGES.OLLAMA_NOT_FOUND,
+            )
     else:
         url = app.state.OLLAMA_BASE_URLS[url_idx]
         try:
diff --git a/backend/constants.py b/backend/constants.py
index eacf8a20f..05bdebc54 100644
--- a/backend/constants.py
+++ b/backend/constants.py
@@ -52,3 +52,4 @@ class ERROR_MESSAGES(str, Enum):
 
     MODEL_NOT_FOUND = lambda name="": f"Model '{name}' was not found"
     OPENAI_NOT_FOUND = lambda name="": f"OpenAI API was not found"
+    OLLAMA_NOT_FOUND = "WebUI could not connect to Ollama"