scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge 6d43ff63d6 into 78ff5ac9cd

Browse source
This commit is contained in:
Andrew 2026-09-03 09:42:33 +08:00 • committed by GitHub
commit 62e606dddd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 57 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -829,6 +829,16 @@ class MockResponseIterator: # for returning ai21 streaming responses
prompt_tokens=chunk_usage.prompt_tokens,
completion_tokens=chunk_usage.completion_tokens,
total_tokens=chunk_usage.total_tokens,
prompt_tokens_details=(
chunk_usage.prompt_tokens_details.model_dump(exclude_none=True)
if chunk_usage.prompt_tokens_details is not None
else None
),
completion_tokens_details=(
chunk_usage.completion_tokens_details.model_dump(exclude_none=True)
if chunk_usage.completion_tokens_details is not None
else None
),
),
index=0,
)

48
tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
View file

@ -9,10 +9,57 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.bedrock.chat.invoke_handler import (
AWSEventStreamDecoder,
MockResponseIterator,
make_call,
make_sync_call,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.utils import (
Choices,
CompletionTokensDetailsWrapper,
Message,
ModelResponse,
PromptTokensDetailsWrapper,
Usage,
)
def test_mock_response_iterator_preserves_usage_details():
response = ModelResponse(
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(content="done", role="assistant"),
)
],
usage=Usage(
prompt_tokens=500,
completion_tokens=25,
total_tokens=525,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=321,
cache_write_tokens=123,
),
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=11,
),
),
)
chunk = MockResponseIterator(model_response=response)._chunk_parser(response)
assert chunk["usage"] == {
"prompt_tokens": 500,
"completion_tokens": 25,
"total_tokens": 525,
"prompt_tokens_details": {
"cached_tokens": 321,
"cache_write_tokens": 123,
"cache_creation_tokens": 123,
},
"completion_tokens_details": {"reasoning_tokens": 11},
}
def test_transform_thinking_blocks_with_redacted_content():
@ -495,4 +542,3 @@ async def test_async_invoke_streaming_forwards_bedrock_response_headers():
)
assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987"