scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-09 08:18:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

fix(streaming): end an n>1 stream cleanly once every streamed choice finished, and map a missing Azure v1 key to 401

Browse source 
A provider that keeps n but streams fewer choices, like OpenRouter, now ends cleanly when the connection
drops after its last streamed choice finished, as the SDK 2 image did. A choice that finishes in a chunk
led by another choice now counts too. An Azure v1 deployment with no credential answers 401
AuthenticationError again instead of a retried 500
This commit is contained in:
mateo-berri 2026-10-29 22:30:26 -07:00
parent 1c117a534e
commit 801cffe005
4 changed files with 118 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -2120,7 +2120,7 @@ def _map_azure_exception(
response=getattr(original_exception, "response", None),
body=getattr(original_exception, "body", None),
)
elif "The api_key client option must be set either by passing api_key to the client or by setting" in error_str:
elif any(marker in error_str for marker in _OPENAI_CLIENT_MISSING_CREDENTIALS_MESSAGES):
raise AuthenticationError(
message=f"{exception_provider} AuthenticationError - {message}",
llm_provider=custom_llm_provider,

60
litellm/litellm_core_utils/streaming_handler.py
View file

@ -97,39 +97,20 @@ def _next_sync_or_exhausted(it: Iterator[object]) -> object:
_CHOICE_CHUNK_TYPES: Final = (ChatCompletionChunk, ModelResponseStream, Completion)
def _first_choice_index(chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> int:
match chunk:
case ChatCompletionChunk(choices=[first, *_]):
return first.index
case ModelResponseStream(choices=[first, *_]):
return first.index
case Completion(choices=[first, *_]):
return first.index
case _:
return 0
def _choice_finish_reasons(
chunk: ChatCompletionChunk | ModelResponseStream | Completion | None,
) -> tuple[tuple[int, str | None], ...]:
match chunk:
case ChatCompletionChunk(choices=choices):
return tuple((choice.index, choice.finish_reason) for choice in choices)
return tuple((choice.index, choice.finish_reason) for choice in choices or ())
case ModelResponseStream(choices=choices):
return tuple((choice.index, choice.finish_reason) for choice in choices)
return tuple((choice.index, choice.finish_reason) for choice in choices or ())
case Completion(choices=choices):
return tuple((choice.index, choice.finish_reason) for choice in choices)
return tuple((choice.index, choice.finish_reason) for choice in choices or ())
case _:
return ()
def _indices_of_finished_choices(
chunk: ChatCompletionChunk | ModelResponseStream | Completion | None,
) -> frozenset[int]:
finished: Final = frozenset(index for index, finish_reason in _choice_finish_reasons(chunk) if finish_reason)
return finished or frozenset({_first_choice_index(chunk)})
def _stamp_usage_cost(
usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream,
complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream,
@ -382,7 +363,8 @@ class CustomStreamWrapper:
self._provider_response_model: str | None = None
self.received_finish_reason: str | None = None
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
self._finished_choice_indices: frozenset[int] | None = None
self._streamed_choice_indices: frozenset[int] = frozenset()
self._finished_choice_indices: frozenset[int] = frozenset()
self.special_tokens = [
"<|assistant|>",
"<|system|>",
@ -1543,11 +1525,11 @@ class CustomStreamWrapper:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider == "text-completion-openai":
response_obj = self.handle_openai_text_completion_chunk(chunk)
self._record_choice_progress(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
if response_obj["usage"] is not None:
_text_completion_usage: Final[Usage] = response_obj["usage"]
setattr(
@ -1583,11 +1565,11 @@ class CustomStreamWrapper:
)
elif self.custom_llm_provider == "azure_text":
response_obj = self.handle_azure_text_completion_chunk(chunk)
self._record_choice_progress(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}")
if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"]
self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
elif self.custom_llm_provider == "cached_response":
cached_chunk: Final = cast(ModelResponseStream, chunk)
cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None
@ -1627,6 +1609,7 @@ class CustomStreamWrapper:
if isinstance(chunk, BaseModel) and hasattr(chunk, "model"):
# for azure, we need to pass the model from the original chunk
self.model = getattr(chunk, "model", self.model)
self._record_choice_progress(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
response_obj = self.handle_openai_chat_completion_chunk(chunk)
if response_obj is None:
return _ProviderChunkEarlyReturn(None)
@ -1638,7 +1621,6 @@ class CustomStreamWrapper:
f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}"
)
self.received_finish_reason = response_obj["finish_reason"]
self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None)
if response_obj.get("original_chunk", None) is not None:
if hasattr(response_obj["original_chunk"], "id"):
model_response = self.set_model_id(response_obj["original_chunk"].id, model_response)
@ -2297,27 +2279,17 @@ class CustomStreamWrapper:
raise
raise StopIteration from e
def _record_finished_choices(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None:
self._finished_choice_indices = (self._finished_choice_indices or frozenset()) | _indices_of_finished_choices(
chunk
def _record_choice_progress(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None:
progress: Final = _choice_finish_reasons(chunk)
self._streamed_choice_indices = self._streamed_choice_indices.union(index for index, _ in progress)
self._finished_choice_indices = self._finished_choice_indices.union(
index for index, finish_reason in progress if finish_reason
)
def _every_choice_finished(self) -> bool:
if self.received_finish_reason is None:
return False
requested_choices: Final = self._requested_choice_count()
match self._finished_choice_indices:
case None:
return True
case finished_choices:
return requested_choices <= 1 or len(finished_choices) >= requested_choices
def _requested_choice_count(self) -> int:
match getattr(self.logging_obj, "optional_params", None):
case {"n": int() as requested}:
return requested
case _:
return 1
return (
self.received_finish_reason is not None and self._streamed_choice_indices <= self._finished_choice_indices
)
def fetch_sync_stream(self):
if self.completion_stream is None and self.make_call is not None:

71
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -5345,6 +5345,77 @@ async def test_one_chunk_finishing_both_requested_choices_then_dropped_before_do
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected if chunk.choices) == "HiHel"
def _two_choices_requested_but_only_choice_0_streamed() -> tuple[ModelResponseStream, ...]:
return (_reset_test_chunk(content="Hi", index=0), _reset_test_chunk(finish_reason="stop", index=0))
def _two_text_choices_requested_but_only_choice_0_streamed() -> tuple[Completion, ...]:
return (_text_completion_chunk(0, "Hi"), _text_completion_chunk(0, finish_reason="stop"))
def _choice_1_finishing_beside_choice_0_content() -> tuple[ModelResponseStream, ...]:
return (
_reset_test_chunk(content="Hel", index=1),
ModelResponseStream(
id="chatcmpl-reset-test",
created=1783458104,
model="stub-model",
object="chat.completion.chunk",
choices=[
StreamingChoices(index=0, delta=Delta(content="Hi"), finish_reason=None),
StreamingChoices(index=1, delta=Delta(content=None), finish_reason="stop"),
],
),
_reset_test_chunk(finish_reason="stop", index=0),
)
@pytest.mark.asyncio
@_SYNC_AND_ASYNC
@pytest.mark.parametrize(
("provider", "chunks", "content"),
(
("openai", _two_choices_requested_but_only_choice_0_streamed, "Hi"),
("text-completion-openai", _two_text_choices_requested_but_only_choice_0_streamed, "Hi"),
("openai", _choice_1_finishing_beside_choice_0_content, "HelHi"),
),
ids=("provider_streams_fewer_choices", "text_provider_streams_fewer_choices", "choice_1_finishes_mid_chunk"),
)
async def test_dropped_before_done_after_every_streamed_choice_finished_ends_cleanly(
logging_obj: Logging,
sync_mode: bool,
provider: str,
chunks: Callable[[], Sequence[ModelResponseStream | Completion]],
content: str,
):
logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
stream: Final = _chunks_then_failure(chunks(), _httpx2_read_error())
wrapper: Final = CustomStreamWrapper(
completion_stream=stream if sync_mode else _as_async(stream),
model="stub-model",
custom_llm_provider=provider,
logging_obj=logging_obj,
)
collected: Final = await _collect(wrapper, sync_mode)
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected if chunk.choices) == content
@pytest.mark.asyncio
@_SYNC_AND_ASYNC
async def test_dropped_before_done_while_choice_0_streams_beside_a_finished_choice_1_raises(
logging_obj: Logging, sync_mode: bool
):
logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
wrapper: Final = _openai_stream_over(
_choice_1_finishing_beside_choice_0_content()[:2], _httpx2_read_error(), sync_mode, logging_obj
)
with pytest.raises(MidStreamFallbackError):
await _collect(wrapper, sync_mode)
@pytest.mark.asyncio
@_SYNC_AND_ASYNC
@_SDK3_DROPPED_CONNECTIONS

30
tests/unit/llms/azure/test_azure_common_utils.py
View file

@ -2687,3 +2687,33 @@ async def test_azure_v1_deployment_keyed_only_by_the_openai_api_key_env_sends_th
assert response.choices[0].message.content == "hi"
assert route.calls.last.request.headers["authorization"] == "Bearer azure-key-in-openai-env"
@pytest.mark.parametrize("is_async", [False, True])
@pytest.mark.asyncio
async def test_azure_v1_deployment_with_no_credential_is_an_authentication_error(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, is_async: bool
) -> None:
for name in (
"AZURE_API_KEY",
"AZURE_OPENAI_API_KEY",
"AZURE_AD_TOKEN",
"AZURE_OPENAI_AD_TOKEN",
"AZURE_CLIENT_ID",
"OPENAI_API_KEY",
):
monkeypatch.delenv(name, raising=False)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
api_base: Final = f"https://nokey-{is_async}.openai.azure.com"
route: Final = respx_mock.post(f"{api_base}/openai/v1/embeddings").respond(200, json={})
request: Final = {
"model": "azure/text-embedding-3-small",
"input": ["hi"],
"api_base": api_base,
"api_version": "v1",
}
with pytest.raises(litellm.AuthenticationError):
await litellm.aembedding(**request) if is_async else litellm.embedding(**request)
assert not route.called