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

fix(openai): end a dropped stream only after every requested choice finished, keep Azure deployment URL query params, and narrow scope

Browse source 
A dropped connection now ends an n>1 stream cleanly only once each requested choice has sent its finish_reason. A pasted Azure deployment URL keeps its extra query params on SDK requests. The new SDK 3 refusal codes map to 500, and the file download change goes back to main.
This commit is contained in:
mateo-berri 2026-10-28 08:42:15 -07:00
parent 2915b1f0a8
commit 924ac590b1
8 changed files with 138 additions and 140 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/files/streaming.py
View file

@ -5,9 +5,7 @@ from typing import TYPE_CHECKING, Final, Optional, cast
import anyio
from litellm.exceptions import APIConnectionError, Timeout
from litellm.files.types import FileContentProvider
from litellm.llms.custom_httpx.transport_errors import TransportErrorKind, chained_transport_error_kind
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
@ -61,11 +59,8 @@ class FileContentStreamingResponse:
self._log_success_sync()
raise
except Exception as e:
public_error: Final = self._public_error(e)
self._log_failure_sync(public_error)
if public_error is e:
raise
raise public_error from e
self._log_failure_sync(e)
raise
def __aiter__(self) -> "FileContentStreamingResponse":
if not hasattr(self.stream_iterator, "__anext__"):
@ -82,28 +77,8 @@ class FileContentStreamingResponse:
await self._log_success_async()
raise
except Exception as e:
public_error: Final = self._public_error(e)
await self._log_failure_async(public_error)
if public_error is e:
raise
raise public_error from e
def _public_error(self, error: Exception) -> Exception:
match chained_transport_error_kind(error):
case TransportErrorKind.TIMEOUT:
return Timeout(
message=f"timed out reading file content: {error}",
model=self.model,
llm_provider=self.custom_llm_provider,
)
case TransportErrorKind.DROPPED_CONNECTION:
return APIConnectionError(
message=f"connection dropped while reading file content: {error}",
llm_provider=self.custom_llm_provider,
model=self.model,
)
case None:
return error
await self._log_failure_async(e)
raise
async def aclose(self) -> None:
if self._close_completed:

33
litellm/litellm_core_utils/streaming_handler.py
View file

