From 2f0746033319373e6ee85abc1798c5bde167b8c5 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 10 Nov 2023 08:54:03 -0800
Subject: [PATCH] (feat) completion ollama raise exception when ollama resp !=
 200

---
 litellm/llms/ollama.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py
index add9c8d7f5b..c46a343b48b 100644
--- a/litellm/llms/ollama.py
+++ b/litellm/llms/ollama.py
@@ -10,6 +10,14 @@ try:
 except ImportError:
     async_generator_imported = False  # this should not throw an error, it will impact the 'import litellm' statement
 
+class OllamaError(Exception):
+    def __init__(self, status_code, message):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(
+            self.message
+        )  # Call the base class constructor with the parameters it needs
+
 class OllamaConfig():
     """
     Reference: https://github.com/jmorganca/ollama/blob/main/docs/api.md#parameters
@@ -123,6 +131,8 @@ def get_ollama_response_stream(
     session = requests.Session()
 
     with session.post(url, json=data, stream=True) as resp:
+        if resp.status_code != 200:
+            raise OllamaError(status_code=resp.status_code, message=resp.text)
         for line in resp.iter_lines():
             if line:
                 try:
@@ -174,6 +184,8 @@ if async_generator_imported:
         session = requests.Session()
 
         with session.post(url, json=data, stream=True) as resp:
+            if resp.status_code != 200:
+                raise OllamaError(status_code=resp.status_code, message=resp.text)
             for line in resp.iter_lines():
                 if line:
                     try: