fix: Support models slash prefix in model name extraction

Support "models/" prefix in model name extraction

PiperOrigin-RevId: 827557443
This commit is contained in:
Google Team Member
2025-11-03 10:32:49 -08:00
committed by Copybara-Service
parent 92a7d19573
commit 8dff85099d
2 changed files with 11 additions and 2 deletions
+6 -2
View File
@@ -27,8 +27,8 @@ def extract_model_name(model_string: str) -> str:
"""Extract the actual model name from either simple or path-based format.
Args:
model_string: Either a simple model name like "gemini-2.5-pro" or
a path-based model name like "projects/.../models/gemini-2.0-flash-001"
model_string: Either a simple model name like "gemini-2.5-pro" or a
path-based model name like "projects/.../models/gemini-2.0-flash-001"
Returns:
The extracted model name (e.g., "gemini-2.5-pro")
@@ -41,6 +41,10 @@ def extract_model_name(model_string: str) -> str:
if match:
return match.group(1)
# Handle 'models/' prefixed names like "models/gemini-2.5-pro"
if model_string.startswith('models/'):
return model_string[len('models/') :]
# If it's not a path-based model, return as-is (simple model name)
return model_string