fix(ollama.py): fix ollama async streaming for /completions calls

litellm/llms/ollama.py

@@ -182,40 +182,28 @@ def ollama_completion_stream(url, data):
                     traceback.print_exc()
     session.close()
 
+async def iter_lines(reader):
+    buffer = b""
+    async for chunk in reader.iter_any():
+        buffer += chunk
+        while b'\n' in buffer:
+            line, buffer = buffer.split(b'\n', 1)
+            yield line
+
 async def ollama_async_streaming(url, data, model_response, encoding, logging_obj):
     try:
-        timeout = aiohttp.ClientTimeout(total=600)  # 10 minutes
-        async with aiohttp.ClientSession(timeout=timeout) as session:
-            resp = await session.post(url, json=data)
-
-            if resp.status != 200:
-                text = await resp.text()
-                raise OllamaError(status_code=resp.status, message=text)
-            
-            async for line in resp.content.iter_any():
-                if line:
-                    try:
-                        json_chunk = line.decode("utf-8")
-                        chunks = json_chunk.split("\n")
-                        completion_string = ""
-                        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": j["response"],
-                                    }
-                                    yield completion_obj
-                    except Exception as e:
-                        traceback.print_exc()
+        with httpx.stream(
+                    url=f"{url}",
+                    json=data,
+                    method="POST",
+                    timeout=litellm.request_timeout
+                ) as response: 
+                    if response.status_code != 200:
+                        raise OllamaError(status_code=response.status_code, message=response.text) 
+                    
+                    streamwrapper = litellm.CustomStreamWrapper(completion_stream=response.iter_lines(), model=data['model'], custom_llm_provider="ollama",logging_obj=logging_obj)
+                    for transformed_chunk in streamwrapper:
+                        yield transformed_chunk
     except Exception as e:
         traceback.print_exc()
 
@@ -267,7 +255,6 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
     except Exception as e:
         traceback.print_exc()
 
-if async_generator_imported:
     # ollama implementation
     @async_generator
     async def async_get_ollama_response_stream(

litellm/main.py

@@ -209,7 +209,7 @@ async def _async_streaming(response, model, custom_llm_provider, args):
             print_verbose(f"line in async streaming: {line}")
             yield line
     except Exception as e: 
-        print_verbose(f"error raised _async_streaming: {str(e)}")
+        print_verbose(f"error raised _async_streaming: {traceback.format_exc()}")
         raise exception_type(
                 model=model, custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=args,
             )
@@ -2378,6 +2378,8 @@ def stream_chunk_builder(chunks: list, messages: Optional[list]=None):
         completion_output = combined_content
     elif len(combined_arguments) > 0: 
         completion_output = combined_arguments
+    else: 
+        completion_output = ""
     # # Update usage information if needed
     try:
         response["usage"]["prompt_tokens"] = token_counter(model=model, messages=messages)

litellm/utils.py

@@ -5626,6 +5626,30 @@ class CustomStreamWrapper:
             traceback.print_exc()
             return ""
 
+    def handle_ollama_stream(self, chunk): 
+        try: 
+            json_chunk = json.loads(chunk)
+            if "error" in json_chunk: 
+                raise Exception(f"Ollama Error - {json_chunk}")
+            
+            text = "" 
+            is_finished = False
+            finish_reason = None
+            if json_chunk["done"] == True:
+                text = ""
+                is_finished = True
+                finish_reason = "stop"
+                return {"text": text, "is_finished": is_finished, "finish_reason": finish_reason}
+            elif json_chunk["response"]:
+                print_verbose(f"delta content: {json_chunk}")
+                text = json_chunk["response"]
+                return {"text": text, "is_finished": is_finished, "finish_reason": finish_reason}
+            else: 
+                raise Exception(f"Ollama Error - {json_chunk}")
+        except Exception as e: 
+            raise e
+
+
     def handle_bedrock_stream(self, chunk):
         if hasattr(chunk, "get"):
             chunk = chunk.get('chunk')
@@ -5800,9 +5824,11 @@ class CustomStreamWrapper:
                 self.completion_stream = self.completion_stream[chunk_size:]
                 time.sleep(0.05)
             elif self.custom_llm_provider == "ollama":
-                if "error" in chunk:
-                    exception_type(model=self.model, custom_llm_provider=self.custom_llm_provider, original_exception=chunk["error"])
-                completion_obj = chunk
+                response_obj = self.handle_ollama_stream(chunk)
+                completion_obj["content"] = response_obj["text"]
+                print_verbose(f"completion obj content: {completion_obj['content']}")
+                if response_obj["is_finished"]: 
+                    model_response.choices[0].finish_reason = response_obj["finish_reason"]
             elif self.custom_llm_provider == "text-completion-openai":
                 response_obj = self.handle_openai_text_completion_chunk(chunk)
                 completion_obj["content"] = response_obj["text"]
@@ -5894,7 +5920,7 @@ class CustomStreamWrapper:
     ## needs to handle the empty string case (even starting chunk can be an empty string)
     def __next__(self):
         try:
-            while True: 
+            while True:
                 if isinstance(self.completion_stream, str) or isinstance(self.completion_stream, bytes):
                     chunk = self.completion_stream
                 else:
@@ -5912,7 +5938,7 @@ class CustomStreamWrapper:
         except StopIteration:
             raise  # Re-raise StopIteration
         except Exception as e:
-            print_verbose(f"HITS AN ERROR: {str(e)}")
+            print_verbose(f"HITS AN ERROR: {str(e)}\n\n {traceback.format_exc()}")
             traceback_exception = traceback.format_exc()
             # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
             threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start()
@@ -5969,17 +5995,22 @@ class TextCompletionStreamWrapper:
         return self
     
     def convert_to_text_completion_object(self, chunk: ModelResponse):
-        response = TextCompletionResponse()
-        response["id"] = chunk.get("id", None)
-        response["object"] = "text_completion"
-        response["created"] = response.get("created", None)
-        response["model"] = response.get("model", None)
-        text_choices = TextChoices()
-        text_choices["text"] = chunk["choices"][0]["delta"]["content"]
-        text_choices["index"] = response["choices"][0]["index"]
-        text_choices["finish_reason"] = response["choices"][0]["finish_reason"]
-        response["choices"] = [text_choices]
-        return response
+        try: 
+            response = TextCompletionResponse()
+            response["id"] = chunk.get("id", None)
+            response["object"] = "text_completion"
+            response["created"] = response.get("created", None)
+            response["model"] = response.get("model", None)
+            text_choices = TextChoices()
+            if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices
+                raise Exception
+            text_choices["text"] = chunk["choices"][0]["delta"]["content"]
+            text_choices["index"] = response["choices"][0]["index"]
+            text_choices["finish_reason"] = response["choices"][0]["finish_reason"]
+            response["choices"] = [text_choices]
+            return response
+        except Exception as e:
+            raise Exception(f"Error occurred converting to text completion object - chunk: {chunk}; Error: {str(e)}")
 
     def __next__(self):
         # model_response = ModelResponse(stream=True, model=self.model)