Merge 38215a8a81 into 1c61c2606e

litellm/google_genai/main.py

@@ -198,6 +198,7 @@ class GenerateContentHelper:
             optional_params=dict(generate_content_config_dict),
             litellm_params={
                 "litellm_call_id": litellm_call_id,
+                **generate_content_provider_config.get_generate_content_logging_params(litellm_params),
             },
             custom_llm_provider=custom_llm_provider,
         )

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,

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,

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