diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index efd004c81ae..27154b6e79c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1596,12 +1596,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 diff --git a/tests/litellm/litellm_core_utils/test_streaming_handler.py b/tests/litellm/litellm_core_utils/test_streaming_handler.py index 19948b25dc7..31d541330c6 100644 --- a/tests/litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/litellm/litellm_core_utils/test_streaming_handler.py @@ -1,14 +1,18 @@ 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 ( @@ -32,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 = { @@ -287,60 +367,155 @@ def test_strip_sse_data_from_chunk(): assert CustomStreamWrapper._strip_sse_data_from_chunk(None) is None -def test_chunk_with_usage(initialized_custom_stream_wrapper: CustomStreamWrapper): - """Test that a chunk with usage is properly handled""" +@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": ""}, - "model_response": ModelResponseStream( - id="chatcmpl-e6abdd00-9d27-4be5-9fce-9b68fa97ac01", - created=1742054811, - 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=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, - ), - ), "response_obj": { - "finish_reason": None, - "is_finished": False, + "text": "", + "is_finished": True, + "finish_reason": "length", "logprobs": None, "original_chunk": ModelResponseStream( - id="chatcmpl-e6abdd00-9d27-4be5-9fce-9b68fa97ac01", - created=1742054811, + id="chatcmpl-ad517c2e-c197-48de-a2e6-a559cca48124", + created=1742093326, model=None, object="chat.completion.chunk", system_fingerprint=None, choices=[ StreamingChoices( - finish_reason=None, + finish_reason="length", index=0, delta=Delta( provider_specific_fields=None, @@ -354,684 +529,13 @@ def test_chunk_with_usage(initialized_custom_stream_wrapper: CustomStreamWrapper ) ], provider_specific_fields={}, - usage=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, - ), + usage=None, ), + "usage": None, }, } - assert initialized_custom_stream_wrapper.is_chunk_non_empty(**args) - -def test_streaming_handler_with_usage(): - import time - - 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 - ), + returned_chunk = initialized_custom_stream_wrapper.return_processed_chunk_logic( + **args, model_response=ModelResponseStream() ) - 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-b317b568-e47b-4060-9450-41048008746e", - 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=" assistant made", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-7a209692-6f74-4e5b-b26a-71a815522441", - 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=" by Anthropic", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-1b3618dc-cebf-4220-bc91-d18b6709f882", - 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", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-d3ef70a8-ea08-4069-b8fe-3a9291bf0657", - 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=" don", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-23c20796-804d-48d3-baec-0852e3289a1c", - 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="'t have", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-0dec0bd5-38b7-4fd7-81d6-94b5093fd278", - 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=" a personal", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-7a86d50c-e2ac-4579-ac6c-2b6ae8728792", - 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=" identity like", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-d735bfee-1852-4a0f-a184-892abdc83707", - 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=" humans do, but", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-d30f3ecb-6ad3-4790-a74b-3348399f48bf", - 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 here", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-15de98eb-3655-4f3a-bbbc-447850bf0910", - 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=" to assist", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-4133758f-3304-4fcd-bdcf-0f1ea9025037", - 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=" you with", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-f36186e1-1a58-4cfc-aa1d-4d3b0a60bb37", - 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=" information", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-11754706-9289-4e40-9d79-bbfd0aad0403", - 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=None, - ), - ModelResponseStream( - id="chatcmpl-65c96965-e371-4f53-81c5-70bb55ea029d", - 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=" answer", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-0ea89799-0089-4dbb-a2bd-797cea60654a", - 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=" questions, or", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-7391b475-8010-47f8-8512-b6bde392633d", - 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=" help with various", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-5796d350-849a-44bc-973f-259ab8136873", - 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=" tasks through", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-fd61a450-fc38-48f1-9594-62968d9ee32b", - 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=" conversation", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-aa863a29-1246-45d3-8857-21608582793c", - 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=". How", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-f55ecbc5-f7ef-43e9-8be2-e56fa660676c", - 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=" can I help you", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-d1eab339-9dd3-4412-8609-2a625114c6c7", - 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=" today?", - role="assistant", - function_call=None, - tool_calls=None, - audio=None, - ), - logprobs=None, - ) - ], - provider_specific_fields={}, - usage=None, - ), - ModelResponseStream( - id="chatcmpl-2235d1e6-950e-4653-9549-963d71880d9b", - 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=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, - ), - 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, - ), - ] - - completion_stream = ModelResponseListIterator(model_responses=chunks) - - response = CustomStreamWrapper( - completion_stream=completion_stream, - model="bedrock/claude-3-5-sonnet-20240620-v1:0", - custom_llm_provider="cached_response", - 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", - ), - ) - - with patch("litellm.main.token_counter") as mock_token_counter: - for chunk in response: - if hasattr(chunk, "usage"): - assert chunk.usage == final_usage_block - assert mock_token_counter.assert_not_called() + assert returned_chunk is None