diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 8b23ae135b5..04ab800c069 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -22,9 +22,9 @@ from litellm.types.utils import ( GenericStreamingChunk, Message, ModelResponse, - Usage, ) +from ...anthropic.chat.transformation import AnthropicConfig from ...base_llm.base_model_iterator import BaseModelResponseIterator from ...openai_like.chat.transformation import OpenAIGPTConfig from ..utils import SnowflakeBaseConfig @@ -428,11 +428,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): message=message, ) - usage_data = response_json.get("usage", {}) - usage = Usage( - prompt_tokens=usage_data.get("input_tokens", 0), - completion_tokens=usage_data.get("output_tokens", 0), - total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0), + usage = AnthropicConfig().calculate_usage( + usage_object=response_json.get("usage", {}), + reasoning_content=None, + completion_response=response_json, ) model_response.choices = [choice] diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py index fb21e2e6f6b..0a5be1d1d0b 100644 --- a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py @@ -422,6 +422,41 @@ class TestAnthropicConfigResponse: assert result.usage.completion_tokens == 5 assert result.usage.total_tokens == 15 + def test_prompt_cache_tokens_mapped(self): + body = { + "id": "msg_cache", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-5", + "content": [{"type": "text", "text": "Hello!"}], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 2, + "output_tokens": 256, + "cache_creation_input_tokens": 1024, + "cache_read_input_tokens": 2048, + }, + } + result = self.cfg.transform_response( + model="snowflake/claude-opus-4-5", + raw_response=httpx.Response(200, json=body), + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + usage = result.usage + assert usage.cache_creation_input_tokens == 1024 + assert usage.cache_read_input_tokens == 2048 + assert usage.prompt_tokens == 2 + 1024 + 2048 + assert usage.total_tokens == 2 + 1024 + 2048 + 256 + assert usage.prompt_tokens_details.cached_tokens == 2048 + assert usage.prompt_tokens_details.cache_creation_tokens == 1024 + assert usage.prompt_tokens_details.text_tokens == 2 + def test_stop_reason_end_turn_maps_to_stop(self): raw = _make_anthropic_response() result = self.cfg.transform_response(