From ed62d6c94341e5f72cdcacbb055e5103c0ba3744 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 10 Oct 2025 13:30:43 -0700 Subject: [PATCH] [Fix] VertexAI - gemma model family support (custom endpoints) (#15419) * TestVertexGemmaiCompletion * test vertex Gemma * fix file name * fix file naming * add VertexAIGemmaModels * add cost_router for vertexai * fix main.py * fix VertexGemmaConfig * fix Vertex AI Gemma-AI Models Handler * docs gemma * fix ids * test fix * ruff check fixes * docs fix * docs fix * test_acompletion_basic_request * Revert "test_acompletion_basic_request" This reverts commit fdaa5bc49ebc61773564872ddf255fcf50c2e9e4. * test_acompletion_basic_request * fix: async transform * fix gemma: stream param * test_acompletion_fake_streaming --- .../vertex_gemma_models/transformation.py | 100 +++++++-------- ...odel_prices_and_context_window_backup.json | 33 +++++ .../test_vertex_gemma_transformation.py | 114 ++++++++++++++++-- 3 files changed, 189 insertions(+), 58 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 5541b60f524..47cc429e793 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -13,9 +13,10 @@ from typing import Any, Callable, Dict, List, Optional, Union, cast import httpx from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ModelResponse +from litellm.types.utils import LlmProviders, ModelResponse class VertexGemmaConfig(OpenAIGPTConfig): @@ -29,6 +30,38 @@ class VertexGemmaConfig(OpenAIGPTConfig): def __init__(self) -> None: super().__init__() + def should_fake_stream( + self, + model: Optional[str], + stream: Optional[bool], + custom_llm_provider: Optional[str] = None, + ) -> bool: + """ + Vertex AI Gemma models do not support streaming. + Return True to enable fake streaming on the client side. + """ + return True + + def _handle_fake_stream_response( + self, + model_response: ModelResponse, + stream: bool, + ) -> Union[ModelResponse, Any]: + """ + Helper method to return fake stream iterator if streaming is requested. + + Args: + model_response: The completed model response + stream: Whether streaming was requested + + Returns: + MockResponseIterator if stream=True, otherwise the model_response + """ + if stream: + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + return MockResponseIterator(model_response=model_response) + return model_response + def transform_request( self, model: str, @@ -52,41 +85,9 @@ class VertexGemmaConfig(OpenAIGPTConfig): headers=headers, ) - # Remove 'model' from the request as it's not needed in the instance - openai_request.pop("model", None) - - # Wrap in Vertex Gemma format - return { - "instances": [ - { - "@requestFormat": "chatCompletions", - **openai_request, - } - ] - } - - async def async_transform_request( - self, - model: str, - messages: List[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - """ - Async version of transform_request. - """ - # Get the base OpenAI request from parent class - openai_request = await super().async_transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - - # Remove 'model' from the request as it's not needed in the instance + # Remove params not needed/supported by Vertex Gemma openai_request.pop("model", None) + openai_request.pop("stream", None) # Streaming not supported, will be faked client-side # Wrap in Vertex Gemma format return { @@ -137,16 +138,8 @@ class VertexGemmaConfig(OpenAIGPTConfig): ): """ Make completion request to Vertex Gemma endpoint. - Supports both sync and async requests. + Supports both sync and async requests with fake streaming. """ - # Handle streaming - stream = optional_params.get("stream", False) - if stream: - raise BaseLLMException( - status_code=400, - message="Streaming is not yet supported for Vertex AI Gemma models", - ) - if acompletion: return self._async_completion( model=model, @@ -194,6 +187,9 @@ class VertexGemmaConfig(OpenAIGPTConfig): from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.utils import convert_to_model_response_object + # Check if streaming is requested (will be faked) + stream = optional_params.get("stream", False) + # Transform the request using parent class methods request_data = self.transform_request( model=model, @@ -260,7 +256,8 @@ class VertexGemmaConfig(OpenAIGPTConfig): additional_args={"complete_input_dict": request_data}, ) - return model_response + # Return fake stream iterator if streaming was requested + return self._handle_fake_stream_response(model_response=model_response, stream=stream) async def _async_completion( self, @@ -277,9 +274,13 @@ class VertexGemmaConfig(OpenAIGPTConfig): encoding: Any, ): """Asynchronous completion request""" - from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.utils import LlmProviders from litellm.utils import convert_to_model_response_object + # Check if streaming is requested (will be faked) + stream = optional_params.get("stream", False) + # Transform the request using parent class async methods request_data = await self.async_transform_request( model=model, @@ -306,7 +307,9 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) # Make the HTTP request - http_handler = AsyncHTTPHandler(concurrent_limit=1) + http_handler = get_async_httpx_client( + llm_provider=LlmProviders.VERTEX_AI, + ) response = await http_handler.post( url=api_base, headers=headers, @@ -346,5 +349,6 @@ class VertexGemmaConfig(OpenAIGPTConfig): additional_args={"complete_input_dict": request_data}, ) - return model_response + # Return fake stream iterator if streaming was requested + return self._handle_fake_stream_response(model_response=model_response, stream=stream) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 499a16f967f..826b9b8af57 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19940,6 +19940,39 @@ "supports_parallel_function_calling": true, "supports_tool_choice": true }, + "together_ai/moonshotai/Kimi-K2-Instruct-0905": { + "input_cost_per_token": 1e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://www.together.ai/models/kimi-k2-0905", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true + }, + "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 1ff9009bb8f..3eed79fb161 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -135,6 +135,7 @@ class TestVertexGemmaCompletion: assert call_args is not None, "HTTP handler was not called" request_data = call_args.kwargs["json"] + print("request body=", json.dumps(request_data, indent=4)) request_url = call_args.kwargs["url"] # Validate exact URL matches what we sent @@ -145,17 +146,17 @@ class TestVertexGemmaCompletion: assert "instances" in request_data assert len(request_data["instances"]) == 1 - outer_instance = request_data["instances"][0] - assert outer_instance["@requestFormat"] == "chatCompletions" + instance = request_data["instances"][0] + assert instance["@requestFormat"] == "chatCompletions" - # The actual instance with messages is nested inside - assert "instances" in outer_instance - inner_instance = outer_instance["instances"][0] - assert inner_instance["@requestFormat"] == "chatCompletions" - assert "messages" in inner_instance - assert inner_instance["messages"][0]["role"] == "user" - assert inner_instance["messages"][0]["content"] == "What is machine learning?" - assert inner_instance["max_tokens"] == 100 + # Messages should be directly in the instance, not double-nested + assert "messages" in instance + assert instance["messages"][0]["role"] == "user" + assert instance["messages"][0]["content"] == "What is machine learning?" + assert instance["max_tokens"] == 100 + + # Verify stream parameter is NOT sent to Vertex (will be faked client-side) + assert "stream" not in instance # Validate LiteLLM Response (OpenAI format) assert response.id == "chatcmpl-aaa4288f-2b8e-4bc0-8b14-4e444decd2c4" @@ -215,3 +216,96 @@ class TestVertexGemmaCompletion: # Verify the error message contains the original error assert "missing 'predictions' field" in str(exc_info.value) + @pytest.mark.asyncio + async def test_acompletion_fake_streaming(self): + """ + Test that streaming requests are faked properly for Vertex AI Gemma models. + + Verifies: + 1. Request body does NOT include 'stream' parameter (model doesn't support it) + 2. Response returns a MockResponseIterator that yields chunks + """ + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + + # Mock Vertex response + mock_vertex_response = { + "deployedModelId": "1207280419999999999", + "model": "projects/993702345710/locations/us-central1/models/gemma-3-12b-it-1222199011122", + "modelDisplayName": "gemma-3-12b-it-1222199011122", + "modelVersionId": "1", + "predictions": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "content": "Streaming test response", + "reasoning_content": None, + "role": "assistant", + "tool_calls": [], + }, + "stop_reason": None, + } + ], + "created": 1759863903, + "id": "chatcmpl-test-stream", + "model": "google/gemma-3-12b-it", + "object": "chat.completion", + "prompt_logprobs": None, + "usage": { + "completion_tokens": 3, + "prompt_tokens": 10, + "prompt_tokens_details": None, + "total_tokens": 13, + }, + }, + } + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client: + mock_client = Mock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_vertex_response + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + # Call litellm.acompletion() with stream=True + response = await litellm.acompletion( + model="vertex_ai/gemma/gemma-3-12b-it-1222199011122", + messages=[{"role": "user", "content": "Test streaming"}], + stream=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", + ) + + # Verify the response is a MockResponseIterator + assert isinstance(response, MockResponseIterator), f"Expected MockResponseIterator, got {type(response)}" + + # Verify the request sent to Vertex does NOT include 'stream' + call_args = mock_client.post.call_args + assert call_args is not None, "HTTP client was not called" + + request_data = call_args.kwargs["json"] + instance = request_data["instances"][0] + + # Critical: Verify stream parameter is NOT sent to Vertex API + assert "stream" not in instance, "stream parameter should not be sent to Vertex API" + + # Verify we can iterate the fake stream and get the response + chunks = [] + async for chunk in response: + chunks.append(chunk) + + # Should get exactly one chunk (fake streaming) + assert len(chunks) == 1, f"Expected 1 chunk from fake stream, got {len(chunks)}" + + # Verify the chunk has the expected content + chunk = chunks[0] + assert hasattr(chunk, "choices") + assert len(chunk.choices) > 0 + assert chunk.choices[0].delta.content == "Streaming test response" +