From d2c1ed9af85d8362ec29adf85bf2cbe7ad1c2c91 Mon Sep 17 00:00:00 2001 From: rxchitrx Date: Sat, 17 Oct 2026 21:58:38 +0530 Subject: [PATCH] fix(ollama): report prompt cache hits as cached_tokens Map prompt_eval_cached_count to prompt_tokens_details.cached_tokens for chat and completion, streaming and non-streaming, so cached prompt tokens are accounted without double billing. --- litellm/llms/ollama/chat/transformation.py | 4 ++ .../llms/ollama/completion/transformation.py | 4 ++ .../ollama/test_ollama_chat_transformation.py | 54 +++++++++++++++++++ .../test_ollama_completion_transformation.py | 50 +++++++++++++++++ 4 files changed, 112 insertions(+) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 644ab187f15..fcb2e990cad 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -388,6 +388,7 @@ class OllamaChatConfig(BaseConfig): "eval_count", litellm.token_counter(text=response_json["message"]["content"]), ) + cached_tokens: Final = response_json.get("prompt_eval_cached_count", 0) setattr( model_response, "usage", @@ -395,6 +396,7 @@ class OllamaChatConfig(BaseConfig): prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details={"cached_tokens": cached_tokens}, ), ) return model_response @@ -420,10 +422,12 @@ def _done_chunk_usage(chunk: Mapping[str, object]) -> ChatCompletionUsageBlock | eval_count: Final = chunk.get("eval_count") if chunk.get("done") is not True or not isinstance(prompt_eval_count, int) or not isinstance(eval_count, int): return None + cached_tokens: Final = chunk.get("prompt_eval_cached_count", 0) return ChatCompletionUsageBlock( prompt_tokens=prompt_eval_count, completion_tokens=eval_count, total_tokens=prompt_eval_count + eval_count, + prompt_tokens_details={"cached_tokens": cached_tokens}, ) diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index ac41310fb61..7b7cb0f6e31 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -392,6 +392,7 @@ class OllamaConfig(BaseConfig): completion_tokens: Final = response_json.get( "eval_count", len((response_json.get("message") or {}).get("content", "")) ) + cached_tokens: Final = response_json.get("prompt_eval_cached_count", 0) setattr( model_response, "usage", @@ -399,6 +400,7 @@ class OllamaConfig(BaseConfig): prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details={"cached_tokens": cached_tokens}, ), ) return model_response @@ -589,10 +591,12 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): usage: ChatCompletionUsageBlock | None = None if prompt_eval_count is not None and eval_count is not None: + cached_tokens: Final = chunk.get("prompt_eval_cached_count", 0) usage = ChatCompletionUsageBlock( prompt_tokens=prompt_eval_count, completion_tokens=eval_count, total_tokens=prompt_eval_count + eval_count, + prompt_tokens_details={"cached_tokens": cached_tokens}, ) if self.held_content: return self._flush_held_content(usage) diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 1a0bbe171a9..57bcb2301eb 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -969,6 +969,24 @@ class TestOllamaStreamingUsage: assert result.usage is not None assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (100, 50, 150) + def test_done_chunk_reports_cached_tokens_ollama_sent(self): + result: Final = self._parse( + { + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 6024, + "prompt_eval_cached_count": 6016, + "eval_count": 50, + } + ) + + assert result.usage is not None + assert result.usage.prompt_tokens == 6024 + assert result.usage.prompt_tokens_details is not None + assert result.usage.prompt_tokens_details.cached_tokens == 6016 + def test_done_chunk_without_counts_reports_no_usage_instead_of_zeros(self): result = self._parse( { @@ -993,6 +1011,42 @@ class TestOllamaStreamingUsage: assert result.usage is None +def test_non_streaming_reports_cached_tokens_ollama_sent() -> None: + config: Final = OllamaChatConfig() + + mock_response: Final = MagicMock() + mock_response.json.return_value = { + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": "Hi"}, + "done": True, + "prompt_eval_count": 6024, + "prompt_eval_cached_count": 6016, + "eval_count": 50, + } + + model_response: Final = ModelResponse() + model_response.choices = [Choices(message=Message(content=""), index=0)] + + result: Final = config.transform_response( + model="qwen3:0.6b", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=False, + ) + + assert result.usage is not None + assert result.usage.prompt_tokens == 6024 + assert result.usage.prompt_tokens_details is not None + assert result.usage.prompt_tokens_details.cached_tokens == 6016 + + def test_get_ollama_params(): try: converted_params = get_optional_params( diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index 248e745ebff..54b58c47aaf 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -65,6 +65,38 @@ class TestOllamaConfig: assert result["usage"]["completion_tokens"] == 5 assert result["usage"]["total_tokens"] == 15 + def test_transform_response_reports_cached_tokens_ollama_sent(self): + config: Final = OllamaConfig() + raw_response: Final = MagicMock() + raw_response.json.return_value = { + "response": "Hello", + "prompt_eval_count": 6024, + "prompt_eval_cached_count": 6016, + "eval_count": 5, + } + model_response: Final = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + mock_encoding: Final = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + result: Final = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + assert result["usage"]["prompt_tokens"] == 6024 + assert result.usage.prompt_tokens_details is not None + assert result.usage.prompt_tokens_details.cached_tokens == 6016 + @patch("uuid.uuid4") def test_transform_response_json_function_call(self, mock_uuid4): # Setup mock UUID @@ -711,6 +743,24 @@ class TestOllamaTextCompletionResponseIterator: assert result["usage"]["completion_tokens"] == 5 assert result["usage"]["total_tokens"] == 15 + def test_chunk_parser_done_chunk_reports_cached_tokens(self): + iterator: Final = OllamaTextCompletionResponseIterator( + streaming_response=iter([]), sync_stream=True, json_mode=False + ) + done_chunk: Final = { + "model": "llama2", + "created_at": "2025-08-06T14:34:31.5276077Z", + "response": "", + "done": True, + "prompt_eval_count": 6024, + "prompt_eval_cached_count": 6016, + "eval_count": 5, + } + result: Final = iterator.chunk_parser(done_chunk) + assert result["usage"] is not None + assert result["usage"]["prompt_tokens"] == 6024 + assert result["usage"]["prompt_tokens_details"]["cached_tokens"] == 6016 + @pytest.mark.parametrize( ("response_chunks", "expected_text", "expected_tool_name", "expected_finish_reason"), [