From cab870f73ac184b776252c8d9cd5327230667e3a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 15 Dec 2023 09:28:18 -0800 Subject: [PATCH] fix(ollama.py): fix ollama async streaming for /completions calls --- litellm/llms/ollama.py | 53 ++++++++++++++--------------------- litellm/main.py | 4 ++- litellm/utils.py | 63 +++++++++++++++++++++++++++++++----------- 3 files changed, 70 insertions(+), 50 deletions(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index b4160088d41..6d6d774a4bb 100644 --- a/litellm/llms/ollama.py +++ b/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( diff --git a/litellm/main.py b/litellm/main.py index bc08dee9d1b..16ce2850cf8 100644 --- a/litellm/main.py +++ b/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) diff --git a/litellm/utils.py b/litellm/utils.py index c004924d25a..8ab0da414b6 100644 --- a/litellm/utils.py +++ b/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)