mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
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:
parent
c143954493
commit
d2c1ed9af8
4 changed files with 112 additions and 0 deletions
|
|
@ -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},
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue