Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions src/google/adk/models/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@

_LazyEntry = tuple[str, str]
_llm_registry_dict: dict[str, Union[type['BaseLlm'], _LazyEntry]] = {}
_deprecated_model_aliases = {
'gemini-flash-latest': 'gemini-2.5-flash',
}


class LLMRegistry:
Expand All @@ -48,6 +51,7 @@ def new_llm(model: str) -> BaseLlm:
"""

prefix, actual_model = LLMRegistry._parse_model(model)
LLMRegistry._warn_if_deprecated_alias(actual_model)
cls = LLMRegistry.resolve(model)

if prefix and LLMRegistry._match_prefix(prefix, cls.__name__):
Expand All @@ -67,6 +71,21 @@ def _parse_model(model: str) -> tuple[str | None, str]:
return prefix, actual_model
return None, model

@staticmethod
def _warn_if_deprecated_alias(model: str) -> None:
replacement = _deprecated_model_aliases.get(model)
if replacement is None:
return

logger.warning(
(
'Model alias %r is deprecated and may resolve to an unavailable'
' Gemini model. Use %r or another versioned model id instead.'
),
model,
replacement,
)

@staticmethod
def _match_prefix(prefix: str, class_name: str) -> bool:
"""Checks if a prefix matches a class name."""
Expand Down
16 changes: 16 additions & 0 deletions tests/unittests/models/test_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,22 @@ def test_new_llm_with_prefix(mocker):
mock_class.assert_called_once_with(model='gpt-4')


@pytest.mark.parametrize(
'model_name',
['gemini-flash-latest', 'gemini:gemini-flash-latest'],
)
def test_new_llm_warns_for_deprecated_model_alias(caplog, model_name):
"""Test that deprecated Gemini aliases surface a useful warning."""
caplog.set_level('WARNING', logger='google_adk.google.adk.models.registry')

llm = models.LLMRegistry.new_llm(model_name)

assert isinstance(llm, Gemini)
assert llm.model == 'gemini-flash-latest'
assert "Model alias 'gemini-flash-latest' is deprecated" in caplog.text
assert "'gemini-2.5-flash'" in caplog.text


def test_new_llm_with_non_matching_prefix(mocker):
"""Test that new_llm keeps prefix if it does not match class."""
mock_class = mocker.MagicMock()
Expand Down