From ffe2bdbe4c2ea86cc7924eb36e8e3bb5528c0016 Mon Sep 17 00:00:00 2001
From: Calvin Giles <calvin.giles@trademe.co.nz>
Date: Sat, 19 Jul 2025 09:04:41 -0700
Subject: [PATCH] feat: Add support for vertex gemini model optimizer

Merge https://github.com/google/adk-python/pull/1130

This enables the use of the `model-optimizer-*` family of models in vertex, as per the [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/vertex-ai-model-optimizer#using-vertex-ai-model-optimizer).

To use this, ensure your location is set to `global` and pass a model optimizer model to an agent:

```python
root_agent = Agent(
    model="model-optimizer-exp-04-09",
    name="fast_and_slow_agent",
    instruction="Answer any question the user gives you - easy or hard.",
    generate_content_config=types.GenerateContentConfig(
        temperature=0.01,
        model_selection_config=ModelSelectionConfig(
            feature_selection_preference=FeatureSelectionPreference.BALANCED
            # Options: PRIORITIZE_QUALITY, BALANCED, PRIORITIZE_COST
        )
    ),
)
```
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/1130 from calvingiles:feat-model-optimizer 1a76bfa22420edb07d83415dcea6dd0114084e8e
PiperOrigin-RevId: 784921913
---
 src/google/adk/models/google_llm.py       | 2 ++
 tests/unittests/models/test_google_llm.py | 7 ++++---
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/google/adk/models/google_llm.py b/src/google/adk/models/google_llm.py
index 48547c7a..93d82317 100644
--- a/src/google/adk/models/google_llm.py
+++ b/src/google/adk/models/google_llm.py
@@ -68,6 +68,8 @@ class Gemini(BaseLlm):
 
     return [
         r'gemini-.*',
+        # model optimizer pattern
+        r'model-optimizer-.*',
         # fine-tuned vertex endpoint pattern
         r'projects\/.+\/locations\/.+\/endpoints\/.+',
         # vertex gemini long name
diff --git a/tests/unittests/models/test_google_llm.py b/tests/unittests/models/test_google_llm.py
index fa01dadd..f757b14d 100644
--- a/tests/unittests/models/test_google_llm.py
+++ b/tests/unittests/models/test_google_llm.py
@@ -75,11 +75,12 @@ def mock_os_environ():
 
 def test_supported_models():
   models = Gemini.supported_models()
-  assert len(models) == 3
+  assert len(models) == 4
   assert models[0] == r"gemini-.*"
-  assert models[1] == r"projects\/.+\/locations\/.+\/endpoints\/.+"
+  assert models[1] == r"model-optimizer-.*"
+  assert models[2] == r"projects\/.+\/locations\/.+\/endpoints\/.+"
   assert (
-      models[2]
+      models[3]
       == r"projects\/.+\/locations\/.+\/publishers\/google\/models\/gemini.+"
   )