scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

[Fix] VertexAI - gemma model family support (custom endpoints) (#15419)

Browse source 
* 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 fdaa5bc49e.

* test_acompletion_basic_request

* fix: async transform

* fix gemma: stream param

* test_acompletion_fake_streaming
This commit is contained in:
Ishaan Jaff 2025-10-10 13:30:43 -07:00 • committed by GitHub
parent c8b93c3940
commit ed62d6c943
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 189 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

100
litellm/llms/vertex_ai/vertex_gemma_models/transformation.py
View file

@ -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)

33
litellm/model_prices_and_context_window_backup.json
View file

@ -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",

114
tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py
View file

@ -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"