File tree Expand file tree Collapse file tree 2 files changed +11
-2
lines changed
Expand file tree Collapse file tree 2 files changed +11
-2
lines changed Original file line number Diff line number Diff 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
Original file line number Diff line number Diff 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 = [
You can’t perform that action at this time.
0 commit comments