diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index f92961df68f..d2134b8b5f1 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -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, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8b7a4817a10..fccea6c752d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index c0383ffbd4f..6ab59b414e6 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -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 diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 260f61ab838..5419ce48d21 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -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