Skip to content

Commit 8dff850

Browse files
google-genai-botcopybara-github
authored andcommitted
fix: Support models slash prefix in model name extraction
Support "models/" prefix in model name extraction PiperOrigin-RevId: 827557443
1 parent 92a7d19 commit 8dff850

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

src/google/adk/utils/model_name_utils.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ def extract_model_name(model_string: str) -> str:
2727
"""Extract the actual model name from either simple or path-based format.
2828
2929
Args:
30-
model_string: Either a simple model name like "gemini-2.5-pro" or
31-
a path-based model name like "projects/.../models/gemini-2.0-flash-001"
30+
model_string: Either a simple model name like "gemini-2.5-pro" or a
31+
path-based model name like "projects/.../models/gemini-2.0-flash-001"
3232
3333
Returns:
3434
The extracted model name (e.g., "gemini-2.5-pro")
@@ -41,6 +41,10 @@ def extract_model_name(model_string: str) -> str:
4141
if match:
4242
return match.group(1)
4343

44+
# Handle 'models/' prefixed names like "models/gemini-2.5-pro"
45+
if model_string.startswith('models/'):
46+
return model_string[len('models/') :]
47+
4448
# If it's not a path-based model, return as-is (simple model name)
4549
return model_string
4650

tests/unittests/utils/test_model_name_utils.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ def test_extract_model_name_path_based_model(self):
4242
path_model_3 = 'projects/test-project/locations/europe-west1/publishers/google/models/claude-3-sonnet'
4343
assert extract_model_name(path_model_3) == 'claude-3-sonnet'
4444

45+
def test_extract_model_name_with_models_prefix(self):
46+
"""Test extraction of model names with 'models/' prefix."""
47+
assert extract_model_name('models/gemini-2.5-pro') == 'gemini-2.5-pro'
48+
assert extract_model_name('models/gemini-1.5-flash') == 'gemini-1.5-flash'
49+
4550
def test_extract_model_name_invalid_path(self):
4651
"""Test that invalid path formats return the original string."""
4752
invalid_paths = [

0 commit comments

Comments
 (0)