mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
test(test_streaming_handler.py): add unit test to ensure model response stream with usage is always used
This commit is contained in:
Krrish Dholakia
parent
82252ecac1
commit
cc82d42d25
1 changed files with 663 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue