From cc65ba582d8d9cf84a2a7166b0bc2f481dfd7242 Mon Sep 17 00:00:00 2001 From: Miyar <14232275+gakugaku@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:20:14 +0000 Subject: [PATCH] perf(ollama_chat): estimate token counts only when Ollama omits them `litellm.token_counter` was passed as the default argument of `dict.get`, so it ran on every response even though Ollama reports `prompt_eval_count` and `eval_count`. That spends a tokenization pass per call, and a counter failure discards a response Ollama already produced. --- litellm/llms/ollama/chat/transformation.py | 17 +++- .../ollama/test_ollama_chat_transformation.py | 99 +++++++++++++++++++ 2 files changed, 112 insertions(+), 4 deletions(-) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index d6aa1f1743b..4be3f35f6fa 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -386,10 +386,19 @@ class OllamaChatConfig(BaseConfig): model_response.choices[0].finish_reason = "tool_calls" model_response.created = int(time.time()) model_response.model = "ollama_chat/" + model - prompt_tokens = response_json.get("prompt_eval_count", litellm.token_counter(messages=messages)) - completion_tokens: Final = response_json.get( - "eval_count", - litellm.token_counter(text=response_json["message"]["content"]), + # Ollama reports both counts, so only fall back to the estimator when a count is + # missing. Passing token_counter as the default argument of dict.get evaluates it + # on every response, which costs a tokenization pass per call and lets a counter + # failure discard a response Ollama already produced. + reported_prompt_tokens: Final = response_json.get("prompt_eval_count") + prompt_tokens: Final = ( + reported_prompt_tokens if reported_prompt_tokens is not None else litellm.token_counter(messages=messages) + ) + reported_completion_tokens: Final = response_json.get("eval_count") + completion_tokens: Final = ( + reported_completion_tokens + if reported_completion_tokens is not None + else litellm.token_counter(text=response_json["message"]["content"]) ) setattr( model_response, diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py index 906c51d8064..eb40d24406b 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py @@ -906,3 +906,102 @@ class TestOllamaToolCallTransformation: assert tool_msg["content"] == "Sunny, 72°F" assert "tool_call_id" in tool_msg, "tool_call_id must be forwarded to Ollama" assert tool_msg["tool_call_id"] == "call_abc123" + + +class TestOllamaChatUsageCounts: + """Tests for how usage is taken from Ollama's response.""" + + @staticmethod + def _transform(config, ollama_response, messages): + import json + from unittest.mock import MagicMock + + from litellm.types.utils import Choices, Message, ModelResponse + + mock_response = MagicMock() + mock_response.json.return_value = ollama_response + mock_response.text = json.dumps(ollama_response) + + model_response = ModelResponse() + model_response.choices = [Choices(message=Message(content=""), index=0)] + + return config.transform_response( + model="qwen3:14b", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=messages, + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=False, + ) + + def test_reported_counts_skip_the_estimator(self): + """ + When Ollama reports both counts, `litellm.token_counter` must not run. + + Passing it as the default argument of `dict.get` evaluates it on every + response: it costs a tokenization pass per call, and a counter failure + discards a response Ollama already produced. Content types the counter + does not handle (a `video_url` block reaches Ollama's route because + `extract_images_from_message` only collects `image_url`) therefore turned + a 200 into a 500. + """ + from unittest.mock import patch + + import litellm + + ollama_response = { + "model": "qwen3:14b", + "created_at": "2025-01-11T00:00:00.000000Z", + "message": {"role": "assistant", "content": "Hello!"}, + "done": True, + "prompt_eval_count": 100, + "eval_count": 50, + } + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this."}, + { + "type": "video_url", + "video_url": {"url": "data:video/mp4;base64,AAAA"}, + }, + ], + } + ] + + with patch.object( + litellm, "token_counter", side_effect=AssertionError("estimator ran") + ): + result = self._transform(OllamaChatConfig(), ollama_response, messages) + + assert result.usage.prompt_tokens == 100 + assert result.usage.completion_tokens == 50 + assert result.usage.total_tokens == 150 + + def test_missing_counts_fall_back_to_the_estimator(self): + """When Ollama omits a count, the estimator fills it in.""" + from unittest.mock import patch + + import litellm + + ollama_response = { + "model": "qwen3:14b", + "created_at": "2025-01-11T00:00:00.000000Z", + "message": {"role": "assistant", "content": "Hello!"}, + "done": True, + } + messages = [{"role": "user", "content": "Hi"}] + + with patch.object(litellm, "token_counter", return_value=7) as counter: + result = self._transform(OllamaChatConfig(), ollama_response, messages) + + assert counter.call_count == 2 + assert result.usage.prompt_tokens == 7 + assert result.usage.completion_tokens == 7 +