@ -93,6 +93,16 @@ def _next_sync_or_exhausted(it: Iterator[object]) -> object:
return _SYNC_ITER_EXHAUSTED
def _first_choice_index(chunk: ChatCompletionChunk | ModelResponseStream) -> int:
match chunk:
case ChatCompletionChunk(choices=[first, *_]):
return first.index
case ModelResponseStream(choices=[first, *_]):
return first.index
case _:
return 0
def _stamp_served_service_tier(response: ModelResponseStream, complete_streaming_response: ModelResponse) -> None:
served_tier: Final = complete_streaming_response.model_dump().get("service_tier")
if isinstance(served_tier, str) and served_tier:
@ -333,6 +343,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.special_tokens = [
"<|assistant|>",
"<|system|>",
@ -1091,7 +1102,7 @@ class CustomStreamWrapper:
"""
return delta is not None and stored_attribute(delta, attribute_name) is not None
def _copy_delta_attribute(self, source_delta, target_delta, attribute_name: str) -> None:
def _copy_delta_attribute(self, source_delta: object, target_delta: object, attribute_name: str) -> None:
"""
Copy a specific attribute from source delta to target delta.
"""
@ -1575,6 +1586,9 @@ 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
}
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)
@ -2205,12 +2219,25 @@ class CustomStreamWrapper:
return next(self.completion_stream)
except Exception as e:
if (
self.received_finish_reason is None
not self._every_choice_finished()
or chained_transport_error_kind(e) is not TransportErrorKind.DROPPED_CONNECTION
):
raise
raise StopIteration from e
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
def _requested_choice_count(self) -> int:
match getattr(self.logging_obj, "optional_params", None):
case {"n": int() as requested}:
return requested
case _:
return 1
def fetch_sync_stream(self):
if self.completion_stream is None and self.make_call is not None:
# Call make_call to get the completion stream
@ -2321,7 +2348,7 @@ class CustomStreamWrapper:
return await self._finalize_completed_stream(cache_hit=cache_hit)
except Exception as e:
match chained_transport_error_kind(e):
case TransportErrorKind.DROPPED_CONNECTION if self.received_finish_reason is not None:
case TransportErrorKind.DROPPED_CONNECTION if self._every_choice_finished():
return await self._finalize_completed_stream(cache_hit=cache_hit)
case TransportErrorKind.TIMEOUT:
traceback_exception: Final = (

35
litellm/llms/azure/common_utils.py
View file

@ -8,7 +8,7 @@ from contextlib import closing
from functools import lru_cache
from types import MappingProxyType
from typing import Final, Literal, NamedTuple, cast
from urllib.parse import urlsplit, urlunsplit
from urllib.parse import SplitResult, parse_qsl, urlsplit, urlunsplit
import httpx
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
@ -334,23 +334,28 @@ _DEPLOYMENT_OPERATION_PATH: Final = re.compile(
)
def _azure_deployment_base_url(deployment_url: str) -> str:
parts: Final = urlsplit(deployment_url)
operation_path: Final = _DEPLOYMENT_OPERATION_PATH.fullmatch(parts.path)
deployment_path: Final = parts.path if operation_path is None else operation_path["deployment"]
return urlunsplit(parts._replace(path=deployment_path, query="", fragment=""))
def _azure_deployment_base_url(deployment_url: SplitResult) -> str:
operation_path: Final = _DEPLOYMENT_OPERATION_PATH.fullmatch(deployment_url.path)
deployment_path: Final = deployment_url.path if operation_path is None else operation_path["deployment"]
return urlunsplit(deployment_url._replace(path=deployment_path, query="", fragment=""))
def select_azure_base_url_or_endpoint(azure_client_params: dict):
def _azure_deployment_extra_query(deployment_url: SplitResult) -> Mapping[str, str]:
return {name: value for name, value in parse_qsl(deployment_url.query) if name != "api-version"}
def select_azure_base_url_or_endpoint(azure_client_params: Mapping[str, object]):
azure_endpoint: Final = azure_client_params.get("azure_endpoint", None)
if isinstance(azure_endpoint, str):
# see : https://github.com/openai/openai-python/blob/3d61ed42aba652b547029095a7eb269ad4e1e957/src/openai/lib/azure.py#L192
if "/openai/deployments" in azure_endpoint:
# this is base_url, not an azure_endpoint
azure_client_params["base_url"] = _azure_deployment_base_url(azure_endpoint)
azure_client_params.pop("azure_endpoint")
return azure_client_params
# see : https://github.com/openai/openai-python/blob/3d61ed42aba652b547029095a7eb269ad4e1e957/src/openai/lib/azure.py#L192
if not isinstance(azure_endpoint, str) or "/openai/deployments" not in azure_endpoint:
return dict(azure_client_params)
deployment_url: Final = urlsplit(azure_endpoint)
extra_query: Final = _azure_deployment_extra_query(deployment_url)
return {
**{name: value for name, value in azure_client_params.items() if name != "azure_endpoint"},
"base_url": _azure_deployment_base_url(deployment_url),
**({"default_query": extra_query} if extra_query else {}),
}
def get_azure_ad_token(

6
litellm/responses/streaming_iterator.py
View file

@ -228,9 +228,9 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(
"insufficient_quota": 429,
"vector_store_timeout": 504,
"invalid_prompt": 400,
"data_residency_mismatch": 400,
"bio_policy": 400,
"misalignment_policy_violation": 403,
"data_residency_mismatch": 500,
"bio_policy": 500,
"misalignment_policy_violation": 500,
"invalid_image": 400,
"invalid_image_format": 400,
"invalid_base64_image": 400,

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

@ -4760,7 +4760,9 @@ class TransportErrorAfterChunksIterator:
raise StopAsyncIteration
def _reset_test_chunk(content: Optional[str] = None, finish_reason: Optional[str] = None) -> ModelResponseStream:
def _reset_test_chunk(
content: Optional[str] = None, finish_reason: Optional[str] = None, index: int = 0
) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-reset-test",
created=1783458104,
@ -4768,7 +4770,7 @@ def _reset_test_chunk(content: Optional[str] = None, finish_reason: Optional[str
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
index=index,
delta=Delta(content=content),
finish_reason=finish_reason,
)
@ -4940,6 +4942,47 @@ async def test_sdk3_dropped_connection_after_finish_reason_ends_the_stream_clean
assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"]
def _two_choice_stream(
finished_choices: Sequence[int], sync_mode: bool, logging_obj: Logging
) -> CustomStreamWrapper:
logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
return _openai_stream_over(
(
_reset_test_chunk(content="Hi", index=0),
_reset_test_chunk(content="Hel", index=1),
*(_reset_test_chunk(finish_reason="stop", index=index) for index in finished_choices),
),
_httpx2_read_error(),
sync_mode,
logging_obj,
)
@pytest.mark.asyncio
@_SYNC_AND_ASYNC
async def test_sdk3_dropped_connection_while_another_requested_choice_streams_raises(
logging_obj: Logging, sync_mode: bool
):
wrapper: Final = _two_choice_stream((0,), sync_mode, logging_obj)
with pytest.raises(MidStreamFallbackError) as excinfo:
await _collect(wrapper, sync_mode)
assert isinstance(excinfo.value.original_exception, litellm.APIConnectionError)
@pytest.mark.asyncio
@_SYNC_AND_ASYNC
async def test_sdk3_dropped_connection_after_every_requested_choice_finished_ends_cleanly(
logging_obj: Logging, sync_mode: bool
):
wrapper: Final = _two_choice_stream((0, 1), sync_mode, logging_obj)
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

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

@ -2360,3 +2360,35 @@ def test_an_azure_client_litellm_built_its_own_http_client_for_is_still_closed(m
closer.reap()
assert wrapper.is_closed() is True
def test_full_deployment_url_keeps_its_extra_query_params_on_the_sdk_request(monkeypatch: pytest.MonkeyPatch) -> None:
def echo_url(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"id": "chatcmpl-url",
"object": "chat.completion",
"created": 0,
"model": "gpt",
"choices": [
{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": str(request.url)}}
],
},
)
monkeypatch.setattr(litellm, "client_session", httpx.Client(transport=httpx.MockTransport(echo_url)))
client: Final = BaseAzureLLM().get_azure_openai_client(
api_key="azure-key",
api_base="https://gateway.example.com/openai/deployments/gpt/chat/completions"
"?api-version=2024-10-21&subscription-key=abc",
api_version="2025-02-01-preview",
litellm_params={},
_is_async=False,
)
response: Final = client.chat.completions.create(model="gpt", messages=[{"role": "user", "content": "hi"}])
sent: Final = httpx.URL(response.choices[0].message.content or "")
assert sent.path == "/openai/deployments/gpt/chat/completions"
assert dict(sent.params) == {"api-version": "2025-02-01-preview", "subscription-key": "abc"}

86
tests/unit/llms/openai/test_openai_file_content_streaming.py
View file

@ -1,13 +1,6 @@
import importlib
from types import ModuleType
import pytest
from typing import AsyncIterator, Final, Iterator, cast
from typing import AsyncIterator, Iterator, cast
from openai import AsyncOpenAI, OpenAI
from openai._types import Response as SDKResponse
import litellm
from litellm.files import main as files_main
from litellm.files.streaming import FileContentStreamingResponse
from litellm.files.types import FileContentStreamingResult
@ -330,80 +323,3 @@ def test_file_content_streaming_passes_exception_to_context_manager_exit():
assert response_cm.exc_info[0] is RuntimeError
assert response_cm.exc_info[1] is exc_info.value
assert response_cm.exc_info[2] is not None
@pytest.fixture(params=("httpx", "sdk"))
def http_backend(request: pytest.FixtureRequest) -> ModuleType:
return importlib.import_module("httpx" if request.param == "httpx" else SDKResponse.__module__)
_TRANSPORT_FAILURES: Final = pytest.mark.parametrize(
("failure_name", "public_error"),
(
("ReadTimeout", litellm.Timeout),
("ReadError", litellm.APIConnectionError),
("RemoteProtocolError", litellm.APIConnectionError),
),
)
def _body_then_failure_transport(http_backend: ModuleType, failure: Exception) -> object:
class BodyThenFailure(http_backend.SyncByteStream, http_backend.AsyncByteStream):
def __iter__(self) -> Iterator[bytes]:
yield b"hello"
raise failure
async def __aiter__(self) -> AsyncIterator[bytes]:
yield b"hello"
raise failure
return http_backend.MockTransport(
lambda _request: http_backend.Response(
200, headers={"content-type": "application/octet-stream"}, stream=BodyThenFailure()
)
)
@_TRANSPORT_FAILURES
def test_file_content_stream_raises_litellm_errors_when_the_transport_fails_mid_body(
http_backend: ModuleType, failure_name: str, public_error: type[Exception]
):
failure: Final = getattr(http_backend, failure_name)("transport failed mid-body")
with http_backend.Client(transport=_body_then_failure_transport(http_backend, failure)) as http_client:
result: Final = litellm.file_content(
file_id="file-abc123",
custom_llm_provider="openai",
client=OpenAI(api_key="transport-only", http_client=http_client),
stream=True,
chunk_size=len(b"hello"),
)
stream: Final = cast(Iterator[bytes], result.stream_iterator)
assert next(stream) == b"hello"
with pytest.raises(public_error) as excinfo:
next(stream)
assert excinfo.value.__cause__ is failure
@pytest.mark.asyncio
@_TRANSPORT_FAILURES
async def test_afile_content_stream_raises_litellm_errors_when_the_transport_fails_mid_body(
http_backend: ModuleType, failure_name: str, public_error: type[Exception]
):
failure: Final = getattr(http_backend, failure_name)("transport failed mid-body")
async with http_backend.AsyncClient(transport=_body_then_failure_transport(http_backend, failure)) as http_client:
result: Final = await litellm.afile_content(
file_id="file-abc123",
custom_llm_provider="openai",
client=AsyncOpenAI(api_key="transport-only", http_client=http_client),
stream=True,
chunk_size=len(b"hello"),
)
stream: Final = cast(AsyncIterator[bytes], result.stream_iterator)
assert await stream.__anext__() == b"hello"
with pytest.raises(public_error) as excinfo:
await stream.__anext__()
assert excinfo.value.__cause__ is failure

6
tests/unit/responses/test_streaming_iterator_error_events.py
View file

@ -687,9 +687,9 @@ def test_every_openai_sdk_response_error_code_has_explicit_status_mapping():
("empty_image_file", 400),
("failed_to_download_image", 400),
("image_file_not_found", 400),
("data_residency_mismatch", 400),
("bio_policy", 400),
("misalignment_policy_violation", 403),
("data_residency_mismatch", 500),
("bio_policy", 500),
("misalignment_policy_violation", 500),
("totally_unknown_future_code", 500),
("429", 429),
("503", 503),