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

Merge pull request #9314 from BerriAI/litellm_dev_03_16_2025_p1

Browse source 
fix(utils.py): Prevents final chunk w/ usage from being ignored
This commit is contained in:
Krish Dholakia 2025-03-17 10:03:03 -07:00 • committed by GitHub
commit 41d1a88d05
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 309 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/litellm_core_utils/litellm_logging.py
View file

@ -110,7 +110,6 @@ from .exception_mapping_utils import _get_response_headers
from .initialize_dynamic_callback_params import (
initialize_standard_callback_dynamic_params as _initialize_standard_callback_dynamic_params,
)
from .logging_utils import _assemble_complete_response_from_streaming_chunks
from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache
try:
@ -2351,18 +2350,6 @@ class Logging(LiteLLMLoggingBaseClass):
return result
elif isinstance(result, ResponseCompletedEvent):
return result.response
elif isinstance(result, ModelResponseStream):
complete_streaming_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = _assemble_complete_response_from_streaming_chunks(
result=result,
start_time=start_time,
end_time=end_time,
request_kwargs=self.model_call_details,
streaming_chunks=streaming_chunks,
is_async=is_async,
)
return complete_streaming_response
return None
def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse:

1
litellm/litellm_core_utils/logging_utils.py
View file

@ -77,6 +77,7 @@ def _assemble_complete_response_from_streaming_chunks(
complete_streaming_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = None
if result.choices[0].finish_reason is not None: # if it's the last chunk
streaming_chunks.append(result)
try:

7
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -13,6 +13,7 @@ from litellm.types.utils import (
Function,
FunctionCall,
ModelResponse,
ModelResponseStream,
PromptTokensDetails,
Usage,
)
@ -319,8 +320,12 @@ class ChunkProcessor:
usage_chunk: Optional[Usage] = None
if "usage" in chunk:
usage_chunk = chunk["usage"]
elif isinstance(chunk, ModelResponse) and hasattr(chunk, "_hidden_params"):
elif (
isinstance(chunk, ModelResponse)
or isinstance(chunk, ModelResponseStream)
) and hasattr(chunk, "_hidden_params"):
usage_chunk = chunk._hidden_params.get("usage", None)
if usage_chunk is not None:
usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk)
if (

42
litellm/litellm_core_utils/streaming_handler.py
View file

@ -898,6 +898,8 @@ class CustomStreamWrapper:
return model_response
# Default - return StopIteration
if hasattr(model_response, "usage"):
self.chunks.append(model_response)
raise StopIteration
# flush any remaining holding chunk
if len(self.holding_chunk) > 0:
@ -1553,10 +1555,11 @@ class CustomStreamWrapper:
if response is None:
continue
## LOGGING
threading.Thread(
target=self.run_success_logging_and_cache_storage,
args=(response, cache_hit),
).start() # log response
executor.submit(
self.run_success_logging_and_cache_storage,
response,
cache_hit,
) # log response
choice = response.choices[0]
if isinstance(choice, StreamingChoices):
self.response_uptil_now += choice.delta.get("content", "") or ""
@ -1600,13 +1603,21 @@ class CustomStreamWrapper:
"usage",
getattr(complete_streaming_response, "usage"),
)
## LOGGING
threading.Thread(
target=self.logging_obj.success_handler,
args=(response, None, None, cache_hit),
).start() # log response
executor.submit(
self.logging_obj.success_handler,
complete_streaming_response,
None,
None,
cache_hit,
)
else:
executor.submit(
self.logging_obj.success_handler,
response,
None,
None,
cache_hit,
)
if self.sent_stream_usage is False and self.send_stream_usage is True:
self.sent_stream_usage = True
return response
@ -1618,10 +1629,11 @@ class CustomStreamWrapper:
usage = calculate_total_usage(chunks=self.chunks)
processed_chunk._hidden_params["usage"] = usage
## LOGGING
threading.Thread(
target=self.run_success_logging_and_cache_storage,
args=(processed_chunk, cache_hit),
).start() # log response
executor.submit(
self.run_success_logging_and_cache_storage,
processed_chunk,
cache_hit,
) # log response
return processed_chunk
except Exception as e:
traceback_exception = traceback.format_exc()

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

@ -1274,6 +1274,13 @@ class AWSEventStreamDecoder:
def converse_chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
try:
verbose_logger.debug("\n\nRaw Chunk: {}\n\n".format(chunk_data))
chunk_data["usage"] = {
"inputTokens": 3,
"outputTokens": 392,
"totalTokens": 2191,
"cacheReadInputTokens": 1796,
"cacheWriteInputTokens": 0,
}
text = ""
tool_use: Optional[ChatCompletionToolCallChunk] = None
finish_reason = ""
@ -1354,6 +1361,7 @@ class AWSEventStreamDecoder:
finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
elif "usage" in chunk_data:
usage = converse_config._transform_usage(chunk_data.get("usage", {}))
model_response_provider_specific_fields = {}
if "trace" in chunk_data:
trace = chunk_data.get("trace")

266
tests/litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -1,16 +1,28 @@
import json
import os
import sys
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock, Mock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import asyncio
import traceback
from typing import Optional
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import ModelResponseStream
from litellm.types.utils import (
Delta,
ModelResponseStream,
PromptTokensDetailsWrapper,
StreamingChoices,
Usage,
)
from litellm.utils import ModelResponseListIterator
@pytest.fixture
@ -24,6 +36,82 @@ def initialized_custom_stream_wrapper() -> CustomStreamWrapper:
return streaming_handler
bedrock_chunks = [
ModelResponseStream(
id="chatcmpl-d249def8-a78b-464c-87b5-3a6f43565292",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content="I'm Claude",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
ModelResponseStream(
id="chatcmpl-fe559823-b383-4249-ab87-52f6ad9d08c2",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content=", an AI",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
ModelResponseStream(
id="chatcmpl-c1c6cc2f-75b9-4a24-88b9-4e5aacd0268b",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(
provider_specific_fields=None,
content="",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
]
def test_is_chunk_non_empty(initialized_custom_stream_wrapper: CustomStreamWrapper):
"""Unit test if non-empty when reasoning_content is present"""
chunk = {
@ -277,3 +365,177 @@ def test_strip_sse_data_from_chunk():
# Test with None input
assert CustomStreamWrapper._strip_sse_data_from_chunk(None) is None
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_streaming_handler_with_usage(
sync_mode: bool, final_usage_block: Optional[Usage] = None
):
import time
final_usage_block = final_usage_block or Usage(
completion_tokens=392,
prompt_tokens=1799,
total_tokens=2191,
completion_tokens_details=None,
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None, cached_tokens=1796, text_tokens=None, image_tokens=None
),
)
final_chunk = ModelResponseStream(
id="chatcmpl-87291500-d8c5-428e-b187-36fe5a4c97ab",
created=1742056047,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content="",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=final_usage_block,
)
test_chunks = bedrock_chunks + [final_chunk]
completion_stream = ModelResponseListIterator(model_responses=test_chunks)
response = CustomStreamWrapper(
completion_stream=completion_stream,
model="bedrock/claude-3-5-sonnet-20240620-v1:0",
custom_llm_provider="bedrock",
logging_obj=Logging(
model="bedrock/claude-3-5-sonnet-20240620-v1:0",
messages=[{"role": "user", "content": "Hey"}],
stream=True,
call_type="completion",
start_time=time.time(),
litellm_call_id="12345",
function_id="1245",
),
stream_options={"include_usage": True},
)
chunk_has_usage = False
if sync_mode:
for chunk in response:
if hasattr(chunk, "usage"):
assert chunk.usage == final_usage_block
chunk_has_usage = True
else:
async for chunk in response:
if hasattr(chunk, "usage"):
assert chunk.usage == final_usage_block
chunk_has_usage = True
assert chunk_has_usage
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_streaming_with_usage_and_logging(sync_mode: bool):
import time
from litellm.integrations.custom_logger import CustomLogger
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
litellm.success_callback = [mock_callback]
litellm._async_success_callback = [mock_callback]
final_usage_block = Usage(
completion_tokens=392,
prompt_tokens=1799,
total_tokens=2191,
completion_tokens_details=None,
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None,
cached_tokens=1796,
text_tokens=None,
image_tokens=None,
),
cache_creation_input_tokens=0,
cache_read_input_tokens=1796,
)
with patch.object(
mock_callback, "log_success_event"
) as mock_log_success_event, patch.object(
mock_callback, "log_stream_event"
) as mock_log_stream_event, patch.object(
mock_callback, "async_log_success_event"
) as mock_async_log_success_event, patch.object(
mock_callback, "async_log_stream_event"
) as mock_async_log_stream_event:
await test_streaming_handler_with_usage(
sync_mode=sync_mode, final_usage_block=final_usage_block
)
if sync_mode:
time.sleep(1)
mock_log_success_event.assert_called_once()
# mock_log_stream_event.assert_called()
else:
await asyncio.sleep(1)
mock_async_log_success_event.assert_called_once()
# mock_async_log_stream_event.assert_called()
print(mock_log_success_event.call_args.kwargs.keys())
mock_log_success_event.call_args.kwargs[
"response_obj"
].usage == final_usage_block
def test_streaming_handler_with_stop_chunk(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
args = {
"completion_obj": {"content": ""},
"response_obj": {
"text": "",
"is_finished": True,
"finish_reason": "length",
"logprobs": None,
"original_chunk": ModelResponseStream(
id="chatcmpl-ad517c2e-c197-48de-a2e6-a559cca48124",
created=1742093326,
model=None,
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason="length",
index=0,
delta=Delta(
provider_specific_fields=None,
content="",
role="assistant",
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields={},
usage=None,
),
"usage": None,
},
}
returned_chunk = initialized_custom_stream_wrapper.return_processed_chunk_logic(
**args, model_response=ModelResponseStream()
)
assert returned_chunk is None

4
tests/local_testing/test_custom_callback_input.py
View file

@ -1339,7 +1339,7 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream):
continue
time.sleep(2)
mock_client.assert_called_once()
mock_client.assert_called()
print(
f"mock_client_post.call_args: {mock_client.call_args.kwargs['kwargs'].keys()}"
@ -1559,7 +1559,7 @@ def test_logging_standard_payload_llm_headers(stream):
continue
time.sleep(2)
mock_client.assert_called_once()
mock_client.assert_called()
standard_logging_object: StandardLoggingPayload = mock_client.call_args.kwargs[
"kwargs"

2
tests/logging_callback_tests/test_assemble_streaming_responses.py
View file

@ -26,7 +26,7 @@ from respx import MockRouter
import litellm
from litellm import Choices, Message, ModelResponse, TextCompletionResponse, TextChoices
from litellm.litellm_core_utils.litellm_logging import (
from litellm.litellm_core_utils.logging_utils import (
_assemble_complete_response_from_streaming_chunks,
)