diff --git a/tests/litellm/litellm_core_utils/test_streaming_handler.py b/tests/litellm/litellm_core_utils/test_streaming_handler.py index e957d25f4b0..76de59b57b8 100644 --- a/tests/litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/litellm/litellm_core_utils/test_streaming_handler.py @@ -9,6 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ( Delta, @@ -17,6 +18,7 @@ from litellm.types.utils import ( StreamingChoices, Usage, ) +from litellm.utils import ModelResponseListIterator @pytest.fixture @@ -370,3 +372,664 @@ def test_chunk_with_usage(initialized_custom_stream_wrapper: CustomStreamWrapper }, } 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 + ), + ) + 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", + ), + ) + + for chunk in response: + if hasattr(chunk, "usage"): + assert chunk.usage == final_usage_block