diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f6340426c1b..0f46f1b718c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -8,11 +8,12 @@ import time import traceback from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from typing_extensions import NotRequired, TypedDict import litellm @@ -182,6 +183,23 @@ class _VertexChunkLike(Protocol): candidates: Sequence[_VertexCandidateLike] +class _ParsedChunkHiddenParams(BaseModel): + provider_specific_fields: Mapping[str, object] | None = None + + +def _provider_hidden_params(chunk: object) -> Mapping[str, object] | None: + hidden: Final[object] = getattr(chunk, "_hidden_params", None) + if not isinstance(hidden, dict): + return None + try: + parsed: Final = _ParsedChunkHiddenParams.model_validate(hidden) + except ValidationError: + return None + if not parsed.provider_specific_fields: + return None + return MappingProxyType({"provider_specific_fields": dict(parsed.provider_specific_fields)}) + + class CustomStreamWrapper: def __init__( self, @@ -801,7 +819,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def model_response_creator(self, chunk: dict | None = None, hidden_params: dict | None = None): + def model_response_creator(self, chunk: dict | None = None, hidden_params: Mapping[str, object] | None = None): _model: Final = self._cached_model_name _logging_obj_llm_provider: Final = self._cached_logging_llm_provider @@ -1504,7 +1522,7 @@ class CustomStreamWrapper: def chunk_creator(self, chunk: Any): if hasattr(chunk, "id"): self.response_id = chunk.id - model_response = self.model_response_creator() + model_response = self.model_response_creator(hidden_params=_provider_hidden_params(chunk)) response_obj: dict[str, Any] = {} try: # return this for all models diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index b5e33a4e421..5329edce47e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4460,3 +4460,51 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp finally: trace_id_var.set("") session_id_var.set("") + + +def test_chunk_creator_preserves_hidden_provider_specific_fields_from_parsed_chunk(): + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + ) + parsed_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="hello", role="assistant"), finish_reason=None)], + ) + parsed_chunk._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND_FLEX"} + + result = wrapper.chunk_creator(chunk=parsed_chunk) + + assert result is not None + assert result._hidden_params["provider_specific_fields"] == {"traffic_type": "ON_DEMAND_FLEX"} + + +@pytest.mark.asyncio +async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging_obj: Logging): + content_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="hello", role="assistant"), finish_reason=None)], + ) + final_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + setattr(final_chunk, "usage", Usage(prompt_tokens=7, completion_tokens=5, total_tokens=12)) + final_chunk._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND_FLEX"} + + async def _stream(): + yield content_chunk + yield final_chunk + + wrapper = CustomStreamWrapper( + completion_stream=_stream(), + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + stream_options={"include_usage": True}, + ) + + received = [chunk async for chunk in wrapper] + + assembled = litellm.stream_chunk_builder(chunks=received, messages=[{"role": "user", "content": "hi"}]) + assert assembled is not None + assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX"