mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 06:22:48 +00:00
a4a3348801
* fix(vertex): stream Model Garden Gemma/Qwen responses correctly through /v1/messages * test(vertex): cover _CombinedChunkSplitter defensive branches * test(databricks): rename test file to avoid duplicate basename collision * fix(databricks,anthropic): defensive token defaults; document single-mode splitter Address greptile P2 concerns: - databricks: default usage token fields to 0 when constructing ChatCompletionUsageBlock from a partially populated usage block — matches the defensive pattern used in ollama/vertex_ai/cohere/bedrock. - _CombinedChunkSplitter: clarify in the docstring that an instance is single-mode (sync or async, not both), since the two iteration paths hold independent upstream iterator references. Co-authored-by: Claude <claude@anthropic.com> --------- Co-authored-by: Steven Kessler <9701252+stvnksslr@users.noreply.github.com> Co-authored-by: Claude <claude@anthropic.com>
65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
"""
|
|
Regression test for the databricks streaming chunk parser.
|
|
|
|
OpenAI-compatible servers (e.g. Vertex AI Model Garden vLLM endpoints) send a final
|
|
usage-only chunk with an empty `choices` list when `stream_options.include_usage` is
|
|
set. `chunk_parser` previously did `choices[0]` unconditionally, raising
|
|
`IndexError` -> `MidStreamFallbackError` and crashing the stream.
|
|
"""
|
|
|
|
from litellm.llms.databricks.streaming_utils import ModelResponseIterator
|
|
|
|
|
|
def test_chunk_parser_handles_empty_choices_usage_chunk():
|
|
"""A usage-only final chunk (empty choices) must not raise IndexError."""
|
|
iterator = ModelResponseIterator(streaming_response=None, sync_stream=True)
|
|
usage_only_chunk = {
|
|
"id": "chatcmpl-x",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "m",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 20, "completion_tokens": 8, "total_tokens": 28},
|
|
}
|
|
|
|
result = iterator.chunk_parser(chunk=usage_only_chunk)
|
|
|
|
assert result["text"] == ""
|
|
assert result["is_finished"] is False
|
|
assert result["usage"] is not None
|
|
assert result["usage"]["prompt_tokens"] == 20
|
|
assert result["usage"]["completion_tokens"] == 8
|
|
|
|
|
|
def test_chunk_parser_empty_choices_without_usage():
|
|
"""An empty-choices chunk with no usage block returns usage=None, no error."""
|
|
iterator = ModelResponseIterator(streaming_response=None, sync_stream=True)
|
|
chunk = {
|
|
"id": "chatcmpl-x",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "m",
|
|
"choices": [],
|
|
}
|
|
|
|
result = iterator.chunk_parser(chunk=chunk)
|
|
|
|
assert result["text"] == ""
|
|
assert result["usage"] is None
|
|
|
|
|
|
def test_chunk_parser_normal_content_chunk_still_works():
|
|
"""A regular content chunk is unaffected by the empty-choices guard."""
|
|
iterator = ModelResponseIterator(streaming_response=None, sync_stream=True)
|
|
chunk = {
|
|
"id": "chatcmpl-x",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "m",
|
|
"choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}],
|
|
}
|
|
|
|
result = iterator.chunk_parser(chunk=chunk)
|
|
|
|
assert result["text"] == "hi"
|