From becff369dc528b26dea5b16a0a03a16790028df3 Mon Sep 17 00:00:00 2001 From: puffo Date: Thu, 18 Jan 2024 10:47:24 -0600 Subject: [PATCH] fix(ollama_chat.py): use tiktoken as backup for prompt token counting --- litellm/llms/ollama_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/ollama_chat.py b/litellm/llms/ollama_chat.py index 1ff93649f0..31e3f0d16a 100644 --- a/litellm/llms/ollama_chat.py +++ b/litellm/llms/ollama_chat.py @@ -220,7 +220,7 @@ def get_ollama_response( model_response["choices"][0]["message"] = response_json["message"] model_response["created"] = int(time.time()) model_response["model"] = "ollama/" + model - prompt_tokens = response_json["prompt_eval_count"] # type: ignore + prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore completion_tokens = response_json["eval_count"] model_response["usage"] = litellm.Usage( prompt_tokens=prompt_tokens, @@ -320,7 +320,7 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): 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 + prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore completion_tokens = response_json["eval_count"] model_response["usage"] = litellm.Usage( prompt_tokens=prompt_tokens,