diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a369b7f3e36..0945c45491d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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: diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 6782435af62..3c934a42761 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -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: diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index e78b10c2892..7a5ee3e41e3 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -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 ( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5d5a8bf2563..15d94b31a99 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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() diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 9fa791e0699..84ac592c411 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -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") diff --git a/tests/litellm/litellm_core_utils/test_streaming_handler.py b/tests/litellm/litellm_core_utils/test_streaming_handler.py index 10fe1db4abd..31d541330c6 100644 --- a/tests/litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/litellm/litellm_core_utils/test_streaming_handler.py @@ -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 diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index d18668ebf1c..b0ebcf77673 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -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" diff --git a/tests/logging_callback_tests/test_assemble_streaming_responses.py b/tests/logging_callback_tests/test_assemble_streaming_responses.py index 7b28f69917e..1101350fa29 100644 --- a/tests/logging_callback_tests/test_assemble_streaming_responses.py +++ b/tests/logging_callback_tests/test_assemble_streaming_responses.py @@ -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, )