diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index c1822e4720d..3daf43c729e 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -2,6 +2,7 @@ import asyncio import contextvars from collections.abc import Iterator from functools import partial +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final import httpx @@ -192,12 +193,28 @@ class GenerateContentHelper: if litellm_logging_obj is None: raise ValueError("litellm_logging_obj is required, but got None") + # The configured Vertex location decides the regional pricing uplift. The + # completion and anthropic_messages paths both record it here; without it + # the cost calculator falls back to a default region and prices a `global` + # model as us-central1, so the same model costs 10% more on this route. + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + explicit_vertex_location: Final = VertexBase.explicit_vertex_ai_location( + MappingProxyType(litellm_params.model_dump(exclude_none=True)) + ) + vertex_location_params: Final = ( + MappingProxyType({"vertex_location": explicit_vertex_location}) + if explicit_vertex_location + else MappingProxyType({}) + ) + litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, optional_params=dict(generate_content_config_dict), litellm_params={ "litellm_call_id": litellm_call_id, + **vertex_location_params, }, 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..eef7ae02a4b 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -259,3 +259,73 @@ async def test_native_fields_forwarded_on_async_stream(): body = mock_post.call_args.kwargs["json"] assert body["safetySettings"] == safety_settings assert "safetySettings" not in body.get("generationConfig", {}) + + +def _generate_content_logging_obj(call_id: str): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="gemini-2.0-flash", + messages=[], + stream=False, + call_type="generate_content", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + ) + + +@pytest.mark.parametrize("configured_location", ["global", "us-east5"]) +def test_vertex_location_recorded_for_cost_calculation(configured_location): + """ + Regression for https://github.com/BerriAI/litellm/issues/40692 + + The Vertex location decides the regional pricing uplift. The completion and + anthropic_messages paths record it on the logging object so the cost + calculator can price against the configured region. generate_content did + not, so a model configured `vertex_location: global` was priced as + us-central1 and cost 10% more than the same model over /v1/chat/completions. + """ + from litellm.google_genai.main import GenerateContentHelper + + logging_obj = _generate_content_logging_obj("vertex-location-test") + + GenerateContentHelper.setup_generate_content_call( + model="vertex_ai/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "say ok"}]}], + custom_llm_provider="vertex_ai", + litellm_logging_obj=logging_obj, + litellm_call_id="vertex-location-test", + vertex_project="test-project", + vertex_location=configured_location, + ) + + recorded = logging_obj.model_call_details["litellm_params"] + assert recorded.get("vertex_location") == configured_location, ( + "the configured vertex_location must reach the cost calculator; " + f"got {recorded.get('vertex_location')!r}" + ) + + +def test_vertex_location_absent_when_not_configured(): + """ + Nothing configured means nothing recorded, so the cost calculator keeps its + own fallback rather than being handed an empty value here. + """ + from litellm.google_genai.main import GenerateContentHelper + + logging_obj = _generate_content_logging_obj("vertex-location-unset") + + GenerateContentHelper.setup_generate_content_call( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "say ok"}]}], + custom_llm_provider="gemini", + litellm_logging_obj=logging_obj, + litellm_call_id="vertex-location-unset", + api_key="test-key", + ) + + recorded = logging_obj.model_call_details["litellm_params"] + assert "vertex_location" not in recorded