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.
This commit is contained in:
rxchitrx 2026-10-17 21:58:38 +05:30
parent c143954493
commit d2c1ed9af8
4 changed files with 112 additions and 0 deletions

View file

@ -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},
)

View file

@ -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)

View file

@ -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(

View file

@ -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"),
[