scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(vertex_ai): append rawPredict suffix for custom api_base on /v1/messages (#31529)

Browse source
This commit is contained in:
Shivam Rawat 2026-06-27 14:54:57 -07:00 • committed by GitHub
parent 71e69d3485
commit d515e5bf05
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 76 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py
View file

@ -49,17 +49,15 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
)
headers["Authorization"] = f"Bearer {access_token}"
# Calculate api_base if not provided
if api_base is None:
api_base = self.get_complete_vertex_url(
custom_api_base=api_base,
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
project_id=project_id or "",
partner=VertexPartnerProvider.claude,
stream=optional_params.get("stream", False),
model=model,
)
api_base = self.get_complete_vertex_url(
custom_api_base=api_base,
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
project_id=project_id or "",
partner=VertexPartnerProvider.claude,
stream=optional_params.get("stream", False),
model=model,
)
headers["content-type"] = "application/json"

67
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py
View file

@ -266,6 +266,73 @@ def test_validate_environment_always_refreshes_token_ignoring_stale_bearer():
assert api_base == "https://mock-vertex-url"
def test_validate_environment_appends_stream_raw_predict_with_custom_api_base():
"""Regression: a custom api_base on /v1/messages must still get the endpoint
suffix appended. The old `if api_base is None` guard skipped
get_complete_vertex_url entirely, leaving the api_base without
`:streamRawPredict`."""
config = VertexAIPartnerModelsAnthropicMessagesConfig()
litellm_params = {
"vertex_ai_project": "test-project",
"vertex_ai_location": "us-central1",
}
with (
patch.object(
config,
"get_complete_vertex_url",
wraps=config.get_complete_vertex_url,
) as spy_get_url,
patch.object(
config, "_ensure_access_token", return_value=("token", "test-project")
),
):
_, api_base = config.validate_anthropic_messages_environment(
headers={},
model="claude-sonnet-4",
messages=[],
optional_params={"stream": True},
litellm_params=litellm_params,
api_base="https://my-proxy.example.com",
)
spy_get_url.assert_called_once()
assert api_base is not None
assert ":streamRawPredict" in api_base
def test_validate_environment_appends_raw_predict_with_custom_api_base():
"""Regression: non-streaming custom api_base must end with `:rawPredict`."""
config = VertexAIPartnerModelsAnthropicMessagesConfig()
litellm_params = {
"vertex_ai_project": "test-project",
"vertex_ai_location": "us-central1",
}
with (
patch.object(
config,
"get_complete_vertex_url",
wraps=config.get_complete_vertex_url,
) as spy_get_url,
patch.object(
config, "_ensure_access_token", return_value=("token", "test-project")
),
):
_, api_base = config.validate_anthropic_messages_environment(
headers={},
model="claude-sonnet-4",
messages=[],
optional_params={},
litellm_params=litellm_params,
api_base="https://my-proxy.example.com",
)
spy_get_url.assert_called_once()
assert api_base is not None
assert api_base.endswith(":rawPredict")
def test_transform_anthropic_messages_request_removes_scope_from_cache_control():
"""Ensure scope field is removed from cache_control for Vertex AI (not supported)."""
config = VertexAIPartnerModelsAnthropicMessagesConfig()