Skip to content

Commit ffe2bdb

Browse files
calvingilescopybara-github
authored andcommitted
feat: Add support for vertex gemini model optimizer
Merge #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=#1130 from calvingiles:feat-model-optimizer 1a76bfa PiperOrigin-RevId: 784921913
1 parent 67284fc commit ffe2bdb

File tree

2 files changed

+6
-3
lines changed

2 files changed

+6
-3
lines changed

src/google/adk/models/google_llm.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@ def supported_models() -> list[str]:
6868

6969
return [
7070
r'gemini-.*',
71+
# model optimizer pattern
72+
r'model-optimizer-.*',
7173
# fine-tuned vertex endpoint pattern
7274
r'projects\/.+\/locations\/.+\/endpoints\/.+',
7375
# vertex gemini long name

tests/unittests/models/test_google_llm.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,11 +75,12 @@ def mock_os_environ():
7575

7676
def test_supported_models():
7777
models = Gemini.supported_models()
78-
assert len(models) == 3
78+
assert len(models) == 4
7979
assert models[0] == r"gemini-.*"
80-
assert models[1] == r"projects\/.+\/locations\/.+\/endpoints\/.+"
80+
assert models[1] == r"model-optimizer-.*"
81+
assert models[2] == r"projects\/.+\/locations\/.+\/endpoints\/.+"
8182
assert (
82-
models[2]
83+
models[3]
8384
== r"projects\/.+\/locations\/.+\/publishers\/google\/models\/gemini.+"
8485
)
8586

0 commit comments

Comments
 (0)