From 3839213d28dcfdf500c321683467e2793e1b02c8 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 26 Dec 2023 20:01:15 +0530 Subject: [PATCH] (feat) ollama_chat acompletion without streaming --- litellm/llms/ollama_chat.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/litellm/llms/ollama_chat.py b/litellm/llms/ollama_chat.py index 16d234f7a..be5a20cea 100644 --- a/litellm/llms/ollama_chat.py +++ b/litellm/llms/ollama_chat.py @@ -283,18 +283,19 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): text = await resp.text() raise OllamaError(status_code=resp.status, message=text) + response_json = await resp.json() + ## LOGGING logging_obj.post_call( - input=data["prompt"], + input=data, api_key="", - original_response=resp.text, + original_response=response_json, additional_args={ "headers": None, "api_base": url, }, ) - response_json = await resp.json() ## RESPONSE OBJECT model_response["choices"][0]["finish_reason"] = "stop" if data.get("format", "") == "json": @@ -313,9 +314,7 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): ) model_response["choices"][0]["message"] = message else: - model_response["choices"][0]["message"]["content"] = response_json[ - "response" - ] + model_response["choices"][0]["message"] = response_json["message"] model_response["created"] = int(time.time()) model_response["model"] = "ollama/" + data["model"] prompt_tokens = response_json["prompt_eval_count"] # type: ignore