scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(snowflake): preserve prompt cache token usage on Claude responses

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-30 11:06:45 +00:00
parent 71b825a7f0
commit 8b52e90180
2 changed files with 40 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/snowflake/chat/transformation.py
View file

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

35
tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py
View file

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