mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
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:
parent
e85da33240
commit
e9c99f41bd
3 changed files with 45 additions and 0 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -682,6 +682,7 @@ def responses(
|
|||
_is_async=_is_async,
|
||||
stream=stream,
|
||||
extra_headers=extra_headers,
|
||||
extra_body=extra_body,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue