From a419d59542d4d557e9e35f6791745019a727e148 Mon Sep 17 00:00:00 2001 From: Joel Eriksson Date: Sun, 17 Dec 2023 17:27:47 +0200 Subject: [PATCH 1/2] Fix for issue that occured when proxying to ollama In the text_completion() function, it previously threw an exception at: raw_response = response._hidden_params.get("original_response", None) Due to response being an coroutine object to an ollama_acompletion call, so I added an asyncio.iscoroutine() check for the response and handle it by calling response = asyncio.run(response) I also had to fix atext_completion(), where init_response was an instance of TextCompletionResponse. Since this case was not handled by the if-elif that checks if init_response is a coroutine, a dict or a ModelResponse instance, response was unbound which threw an exception on the "return response" line. Note that a regular pyright based linter detects that response is possibly unbound, and that the same code pattern is used in multiple other places in main.py. I would suggest that you either change these cases: init_response = await loop.run_in_executor(... if isinstance(init_response, ... response = init_response elif asyncio.iscoroutine(init_response): response = await init_response To either just: response = await loop.run_in_executor( if asyncio.iscoroutine(response): response = await response Or at the very least, include an else statement and set response = init_response, so that response is never unbound when the code proceeds. --- litellm/main.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index b7e9ccce222..878d0fa5a4e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2016,11 +2016,9 @@ async def atext_completion(*args, **kwargs): response = text_completion(*args, **kwargs) else: # Await normally - init_response = await loop.run_in_executor(None, func_with_context) - if isinstance(init_response, dict) or isinstance(init_response, ModelResponse): ## CACHING SCENARIO - response = init_response - elif asyncio.iscoroutine(init_response): - response = await init_response + response = await loop.run_in_executor(None, func_with_context) + if asyncio.iscoroutine(response): + response = await response else: # Call the synchronous function using run_in_executor response = await loop.run_in_executor(None, func_with_context) @@ -2196,6 +2194,9 @@ def text_completion( response = TextCompletionStreamWrapper(completion_stream=response, model=model) return response + if asyncio.iscoroutine(response): + response = asyncio.run(response) + transformed_logprobs = None # only supported for TGI models try: From e214e6ab47c6ac9f3349ce07f6ba41a367cc5b69 Mon Sep 17 00:00:00 2001 From: Joel Eriksson Date: Sun, 17 Dec 2023 20:23:26 +0200 Subject: [PATCH 2/2] Fix bug when iterating over lines in ollama response async for line in resp.content.iter_any() will return incomplete lines when the lines are long, and that results in an exception being thrown by json.loads() when it tries to parse the incomplete JSON The default behavior of the stream reader for aiohttp response objects is to iterate over lines, so just removing .iter_any() fixes the bug --- litellm/llms/ollama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index f2a9b0df4d1..e2be1c2d5de 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -195,7 +195,7 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): raise OllamaError(status_code=resp.status, message=text) completion_string = "" - async for line in resp.content.iter_any(): + async for line in resp.content: if line: try: json_chunk = line.decode("utf-8")