diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 9404dd920a..b1bf77f0e3 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -141,8 +141,12 @@ def get_ollama_response_stream( additional_args={"api_base": url, "complete_input_dict": data, "headers": {}, "acompletion": acompletion,}, ) if acompletion is True: - response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj) + if optional_params.get("stream", False): + response = ollama_async_streaming(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj) + else: + response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj) return response + else: return ollama_completion_stream(url=url, data=data) @@ -178,8 +182,45 @@ def ollama_completion_stream(url, data): traceback.print_exc() session.close() -async def ollama_acompletion(url, data, model_response, encoding, logging_obj): +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() + except Exception as e: + traceback.print_exc() + + +async def ollama_acompletion(url, data, model_response, encoding, logging_obj): try: timeout = aiohttp.ClientTimeout(total=600) # 10 minutes async with aiohttp.ClientSession(timeout=timeout) as session: