From 689acb8a08724c66cd1daf15187f09a585af2b10 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 11 Oct 2023 16:47:52 -0700
Subject: [PATCH] (feat) add CustomStreamWrapper for Ollama - match OpenAI
 streaming

---
 litellm/main.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/litellm/main.py b/litellm/main.py
index c5189ea46..353438ea5 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -1033,7 +1033,10 @@ def completion(
             generator = ollama.get_ollama_response_stream(api_base, model, prompt, optional_params)
             if optional_params.get("stream", False) == True:
                 # assume all ollama responses are streamed
-                return generator
+                response = CustomStreamWrapper(
+                        generator, model, custom_llm_provider="ollama", logging_obj=logging
+                )
+                return response
             else:
                 response_string = ""
                 for chunk in generator: