mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(streaming_handler.py): pass complete streaming response on completion
This commit is contained in:
Krrish Dholakia
parent
08b297230e
commit
c0a76427d2
2 changed files with 240 additions and 727 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue