fix: end text completion n=2 streams cleanly, keep Azure image gateway query params, tighten SDK 3 test fixtures

.circleci/scripts/run_integration.sh

@@ -70,6 +70,12 @@ trap 'exit 143' TERM
 
 export PATH="$PWD/.venv/bin:$PATH"
 export PYTHONPATH="$PWD:$PWD/tests:$PWD/tests/e2e"
+if [ "$suite" = extensions ]; then
+  retained_openai="$(mktemp -d)/retained-openai"
+  uv venv --quiet --python 3.12 "$retained_openai"
+  uv pip install --quiet --python "$retained_openai/bin/python" --exclude-newer 2026-04-29T00:00:00Z openai==2.33.0
+  export RETAINED_OPENAI_PYTHON="$retained_openai/bin/python"
+fi
 export DATABASE_URL="postgresql://postgres:postgres@127.0.0.1:5432/circle_test"
 export REDIS_HOST=127.0.0.1 REDIS_PORT=6379
 export LITELLM_MASTER_KEY="sk-$(openssl rand -hex 16)" LITELLM_SALT_KEY=sk-integration-salt

litellm/litellm_core_utils/streaming_handler.py

@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast
 
 import anyio
 import httpx
+from openai.types import Completion
 from openai.types.chat import ChatCompletionChunk
 from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 from openai.types.chat.chat_completion_chunk import ChoiceDelta
@@ -93,12 +94,17 @@ def _next_sync_or_exhausted(it: Iterator[object]) -> object:
         return _SYNC_ITER_EXHAUSTED
 
 
-def _first_choice_index(chunk: ChatCompletionChunk | ModelResponseStream) -> int:
+_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
 
@@ -343,7 +349,7 @@ 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] = frozenset()
+        self._finished_choice_indices: frozenset[int] | None = None
         self.special_tokens = [
             "<|assistant|>",
             "<|system|>",
@@ -1503,6 +1509,7 @@ class CustomStreamWrapper:
             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_choice(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(
@@ -1542,6 +1549,7 @@ class CustomStreamWrapper:
             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_choice(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
@@ -1586,9 +1594,7 @@ 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._finished_choice_indices = self._finished_choice_indices | {
-                    _first_choice_index(chunk) if isinstance(chunk, (ChatCompletionChunk, ModelResponseStream)) else 0
-                }
+                self._record_finished_choice(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)
@@ -2225,11 +2231,18 @@ class CustomStreamWrapper:
                 raise
             raise StopIteration from e
 
+    def _record_finished_choice(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None:
+        self._finished_choice_indices = (self._finished_choice_indices or frozenset()) | {_first_choice_index(chunk)}
+
     def _every_choice_finished(self) -> bool:
         if self.received_finish_reason is None:
             return False
         requested_choices: Final = self._requested_choice_count()
-        return requested_choices <= 1 or len(self._finished_choice_indices) >= requested_choices
+        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):

litellm/llms/azure/azure.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 import time
-from collections.abc import Callable, Coroutine
+from collections.abc import Callable, Coroutine, Mapping
 from typing import Final
 
 import httpx
@@ -13,7 +13,7 @@ from openai import (
     BadRequestError,
     OpenAI,
 )
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, ConfigDict, TypeAdapter
 
 import litellm
 from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES
@@ -140,6 +140,16 @@ def check_dynamic_azure_params(
 _check_dynamic_azure_params = check_dynamic_azure_params
 
 
+class _AzureDefaultQuery(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    default_query: Mapping[str, str] | None = None
+
+
+def _with_default_query(url: str, default_query: Mapping[str, str] | None) -> str:
+    return str(httpx.URL(url).copy_merge_params(default_query)) if default_query else url
+
+
 class AzureChatCompletion(BaseAzureLLM, BaseLLM):
     def __init__(self) -> None:
         super().__init__()
@@ -1110,6 +1120,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
             "/"
         )
         api_version: Final[str] = azure_client_params.get("api_version", "")
+        default_query: Final = _AzureDefaultQuery.model_validate(azure_client_params).default_query
         if model is None:
             model = ""
 
@@ -1134,17 +1145,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
             route="/openai/images/generations",
         )
         if v1_url is not None:
-            return v1_url
+            return _with_default_query(v1_url, default_query)
 
-        if "/openai/deployments/" in api_base:
-            base_url_with_deployment = api_base
-        else:
-            base_url_with_deployment = api_base + "/openai/deployments/" + model
-
-        base_url_with_deployment += "/images/generations"
-        base_url_with_deployment += "?api-version=" + api_version
-
-        return base_url_with_deployment
+        deployment_url: Final = (
+            api_base if "/openai/deployments/" in api_base else f"{api_base}/openai/deployments/{model}"
+        )
+        return _with_default_query(f"{deployment_url}/images/generations?api-version={api_version}", default_query)
 
     async def aimage_generation(
         self,

pyproject.toml

@@ -369,7 +369,7 @@ constraint-dependencies = [
 override-dependencies = [
     # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0.
     "packaging>=24.0",
-    # semantic-router (through 1.1.0) caps openai at <3.0.0 but only imports it for encoders LiteLLM never builds.
+    # semantic-router (through 1.1.0) caps openai at <3.0.0, yet its OpenAI encoder for auto router files runs on SDK 3.
     "openai>=2.20.0,<4.0.0",
     "cryptography>=50.0.0,<51.0",
     "mcp>=2.3.0,<3",

tests/integration/compatibility/_retained_openai_consumer.py

@@ -0,0 +1,56 @@
+import asyncio
+import json
+import sys
+from importlib.metadata import version
+from typing import Final
+
+import httpx
+from openai import AsyncOpenAI, OpenAI
+from openai.types.chat import ChatCompletion
+
+
+def _summary(response: ChatCompletion) -> dict[str, object]:
+    choice: Final = response.choices[0]
+    calls: Final = choice.message.tool_calls or ()
+    return {
+        "id": response.id,
+        "object": response.object,
+        "usage": [response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens]
+        if response.usage
+        else None,
+        "finish_reason": choice.finish_reason,
+        "content": choice.message.content,
+        "tool_calls": [[call.id, call.function.name, json.loads(call.function.arguments)] for call in calls],
+    }
+
+
+async def _asynchronous(base_url: str, key: str, parameters: dict[str, object]) -> tuple[ChatCompletion, ...]:
+    plain: Final = {name: value for name, value in parameters.items() if name != "tools"}
+    async with AsyncOpenAI(
+        api_key=key, base_url=base_url, max_retries=0, http_client=httpx.AsyncClient(timeout=10, trust_env=False)
+    ) as client:
+        return (await client.chat.completions.create(**parameters), await client.chat.completions.create(**plain))
+
+
+def main(base_url: str, key: str, parameters_json: str) -> None:
+    parameters: Final = json.loads(parameters_json)
+    plain: Final = {name: value for name, value in parameters.items() if name != "tools"}
+    with OpenAI(
+        api_key=key, base_url=base_url, max_retries=0, http_client=httpx.Client(timeout=10, trust_env=False)
+    ) as client:
+        sync_tools: Final[ChatCompletion] = client.chat.completions.create(**parameters)
+        sync_text: Final[ChatCompletion] = client.chat.completions.create(**plain)
+    async_tools, async_text = asyncio.run(_asynchronous(base_url, key, parameters))
+    print(
+        json.dumps(
+            {
+                "openai": version("openai"),
+                "tools": [_summary(sync_tools), _summary(async_tools)],
+                "text": [_summary(sync_text), _summary(async_text)],
+            }
+        )
+    )
+
+
+if __name__ == "__main__":
+    main(*sys.argv[1:])

tests/integration/compatibility/test_openai_consumer.py

@@ -1,7 +1,9 @@
 import json
+import os
+import subprocess
 import uuid
 from collections.abc import Callable
-from importlib.metadata import version
+from pathlib import Path
 from typing import Final
 
 import httpx
@@ -11,12 +13,22 @@ from integration._support.wire import Reply, Request, wire_server
 from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 
 
-@pytest.mark.covers("other.compatibility.openai.retained_client_parses_tools_and_usage")
-async def test_retained_openai_clients_parse_real_proxy_tool_and_usage_responses(gateway: Gateway) -> None:
-    assert version("openai") == "2.33.0", (
-        "Retain this consumer version independently before upgrading the candidate lock"
-    )
+_RETAINED_OPENAI_VERSION: Final = "2.33.0"
+_RETAINED_CONSUMER: Final = Path(__file__).with_name("_retained_openai_consumer.py")
 
+
+def _retained_openai_python() -> str:
+    python: Final = os.environ.get("RETAINED_OPENAI_PYTHON")
+    if not python:
+        pytest.fail(
+            f"Set RETAINED_OPENAI_PYTHON to a Python with openai=={_RETAINED_OPENAI_VERSION} installed, "
+            "as .circleci/scripts/run_integration.sh does for the extensions suite"
+        )
+    return python
+
+
+@pytest.mark.covers("other.compatibility.openai.retained_client_parses_tools_and_usage")
+def test_retained_openai_clients_parse_real_proxy_tool_and_usage_responses(gateway: Gateway) -> None:
     def provider(request: Request) -> Reply:
         assert request.method == "POST" and request.target == "/v1/chat/completions"
         body: Final = json.loads(request.body)
@@ -72,40 +84,35 @@ async def test_retained_openai_clients_parse_real_proxy_tool_and_usage_responses
             ],
             "extra_body": {"cache": {"no-cache": True}},
         }
-        plain: Final = {name: value for name, value in parameters.items() if name != "tools"}
-        with OpenAI(
-            api_key=key,
-            base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
-            max_retries=0,
-            http_client=httpx.Client(timeout=10, trust_env=False),
-        ) as sync:
-            first: Final = sync.chat.completions.create(**parameters)
-            first_text: Final = sync.chat.completions.create(**plain)
-        async with AsyncOpenAI(
-            api_key=key,
-            base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
-            max_retries=0,
-            http_client=httpx.AsyncClient(timeout=10, trust_env=False),
-        ) as asynchronous:
-            second: Final = await asynchronous.chat.completions.create(**parameters)
-            second_text: Final = await asynchronous.chat.completions.create(**plain)
-        assert len({response.id for response in (first, second, first_text, second_text)}) == 4
-        for response in (first, second, first_text, second_text):
-            assert response.object == "chat.completion"
-            assert (
-                response.usage.prompt_tokens == 11
-                and response.usage.completion_tokens == 4
-                and response.usage.total_tokens == 15
-            )
-        for response in (first, second):
-            assert response.choices[0].finish_reason == "tool_calls"
-            call: Final = response.choices[0].message.tool_calls[0]
-            assert call.id == "synthetic-call" and call.function.name == "add"
-            assert json.loads(call.function.arguments) == {"a": 3, "b": 5}
-        for response in (first_text, second_text):
-            assert response.choices[0].finish_reason == "stop"
-            assert response.choices[0].message.content == "Synthetic answer: 8"
-            assert not response.choices[0].message.tool_calls
+        consumer: Final = subprocess.run(
+            (
+                _retained_openai_python(),
+                "-I",
+                str(_RETAINED_CONSUMER),
+                str(gateway.client.base_url).rstrip("/") + "/v1",
+                key,
+                json.dumps(parameters),
+            ),
+            capture_output=True,
+            text=True,
+            timeout=60,
+            check=False,
+        )
+        assert consumer.returncode == 0, consumer.stderr
+        observed: Final = json.loads(consumer.stdout)
+        assert observed["openai"] == _RETAINED_OPENAI_VERSION
+        responses: Final = (*observed["tools"], *observed["text"])
+        assert len({response["id"] for response in responses}) == 4
+        for response in responses:
+            assert response["object"] == "chat.completion"
+            assert response["usage"] == [11, 4, 15]
+        for response in observed["tools"]:
+            assert response["finish_reason"] == "tool_calls"
+            assert response["tool_calls"] == [["synthetic-call", "add", {"a": 3, "b": 5}]]
+        for response in observed["text"]:
+            assert response["finish_reason"] == "stop"
+            assert response["content"] == "Synthetic answer: 8"
+            assert response["tool_calls"] == []
         assert len(wire.drain()) == 4
 
 

tests/unit/conftest.py

@@ -6,6 +6,7 @@ import sys
 from collections.abc import Callable, Coroutine, Iterator
 from dataclasses import dataclass, field
 from functools import partial
+from importlib.metadata import version
 from pathlib import Path
 from types import MappingProxyType
 from typing import Final, Literal, TypeAlias
@@ -435,11 +436,12 @@ def fake_provider_credentials(monkeypatch: pytest.MonkeyPatch) -> None:
 
 @pytest.fixture
 def openai_token_exchange_via_httpx(monkeypatch: pytest.MonkeyPatch) -> None:
-    workload_module: Final = sys.modules.get("openai.auth._workload")
-    native_client_factory: Final = getattr(workload_module, "DefaultHttpx2Client", None)
-    if native_client_factory is None:
+    if int(version("openai").split(".")[0]) < 3:
         return
     import httpx2
+    import openai.auth._workload as workload_module
+
+    native_client_factory: Final = workload_module.DefaultHttpx2Client
 
     def send_through_httpx(request: httpx2.Request) -> httpx2.Response:
         with httpx.Client(trust_env=False) as client:

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -14,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
 import httpx
 import openai
 import pytest
+from openai.types import Completion, CompletionChoice
 from openai.types.chat import ChatCompletionChunk
 from pydantic import BaseModel, TypeAdapter
 import litellm
@@ -4983,6 +4984,49 @@ async def test_sdk3_dropped_connection_after_every_requested_choice_finished_end
     assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "HiHel"
 
 
+def _text_completion_chunk(index: int, text: str = "", finish_reason: Optional[str] = None) -> Completion:
+    return Completion.model_construct(
+        id="cmpl-two-choices",
+        object="text_completion",
+        created=1,
+        model="gpt-3.5-turbo-instruct",
+        choices=[CompletionChoice.model_construct(index=index, text=text, finish_reason=finish_reason, logprobs=None)],
+    )
+
+
+@pytest.mark.asyncio
+@_SYNC_AND_ASYNC
+@pytest.mark.parametrize("provider", ("text-completion-openai", "azure_text"))
+@pytest.mark.parametrize(
+    ("finished_choices", "ends_cleanly"), (((0, 1), True), ((1,), False)), ids=("both_finished", "choice_0_open")
+)
+async def test_two_choice_text_completion_stream_dropped_before_done_ends_cleanly_only_once_both_finish(
+    logging_obj: Logging, sync_mode: bool, provider: str, finished_choices: Sequence[int], ends_cleanly: bool
+):
+    logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
+    stream: Final = _chunks_then_failure(
+        (
+            _text_completion_chunk(0, "Hi"),
+            _text_completion_chunk(1, "Hel"),
+            *(_text_completion_chunk(index, finish_reason="stop") for index in finished_choices),
+        ),
+        _httpx2_read_error(),
+    )
+    wrapper: Final = CustomStreamWrapper(
+        completion_stream=stream if sync_mode else _as_async(stream),
+        model="gpt-3.5-turbo-instruct",
+        custom_llm_provider=provider,
+        logging_obj=logging_obj,
+    )
+
+    if not ends_cleanly:
+        with pytest.raises(MidStreamFallbackError):
+            await _collect(wrapper, sync_mode)
+        return
+    chunks: Final = await _collect(wrapper, sync_mode)
+    assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "HiHel"
+
+
 @pytest.mark.asyncio
 @_SYNC_AND_ASYNC
 @_SDK3_DROPPED_CONNECTIONS

tests/unit/llms/azure/test_azure.py

@@ -271,6 +271,27 @@ def test_process_azure_endpoint_url(api_base, model, expected_endpoint):
     assert result == expected_endpoint, "Unexpected endpoint"
 
 
+@respx.mock
+@pytest.mark.parametrize("api_version", ("2025-04-01-preview", "v1"))
+def test_azure_image_generation_keeps_deployment_url_query_params(api_version: str):
+    route: Final = respx.post(url__regex=r"https://azure-gateway\.example\.com/.*").mock(
+        return_value=httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGk="}]})
+    )
+
+    litellm.image_generation(
+        model="azure/gpt-image-1",
+        prompt="a red circle",
+        api_base="https://azure-gateway.example.com/openai/deployments/gpt-image-1/images/generations?subscription-key=gateway-key",
+        api_version=api_version,
+        api_key="test-azure-key",
+    )
+
+    request_url: Final = route.calls.last.request.url
+    assert request_url.path.endswith("/images/generations")
+    assert request_url.params["subscription-key"] == "gateway-key"
+    assert request_url.params["api-version"] == api_version
+
+
 @patch("azure.identity.UsernamePasswordCredential")
 @patch("azure.identity.get_bearer_token_provider")
 def test_get_azure_ad_token_from_username_password(mock_get_bearer_token_provider, mock_credential):