diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 14aebcaabaf..6d050d5a856 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -25,6 +25,21 @@ from litellm.types.llms.vertex_ai import ( from litellm.types.utils import TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages +VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS: Final = frozenset( + { + "audio", + "max_retries", + "modalities", + "prediction", + "prompt_cache_key", + "prompt_cache_retention", + "safety_identifier", + "service_tier", + "store", + "web_search_options", + } +) + class VertexAILyriaModelInfo(TypedDict): vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"]] @@ -370,6 +385,27 @@ def get_vertex_base_model_name(model: str) -> str: return model +def vertex_model_garden_model_id_in_json_body(model: str) -> bool: + """ + Vertex catalog / publisher models are addressed as publisher/model (e.g. + xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body. + + Deployed Model Garden endpoints are typically a single segment (often numeric) + and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field. + """ + return "/" in model + + +def is_vertex_self_deployed_openai_compatible_endpoint(model: str) -> bool: + local_model: Final = model.removeprefix("vertex_ai/") + route: Final = get_vertex_ai_model_route(local_model) + if route == VertexAIModelRoute.GEMMA: + return True + return route == VertexAIModelRoute.MODEL_GARDEN and not vertex_model_garden_model_id_in_json_body( + get_vertex_base_model_name(local_model) + ) + + def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: """ Fine-tuned Gemini deployments are addressed by a numeric endpoint id, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index f2d2c0896d2..ca0bcb74906 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -18,7 +18,11 @@ from litellm.types.utils import ( Usage, ) -from ...common_utils import VertexAIError +from ...common_utils import ( + VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS, + VertexAIError, + is_vertex_self_deployed_openai_compatible_endpoint, +) if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer @@ -66,13 +70,15 @@ class VertexAILlama3Config(OpenAIGPTConfig): and v is not None } - def get_supported_openai_params(self, model: str): - supported_params: Final = super().get_supported_openai_params(model=model) - try: - supported_params.remove("max_retries") - except KeyError: - pass - return supported_params + def get_supported_openai_params(self, model: str) -> list[str]: + unsupported_params: Final = ( + VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS + if is_vertex_self_deployed_openai_compatible_endpoint(model) + else frozenset({"max_retries"}) + ) + return [ # mutable-ok: get_optional_params extends the returned list with allowed_openai_params + param for param in super().get_supported_openai_params(model=model) if param not in unsupported_params + ] def map_openai_params( self, diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index a774dba6cf2..ea97f0a0a9a 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, ) from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.llms.vertex_ai.common_utils import VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai_gemma import VertexGemmaContainerError from litellm.types.utils import ModelResponse @@ -49,6 +50,13 @@ class VertexGemmaConfig(OpenAIGPTConfig): def __init__(self) -> None: super().__init__() + def get_supported_openai_params(self, model: str) -> list[str]: + return [ # mutable-ok: get_optional_params extends the returned list with allowed_openai_params + param + for param in super().get_supported_openai_params(model=model) + if param not in VERTEX_SELF_DEPLOYED_ENDPOINT_UNSUPPORTED_PARAMS + ] + def should_fake_stream( self, model: str | None, diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index f5c9ac623a1..84907f01685 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -24,21 +24,14 @@ import httpx from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.utils import ModelResponse -from ..common_utils import VertexAIError, get_vertex_base_model_name +from ..common_utils import ( + VertexAIError, + get_vertex_base_model_name, + vertex_model_garden_model_id_in_json_body, +) from ..vertex_llm_base import VertexBase -def _vertex_model_garden_model_id_in_json_body(model: str) -> bool: - """ - Vertex catalog / publisher models are addressed as publisher/model (e.g. - xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body. - - Deployed Model Garden endpoints are typically a single segment (often numeric) - and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field. - """ - return "/" in model - - def create_vertex_url( vertex_location: str, vertex_project: str, @@ -48,7 +41,7 @@ def create_vertex_url( ) -> str: """Return the api base for vertex model garden (without /chat/completions).""" base_url: Final = get_vertex_base_url(vertex_location) - if _vertex_model_garden_model_id_in_json_body(model): + if vertex_model_garden_model_id_in_json_body(model): return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi" return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}" @@ -124,7 +117,7 @@ class VertexAIModelGardenModels(VertexBase): ) # Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route). # Single-segment endpoint ids: model is encoded in the URL path; body model stays empty. - if not _vertex_model_garden_model_id_in_json_body(model): + if not vertex_model_garden_model_id_in_json_body(model): model = "" return openai_like_chat_completions.completion( model=model, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index d03174bc2c6..2fe22ba2620 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1046,11 +1046,26 @@ def test_translate_anthropic_to_openai_sets_prompt_cache_key_for_azure(model: st assert openai_request["prompt_cache_key"] == "session-abc" +@pytest.mark.parametrize( + "model", + [ + "vertex_ai/meta/llama-4-maverick-17b-128e-instruct-maas", + "vertex_ai/moonshotai/kimi-k2-thinking-maas", + "vertex_ai/xai/grok-4.1-fast-non-reasoning", + ], +) +def test_translate_anthropic_to_openai_sets_prompt_cache_key_for_vertex_maas_models(model: str): + openai_request = _translate_with_metadata(model, {"user_id": CLAUDE_CODE_USER_ID}, "vertex_ai") + assert openai_request["prompt_cache_key"] == "session-abc" + + @pytest.mark.parametrize( "model, custom_llm_provider", [ ("gemini/gemini-2.5-pro", "gemini"), ("vertex_ai/gemini-2.5-pro", "vertex_ai"), + ("vertex_ai/gemma/gemma-2-2b-it", "vertex_ai"), + ("vertex_ai/openai/mg-endpoint-lit8592", "vertex_ai"), ("anthropic/claude-sonnet-4-5", "anthropic"), ("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", "bedrock"), ("no-such-model-lit5875", "no-such-provider-lit5875"), diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py index 0dcaa4c72c2..b80d4714253 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py @@ -8,10 +8,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm -from litellm.llms.vertex_ai.vertex_model_garden.main import ( - _vertex_model_garden_model_id_in_json_body, - create_vertex_url, +from litellm.llms.vertex_ai.common_utils import ( + vertex_model_garden_model_id_in_json_body, ) +from litellm.llms.vertex_ai.vertex_model_garden.main import create_vertex_url @pytest.mark.parametrize( @@ -43,11 +43,8 @@ def test_create_vertex_url_openapi_vs_deployed_endpoint( def test_model_id_in_json_body_heuristic() -> None: - assert ( - _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") - is True - ) - assert _vertex_model_garden_model_id_in_json_body("5464397967697903616") is False + assert vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") is True + assert vertex_model_garden_model_id_in_json_body("5464397967697903616") is False @pytest.fixture diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py index 3bca51ec6b3..05e4e36edd7 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py @@ -11,6 +11,39 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.llama3.transformation impor ) +OPENAI_PLATFORM_PARAMS = ( + "prompt_cache_key", + "prompt_cache_retention", + "safety_identifier", + "service_tier", + "store", + "web_search_options", + "modalities", + "prediction", + "audio", +) + +SELF_DEPLOYED_ENDPOINT_MODELS = ( + "gemma/gemma-2-2b-it", + "vertex_ai/gemma/gemma-2-2b-it", + "openai/mg-endpoint-lit8592", + "vertex_ai/openai/mg-endpoint-lit8592", + "openai/5464397967697903616", +) + +MAAS_MODELS = ( + "meta/llama-4-maverick-17b-128e-instruct-maas", + "vertex_ai/meta/llama-4-maverick-17b-128e-instruct-maas", + "moonshotai/kimi-k2-thinking-maas", + "qwen/qwen3-next-80b-a3b-instruct-maas", + "google/gemma-4-26b-a4b-it-maas", + "xai/grok-4.1-fast-non-reasoning", + "openai/xai/grok-4.1-fast-reasoning", + "1984786713414729728", + "llama3", +) + + class TestVertexAILlama3Config: def test_transform_choices(self): """ @@ -56,6 +89,52 @@ class TestVertexAILlama3Config: assert response[0].message.tool_calls is not None assert response[0].finish_reason == "tool_calls" + @pytest.mark.parametrize("model", SELF_DEPLOYED_ENDPOINT_MODELS) + @pytest.mark.parametrize("param", OPENAI_PLATFORM_PARAMS) + def test_get_supported_openai_params_omits_platform_params_for_self_deployed_endpoints( + self, model: str, param: str + ): + assert param not in VertexAILlama3Config().get_supported_openai_params(model=model) + + @pytest.mark.parametrize("model", MAAS_MODELS) + @pytest.mark.parametrize("param", OPENAI_PLATFORM_PARAMS) + def test_get_supported_openai_params_keeps_platform_params_for_maas_models(self, model: str, param: str): + assert param in VertexAILlama3Config().get_supported_openai_params(model=model) + + @pytest.mark.parametrize("model", [*SELF_DEPLOYED_ENDPOINT_MODELS, *MAAS_MODELS]) + def test_get_supported_openai_params_never_lists_max_retries(self, model: str): + assert "max_retries" not in VertexAILlama3Config().get_supported_openai_params(model=model) + + @pytest.mark.parametrize("model", [*SELF_DEPLOYED_ENDPOINT_MODELS, *MAAS_MODELS]) + @pytest.mark.parametrize( + "param", + ["max_completion_tokens", "tools", "tool_choice", "response_format", "seed", "logprobs", "parallel_tool_calls"], + ) + def test_get_supported_openai_params_keeps_params_every_vertex_openai_endpoint_accepts( + self, model: str, param: str + ): + assert param in VertexAILlama3Config().get_supported_openai_params(model=model) + + @pytest.mark.parametrize("model", SELF_DEPLOYED_ENDPOINT_MODELS) + def test_map_openai_params_drops_prompt_cache_key_for_self_deployed_endpoints(self, model: str): + mapped = VertexAILlama3Config().map_openai_params( + {"prompt_cache_key": "session-lit8592", "max_completion_tokens": 10}, + {}, + model, + drop_params=True, + ) + assert mapped == {"max_tokens": 10} + + @pytest.mark.parametrize("model", MAAS_MODELS) + def test_map_openai_params_forwards_prompt_cache_key_for_maas_models(self, model: str): + mapped = VertexAILlama3Config().map_openai_params( + {"prompt_cache_key": "session-lit8592", "max_completion_tokens": 10}, + {}, + model, + drop_params=True, + ) + assert mapped == {"prompt_cache_key": "session-lit8592", "max_tokens": 10} + class TestVertexAILlama3StreamingHandler: def test_first_chunk_has_role_assistant_when_missing(self): diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index e9ae5234094..e5ca31833ce 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -694,6 +694,89 @@ class TestVertexGemmaCompletion: assert instance["@requestFormat"] == "chatCompletions" assert "messages" in instance + @pytest.mark.parametrize( + "param", + [ + "prompt_cache_key", + "prompt_cache_retention", + "safety_identifier", + "service_tier", + "store", + "web_search_options", + "modalities", + "prediction", + "audio", + "max_retries", + ], + ) + def test_get_supported_openai_params_omits_params_the_predict_endpoint_rejects(self, param: str): + from litellm.llms.vertex_ai.vertex_gemma_models.transformation import ( + VertexGemmaConfig, + ) + + assert param not in VertexGemmaConfig().get_supported_openai_params(model="gemma-2-2b-it") + + @pytest.mark.asyncio + async def test_acompletion_drops_prompt_cache_key_when_drop_params_is_set(self): + with ( + patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID"), + ), + ): + mock_client = Mock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = _make_gemma_vertex_response() + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + await litellm.acompletion( + model="vertex_ai/gemma/gemma-2-2b-it", + messages=[{"role": "user", "content": "Test"}], + prompt_cache_key="session-lit8592", + service_tier="default", + max_completion_tokens=16, + drop_params=True, + api_base="https://test.us-central1-project.prediction.vertexai.goog/v1/projects/PROJECT_ID/locations/us-central1/endpoints/ENDPOINT_ID:predict", + vertex_project="PROJECT_ID", + vertex_location="us-central1", + ) + + instance = mock_client.post.call_args.kwargs["json"]["instances"][0] + assert "prompt_cache_key" not in instance + assert "service_tier" not in instance + assert instance["max_tokens"] == 16 + assert instance["messages"] == [{"role": "user", "content": "Test"}] + + @pytest.mark.asyncio + async def test_acompletion_rejects_prompt_cache_key_before_calling_vertex(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + with ( + patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client") as mock_get_client, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID"), + ), + ): + mock_client = Mock() + mock_client.post = AsyncMock() + mock_get_client.return_value = mock_client + + with pytest.raises(litellm.UnsupportedParamsError, match="prompt_cache_key"): + await litellm.acompletion( + model="vertex_ai/gemma/gemma-2-2b-it", + messages=[{"role": "user", "content": "Test"}], + prompt_cache_key="session-lit8592", + drop_params=False, + api_base="https://test.us-central1-project.prediction.vertexai.goog/v1/projects/PROJECT_ID/locations/us-central1/endpoints/ENDPOINT_ID:predict", + vertex_project="PROJECT_ID", + vertex_location="us-central1", + ) + + mock_client.post.assert_not_called() + def test_transform_request_strips_context_management(self): """ Direct unit test for VertexGemmaConfig.transform_request: verify that