diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 59020b058a2..46e846765c9 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -397,6 +397,8 @@ class OllamaChatConfig(BaseConfig): "eval_count", litellm.token_counter(text=response_json["message"]["content"]), ) + raw_cached: Final = response_json.get("prompt_eval_cached_count") + cached_tokens: Final = raw_cached if isinstance(raw_cached, int) else 0 setattr( model_response, "usage", @@ -404,6 +406,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 @@ -429,10 +432,13 @@ 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 + raw_cached: Final = chunk.get("prompt_eval_cached_count") + cached_tokens: Final = raw_cached if isinstance(raw_cached, int) else 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 6dbcb050e11..26091ff1739 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -1,8 +1,8 @@ import json import re import time -from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import TYPE_CHECKING, Any, Final, cast # noqa: TID251 # untyped stream chunk from httpx._models import Headers, Response from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -95,6 +95,7 @@ class _OllamaGenerateMessage(TypedDict): class _OllamaGenerateResponse(TypedDict): message: ReadOnly[NotRequired[_OllamaGenerateMessage]] prompt_eval_count: ReadOnly[NotRequired[int]] + prompt_eval_cached_count: ReadOnly[NotRequired[int]] eval_count: ReadOnly[NotRequired[int]] @@ -391,6 +392,8 @@ class OllamaConfig(BaseConfig): completion_tokens: Final = response_json.get( "eval_count", len((response_json.get("message") or {}).get("content", "")) ) + raw_cached: Final = response_json.get("prompt_eval_cached_count") + cached_tokens: Final = raw_cached if isinstance(raw_cached, int) else 0 setattr( model_response, "usage", @@ -398,6 +401,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 @@ -585,10 +589,13 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): usage: ChatCompletionUsageBlock | None = None if prompt_eval_count is not None and eval_count is not None: + raw_cached: Final = cast(Mapping[str, object], chunk).get("prompt_eval_cached_count") + cached_tokens: Final = raw_cached if isinstance(raw_cached, int) else 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 98ca5aee407..b834333c3c1 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1065,6 +1065,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( { @@ -1089,6 +1107,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 26addf1b24c..515a49f4a91 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -67,6 +67,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 @@ -713,6 +745,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"), [