From 7b3ee8d1292f11fab0460f540e720d8b61b067aa Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 11 Oct 2023 16:59:28 -0700
Subject: [PATCH] (feat) ollama raise Exceptions + use LiteLLM stream wrapper

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

diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py
index c3dc762816..3a05308039 100644
--- a/litellm/llms/ollama.py
+++ b/litellm/llms/ollama.py
@@ -131,13 +131,20 @@ def get_ollama_response_stream(
                     for chunk in chunks:
                         if chunk.strip() != "":
                             j = json.loads(chunk)
+                            if "error" in j:
+                                completion_obj = {
+                                    "role": "assistant",
+                                    "content": "",
+                                    "error": j
+                                }
+                                yield completion_obj
                             if "response" in j:
                                 completion_obj = {
                                     "role": "assistant",
                                     "content": "",
                                 }
                                 completion_obj["content"] = j["response"]
-                                yield {"choices": [{"delta": completion_obj}]}
+                                yield completion_obj
                 except Exception as e:
                     traceback.print_exc()
                     print(f"Error decoding JSON: {e}")
@@ -176,6 +183,13 @@ if async_generator_imported:
                         for chunk in chunks:
                             if chunk.strip() != "":
                                 j = json.loads(chunk)
+                                if "error" in j:
+                                    completion_obj = {
+                                        "role": "assistant",
+                                        "content": "",
+                                        "error": j
+                                    }
+                                    await yield_({"choices": [{"delta": completion_obj}]})
                                 if "response" in j:
                                     completion_obj = {
                                         "role": "assistant",