fix(vertex_ai): forward extra_body to completion transformation handler (#20950)

* fix(vertex_ai): forward extra_body to completion transformation handler

The responses() function accepted extra_body as a named parameter but
did not pass it to response_api_handler when responses_api_provider_config
was None (completion transformation path), silently dropping it.

Also adds deep-merge support for extra_body in Vertex AI Gemini
transformation, so dict values like generationConfig are merged rather
than replaced.

* refactor(vertex_ai): extract _merge_extra_body to fix PLR0915 lint

Move the extra_body merge loop into a helper function to keep
_transform_request_body under the 50-statement limit.
This commit is contained in:
Lei Nie 2026-02-11 18:52:56 -08:00 • committed by Sameer Kankute
parent e85da33240
commit e9c99f41bd
3 changed files with 45 additions and 0 deletions

View file

@ -529,6 +529,17 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
raise e
def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None:
"""Pop extra_body from optional_params and shallow-merge into data, deep-merging dict values."""
extra_body: Optional[dict] = optional_params.pop("extra_body", None)
if extra_body is not None:
for k, v in extra_body.items():
if k in data and isinstance(data[k], dict) and isinstance(v, dict):
data[k].update(v)
else:
data[k] = v
def _transform_request_body(
messages: List[AllMessageValues],
model: str,
@ -619,6 +630,7 @@ def _transform_request_body(
# Only add labels for Vertex AI endpoints (not Google GenAI/AI Studio) and only if non-empty
if labels and custom_llm_provider != LlmProviders.GEMINI:
data["labels"] = labels
_pop_and_merge_extra_body(data, optional_params)
except Exception as e:
raise e

View file

@ -682,6 +682,7 @@ def responses(
_is_async=_is_async,
stream=stream,
extra_headers=extra_headers,
extra_body=extra_body,
**kwargs,
)

View file

@ -2,6 +2,7 @@ import base64
import json
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
from fastapi.testclient import TestClient
@ -352,3 +353,34 @@ class TestResponsesAPIProviderSpecificParams:
# Should not raise any exception
result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params)
assert "temperature" in result
def test_responses_extra_body_forwarded_to_completion_transformation_handler():
"""
Regression test: extra_body must be forwarded to response_api_handler
when responses_api_provider_config is None (completion transformation path).
Before the fix, extra_body was a named parameter of responses() but was
not passed to litellm_completion_transformation_handler.response_api_handler(),
so it was silently dropped.
"""
with patch(
"litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config",
return_value=None,
), patch(
"litellm.responses.main.litellm_completion_transformation_handler.response_api_handler",
) as mock_handler:
mock_handler.return_value = MagicMock()
litellm.responses(
model="openai/gpt-4o",
input="Hello",
extra_body={"custom_key": "custom_value"},
)
mock_handler.assert_called_once()
call_kwargs = mock_handler.call_args
# extra_body can be a positional or keyword arg; check both
assert call_kwargs.kwargs.get("extra_body") == {
"custom_key": "custom_value"
}