mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(streaming): reject partial generic streaming chunks missing required fields
Fixes #43487
This commit is contained in:
parent
b94f5bdbed
commit
39961ad57a
2 changed files with 22 additions and 1 deletions
|
|
@ -55,6 +55,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call"
|
|||
_SYNC_ITER_EXHAUSTED: Final = object()
|
||||
|
||||
_GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
|
||||
_GCHUNK_REQUIRED_FIELDS: Final[frozenset] = frozenset(GChunk.__required_keys__)
|
||||
_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value})
|
||||
|
||||
|
||||
|
|
@ -2459,7 +2460,7 @@ def generic_chunk_has_all_required_fields(chunk: dict) -> bool:
|
|||
:param chunk: The dictionary to check.
|
||||
:return: True if all required fields are present, False otherwise.
|
||||
"""
|
||||
return all(key in _GCHUNK_FIELDS for key in chunk)
|
||||
return all(key in chunk for key in _GCHUNK_REQUIRED_FIELDS)
|
||||
|
||||
|
||||
def convert_generic_chunk_to_model_response_stream(
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
CustomStreamWrapper,
|
||||
_ProviderChunkEarlyReturn,
|
||||
_ProviderChunkParsed,
|
||||
convert_generic_chunk_to_model_response_stream,
|
||||
generic_chunk_has_all_required_fields,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
|
|
@ -4983,3 +4985,21 @@ async def test_async_stream_without_usage_counts_tokens_off_the_event_loop():
|
|||
assert chunks[-1].usage.prompt_tokens > 100_000
|
||||
assert chunks[-1].usage.completion_tokens > 100_000
|
||||
assert_loop_stayed_free(took, lags)
|
||||
|
||||
|
||||
def test_generic_chunk_predicate_rejects_partial_chunk():
|
||||
partial_chunk = {"is_finished": False}
|
||||
assert generic_chunk_has_all_required_fields(partial_chunk) is False
|
||||
|
||||
full_chunk = {
|
||||
"text": "hello",
|
||||
"is_finished": True,
|
||||
"finish_reason": "stop",
|
||||
"usage": None,
|
||||
"index": 0,
|
||||
}
|
||||
assert generic_chunk_has_all_required_fields(full_chunk) is True
|
||||
|
||||
stream = convert_generic_chunk_to_model_response_stream(full_chunk)
|
||||
assert stream.choices[0].delta.content == "hello"
|
||||
assert stream.finish_reason == "stop"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue