From 73d8e3e640452c5fad1abe37f1baaa4bd749f978 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 6 Feb 2024 22:18:46 -0800 Subject: [PATCH] fix(ollama_chat.py): fix token counting --- litellm/llms/ollama_chat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/ollama_chat.py b/litellm/llms/ollama_chat.py index 0311931b139..c9d6654c7c3 100644 --- a/litellm/llms/ollama_chat.py +++ b/litellm/llms/ollama_chat.py @@ -230,7 +230,7 @@ def get_ollama_response( model_response["model"] = "ollama/" + model prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=messages)) # type: ignore completion_tokens = response_json.get( - "eval_count", litellm.token_counter(text=response_json["message"]) + "eval_count", litellm.token_counter(text=response_json["message"]["content"]) ) model_response["usage"] = litellm.Usage( prompt_tokens=prompt_tokens,