From 36915da9c8bce3acdd4d3924f8038a834228947f Mon Sep 17 00:00:00 2001 From: Mingyang Wu <129849514+aprylewu@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:05:14 +0800 Subject: [PATCH 1/2] fix(google-genai): preserve Vertex location for cost tracking --- litellm/google_genai/main.py | 2 + .../google_genai/test_google_genai_main.py | 106 +++++++++++++++++- 2 files changed, 106 insertions(+), 2 deletions(-) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index c1822e4720d..403369c8f66 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -17,6 +17,7 @@ from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes from litellm.utils import ProviderConfigManager, client @@ -198,6 +199,7 @@ class GenerateContentHelper: optional_params=dict(generate_content_config_dict), litellm_params={ "litellm_call_id": litellm_call_id, + "vertex_location": VertexBase.explicit_vertex_ai_location(litellm_params.model_dump()), }, custom_llm_provider=custom_llm_provider, ) diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/test_litellm/google_genai/test_google_genai_main.py index 238fff7deca..9da2f213d69 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -4,13 +4,115 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json +from datetime import datetime +from typing import Final import pytest - - import litellm +from litellm.google_genai.main import GenerateContentHelper +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig + + +@pytest.mark.parametrize( + "provider,location_kwargs,environment_location,expected_location,expected_cost", + [ + ("vertex_ai", {"vertex_location": "global"}, "us-central1", "global", 5.25e-6), + ( + "vertex_ai", + {"vertex_location": "europe-west4"}, + "global", + "europe-west4", + 5.775e-6, + ), + ( + "vertex_ai", + {"vertex_ai_location": "global"}, + "us-central1", + "global", + 5.25e-6, + ), + ( + "vertex_ai", + {"vertex_location": "global", "vertex_ai_location": "us-central1"}, + "europe-west4", + "global", + 5.25e-6, + ), + ("vertex_ai", {}, "global", "global", 5.25e-6), + ("vertex_ai", {}, None, "us-central1", 5.775e-6), + ("gemini", {"vertex_location": "us-central1"}, "us-central1", None, 5.25e-6), + ], +) +def test_generate_content_prices_the_request_location( + monkeypatch: pytest.MonkeyPatch, + provider: str, + location_kwargs: dict[str, str], + environment_location: str | None, + expected_location: str | None, + expected_cost: float, +) -> None: + monkeypatch.setattr(litellm, "vertex_location", None) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) + if environment_location is None: + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + else: + monkeypatch.setenv("VERTEXAI_LOCATION", environment_location) + + logging_obj: Final = Logging( + model="gemini-flash", + messages=[], + stream=False, + call_type="agenerate_content", + start_time=datetime.now(), + litellm_call_id="test-vertex-location", + function_id="test-vertex-location", + ) + setup: Final = GenerateContentHelper.setup_generate_content_call( + model=f"{provider}/gemini-3.8-flash", + contents=[{"role": "user", "parts": [{"text": "say ok"}]}], + config={"temperature": 0}, + vertex_project="test-project", + api_key="test-key", + litellm_logging_obj=logging_obj, + **location_kwargs, + ) + config: Final = setup.generate_content_provider_config + assert isinstance(config, GoogleGenAIConfig) + credentials, project, location = config._get_common_auth_components(dict(setup.litellm_params)) + _, url = config._build_final_headers_and_url( + model=setup.model, + auth_header="test-token", + vertex_project=project, + vertex_location=location, + vertex_credentials=credentials, + stream=False, + api_base=None, + litellm_params=dict(setup.litellm_params), + ) + if expected_location is None: + assert url.startswith("https://generativelanguage.googleapis.com/") + else: + assert f"/locations/{expected_location}/" in url + assert setup.generate_content_config_dict == {"temperature": 0} + assert "vertex_location" not in setup.request_body + + response: Final = { + "candidates": [ + { + "content": {"parts": [{"text": "ok"}], "role": "model"}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 2, + "candidatesTokenCount": 1, + "totalTokenCount": 3, + }, + } + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(expected_cost) @pytest.mark.asyncio From 38215a8a811489a266f9e62af03c49f32626ac25 Mon Sep 17 00:00:00 2001 From: Mingyang Wu <129849514+aprylewu@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:46:48 +0800 Subject: [PATCH 2/2] fix(google-genai): keep cost logging parameters in provider config --- litellm/google_genai/main.py | 3 +-- litellm/llms/base_llm/google_genai/transformation.py | 4 ++++ litellm/llms/vertex_ai/google_genai/transformation.py | 5 +++++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 403369c8f66..ba7222c041d 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -17,7 +17,6 @@ from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes from litellm.utils import ProviderConfigManager, client @@ -199,7 +198,7 @@ class GenerateContentHelper: optional_params=dict(generate_content_config_dict), litellm_params={ "litellm_call_id": litellm_call_id, - "vertex_location": VertexBase.explicit_vertex_ai_location(litellm_params.model_dump()), + **generate_content_provider_config.get_generate_content_logging_params(litellm_params), }, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index bd0d29d5ea3..7e85d6ccf80 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -1,5 +1,6 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -73,6 +74,9 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): """ return ("safetySettings", "toolConfig", "cachedContent", "labels") + def get_generate_content_logging_params(self, litellm_params: GenericLiteLLMParams) -> Mapping[str, object]: + return types.MappingProxyType({}) + @abstractmethod def map_generate_content_optional_params( self, diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 8c22ae06af0..cbbed2134f9 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -2,6 +2,8 @@ Transformation for Calling Google models in their native format. """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, Literal from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig @@ -20,6 +22,9 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): def custom_llm_provider(self) -> Literal["gemini", "vertex_ai"]: return "vertex_ai" + def get_generate_content_logging_params(self, litellm_params: GenericLiteLLMParams) -> Mapping[str, object]: + return MappingProxyType({"vertex_location": self.explicit_vertex_ai_location(litellm_params.model_dump())}) + def validate_environment( self, api_key: str | None,