Skip to content

Commit 1410b3f

Browse files
committed
feat(models): allow configuring Vertex AI API version (v1/v1beta1)
The Gemini model integration relied on the google-genai SDK's default API version, which is v1beta1 for the Vertex AI backend. Production users that need a stable, SLA-eligible endpoint had no supported way to opt into the GA v1 Vertex AI API short of subclassing Gemini and overriding api_client. Add an optional api_version field on Gemini and honor the GOOGLE_GENAI_API_VERSION environment variable as a fallback. The resolved version flows through the existing _base_url_and_api_version chokepoint, so it applies consistently to the standard api_client, the live API client, and the per-request http_options patching. Resolution order: a version embedded in base_url wins, then the api_version field, then GOOGLE_GENAI_API_VERSION, then the SDK default. When none is configured the behavior is unchanged, so existing users are unaffected. Fixes #3246.
1 parent 6bc9c9f commit 1410b3f

2 files changed

Lines changed: 150 additions & 1 deletion

File tree

src/google/adk/models/google_llm.py

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import copy
2020
from functools import cached_property
2121
import logging
22+
import os
2223
import re
2324
from typing import Any
2425
from typing import AsyncGenerator
@@ -54,6 +55,7 @@
5455
_NEW_LINE = '\n'
5556
_EXCLUDED_PART_FIELD = {'inline_data': {'data'}}
5657
_GOOGLE_API_VERSION_SUFFIX_PATTERN = re.compile(r'/?(v[0-9][a-z0-9.-]*)/?')
58+
_API_VERSION_ENV_VARIABLE_NAME = 'GOOGLE_GENAI_API_VERSION'
5759

5860

5961
_RESOURCE_EXHAUSTED_POSSIBLE_FIX_MESSAGE = """
@@ -123,6 +125,27 @@ def api_client(self) -> Client:
123125
base_url: Optional[str] = None
124126
"""The base URL for the AI platform service endpoint."""
125127

128+
api_version: Optional[str] = None
129+
"""The API version to use for the AI platform service endpoint.
130+
131+
For the Vertex AI backend the google-genai SDK defaults to ``v1beta1``, which
132+
exposes the latest preview features. Production deployments that require a
133+
stable, SLA-eligible endpoint can set this to ``v1`` to use the GA Vertex AI
134+
API. When unset, the ``GOOGLE_GENAI_API_VERSION`` environment variable is
135+
consulted, and finally the SDK's own default is used so existing behavior is
136+
unchanged.
137+
138+
An API version embedded in ``base_url`` (e.g.
139+
``https://...googleapis.com/v1``) takes precedence over this field.
140+
141+
Sample:
142+
```python
143+
from google.adk.models import Gemini
144+
145+
agent = Agent(model=Gemini(model="gemini-2.5-pro", api_version="v1"))
146+
```
147+
"""
148+
126149
speech_config: Optional[types.SpeechConfig] = None
127150

128151
use_interactions_api: bool = False
@@ -371,9 +394,29 @@ def _api_backend(self) -> GoogleLLMVariant:
371394
def _tracking_headers(self) -> dict[str, str]:
372395
return get_tracking_headers()
373396

397+
def _configured_api_version(self) -> Optional[str]:
398+
"""Returns the explicitly configured API version, if any.
399+
400+
Resolution order:
401+
1. The ``api_version`` field set on this instance.
402+
2. The ``GOOGLE_GENAI_API_VERSION`` environment variable.
403+
404+
Returns ``None`` when neither is set, in which case the google-genai SDK's
405+
own default (``v1beta1`` for Vertex AI) applies, preserving existing
406+
behavior.
407+
"""
408+
if self.api_version:
409+
return self.api_version
410+
return os.environ.get(_API_VERSION_ENV_VARIABLE_NAME) or None
411+
374412
@cached_property
375413
def _base_url_and_api_version(self) -> tuple[Optional[str], Optional[str]]:
376-
return _normalize_base_url_and_api_version(self.base_url)
414+
base_url, api_version = _normalize_base_url_and_api_version(self.base_url)
415+
# A version embedded in the base URL wins; otherwise fall back to the
416+
# explicitly configured api_version (field or environment variable).
417+
if api_version is None:
418+
api_version = self._configured_api_version()
419+
return base_url, api_version
377420

378421
@cached_property
379422
def _live_api_version(self) -> str:

tests/unittests/models/test_google_llm.py

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,61 @@ def test_api_client_preserves_custom_base_url_path():
341341
assert client._api_client._http_options.api_version == "v1beta"
342342

343343

344+
def test_api_client_default_api_version_unchanged(monkeypatch):
345+
"""Without configuration, ADK does not force an api_version (SDK default)."""
346+
monkeypatch.delenv("GOOGLE_GENAI_API_VERSION", raising=False)
347+
model = Gemini(model="gemini-2.5-flash")
348+
349+
# ADK leaves api_version unset so the google-genai SDK applies its own
350+
# default (v1beta1 for Vertex AI), preserving existing behavior.
351+
assert model._base_url_and_api_version == (None, None)
352+
353+
354+
def test_api_client_uses_api_version_field():
355+
"""The api_version field flows into the constructed client's http_options."""
356+
model = Gemini(model="gemini-2.5-flash", api_version="v1")
357+
358+
client = model.api_client
359+
360+
assert client._api_client._http_options.api_version == "v1"
361+
362+
363+
def test_api_client_uses_api_version_env_var(monkeypatch):
364+
"""The GOOGLE_GENAI_API_VERSION env var flows into http_options."""
365+
monkeypatch.setenv("GOOGLE_GENAI_API_VERSION", "v1")
366+
model = Gemini(model="gemini-2.5-flash")
367+
368+
client = model.api_client
369+
370+
assert client._api_client._http_options.api_version == "v1"
371+
372+
373+
def test_api_version_field_overrides_env_var(monkeypatch):
374+
"""The explicit api_version field takes precedence over the env var."""
375+
monkeypatch.setenv("GOOGLE_GENAI_API_VERSION", "v1beta1")
376+
model = Gemini(model="gemini-2.5-flash", api_version="v1")
377+
378+
client = model.api_client
379+
380+
assert client._api_client._http_options.api_version == "v1"
381+
382+
383+
def test_base_url_api_version_overrides_field():
384+
"""A version embedded in base_url wins over the api_version field."""
385+
model = Gemini(
386+
model="gemini-2.5-flash",
387+
base_url="https://generativelanguage.googleapis.com/v1alpha",
388+
api_version="v1",
389+
)
390+
391+
client = model.api_client
392+
393+
assert client._api_client._http_options.base_url == (
394+
"https://generativelanguage.googleapis.com/"
395+
)
396+
assert client._api_client._http_options.api_version == "v1alpha"
397+
398+
344399
def test_maybe_append_user_content(gemini_llm, llm_request):
345400
# Test with user content already present
346401
gemini_llm._maybe_append_user_content(llm_request)
@@ -766,6 +821,35 @@ async def mock_coro():
766821
assert len(responses) == 2 if stream else 1
767822

768823

824+
@pytest.mark.asyncio
825+
async def test_generate_content_async_patches_api_version_from_field(
826+
llm_request, generate_content_response
827+
):
828+
"""The configured api_version field is patched onto the request config."""
829+
gemini_llm = Gemini(model="gemini-2.5-flash", api_version="v1")
830+
llm_request.config.http_options = types.HttpOptions(
831+
headers={"custom-header": "custom-value"}
832+
)
833+
834+
with mock.patch.object(gemini_llm, "api_client") as mock_client:
835+
836+
async def mock_coro():
837+
return generate_content_response
838+
839+
mock_client.aio.models.generate_content.return_value = mock_coro()
840+
841+
_ = [
842+
resp
843+
async for resp in gemini_llm.generate_content_async(
844+
llm_request, stream=False
845+
)
846+
]
847+
848+
call_args = mock_client.aio.models.generate_content.call_args
849+
final_config = call_args.kwargs["config"]
850+
assert final_config.http_options.api_version == "v1"
851+
852+
769853
def test_live_api_version_vertex_ai(gemini_llm):
770854
"""Test that _live_api_version returns 'v1beta1' for Vertex AI backend."""
771855
with mock.patch.object(
@@ -774,6 +858,28 @@ def test_live_api_version_vertex_ai(gemini_llm):
774858
assert gemini_llm._live_api_version == "v1beta1"
775859

776860

861+
def test_live_api_version_uses_configured_field():
862+
"""Test that _live_api_version honors the configured api_version field."""
863+
gemini_llm = Gemini(model="gemini-2.5-flash", api_version="v1")
864+
865+
with mock.patch.object(
866+
gemini_llm, "_api_backend", GoogleLLMVariant.VERTEX_AI
867+
):
868+
assert gemini_llm._live_api_version == "v1"
869+
870+
871+
def test_live_api_client_uses_configured_field():
872+
"""Test that _live_api_client http_options honors the api_version field."""
873+
gemini_llm = Gemini(model="gemini-2.5-flash", api_version="v1")
874+
875+
with mock.patch.object(
876+
gemini_llm, "_api_backend", GoogleLLMVariant.VERTEX_AI
877+
):
878+
client = gemini_llm._live_api_client
879+
880+
assert client._api_client._http_options.api_version == "v1"
881+
882+
777883
def test_live_api_version_uses_google_base_url_version():
778884
gemini_llm = Gemini(
779885
model="gemini-2.5-flash",

0 commit comments

Comments
 (0)