Skip to content

fix(agents): safely decouple model memoization to prevent Agent Engine crashes #6780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
manideep-malyala wants to merge 1 commit into
base: main
Choose a base branch
from
+2 −57
Open

fix(agents): safely decouple model memoization to prevent Agent Engine crashes #6780

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
21 changes: 2 additions & 19 deletions src/google/adk/agents/llm_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@
from pydantic import Field
from pydantic import field_validator
from pydantic import model_validator
from pydantic import PrivateAttr
from typing_extensions import override
from typing_extensions import TypeAlias

Expand Down Expand Up @@ -256,14 +255,6 @@ class LlmAgent(BaseAgent, abc.ABC):
LlmAgent.set_default_model. The built-in default is gemini-3.5-flash.
"""

_resolved_model: Optional[tuple[str, BaseLlm]] = PrivateAttr(default=None)
"""The model name last resolved by canonical_model, with its BaseLlm."""

_resolved_live_model: Optional[tuple[str, BaseLlm]] = PrivateAttr(
default=None
)
"""The model name last resolved by canonical_live_model, with its BaseLlm."""

config_type: ClassVar[Type[BaseAgentConfig]] = LlmAgentConfig
"""The config type for this agent.

Expand Down Expand Up @@ -640,11 +631,7 @@ def canonical_model(self) -> BaseLlm:
if isinstance(self.model, BaseLlm):
return self.model
elif self.model: # model is non-empty str
resolved = self._resolved_model
if resolved is None or resolved[0] != self.model:
resolved = (self.model, LLMRegistry.new_llm(self.model))
self._resolved_model = resolved
return resolved[1]
return LLMRegistry.new_llm(self.model)
else: # find model from ancestors.
ancestor_agent = self.parent_agent
while ancestor_agent is not None:
Expand All @@ -662,11 +649,7 @@ def canonical_live_model(self) -> BaseLlm:
if isinstance(self.model, BaseLlm):
return self.model
elif self.model: # model is non-empty str
resolved = self._resolved_live_model
if resolved is None or resolved[0] != self.model:
resolved = (self.model, LLMRegistry.new_llm(self.model))
self._resolved_live_model = resolved
return resolved[1]
return LLMRegistry.new_llm(self.model)
else: # find model from ancestors.
ancestor_agent = self.parent_agent
while ancestor_agent is not None:
Expand Down
38 changes: 0 additions & 38 deletions tests/unittests/agents/test_llm_agent_fields.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,31 +100,6 @@ def test_canonical_model_inherit():
assert sub_agent.canonical_model == parent_agent.canonical_model


def test_canonical_model_str_resolved_once():
agent = LlmAgent(name='test_agent', model='gemini-pro')

with mock.patch.object(
LLMRegistry, 'new_llm', wraps=LLMRegistry.new_llm
) as new_llm:
first = agent.canonical_model
second = agent.canonical_model
third = agent.canonical_model

assert new_llm.call_count == 1
assert first is second is third


def test_canonical_model_str_resolved_again_after_reassignment():
agent = LlmAgent(name='test_agent', model='gemini-pro')
first = agent.canonical_model

agent.model = 'gemini-2.5-flash'
second = agent.canonical_model

assert second is not first
assert second.model == 'gemini-2.5-flash'


def test_canonical_model_str_not_stale_after_model_copy():
agent = LlmAgent(name='test_agent', model='gemini-pro')
assert agent.canonical_model.model == 'gemini-pro'
Expand All @@ -135,19 +110,6 @@ def test_canonical_model_str_not_stale_after_model_copy():
assert agent.canonical_model.model == 'gemini-pro'


def test_canonical_live_model_str_resolved_once():
agent = LlmAgent(name='test_agent', model='gemini-pro')

with mock.patch.object(
LLMRegistry, 'new_llm', wraps=LLMRegistry.new_llm
) as new_llm:
first = agent.canonical_live_model
second = agent.canonical_live_model

assert new_llm.call_count == 1
assert first is second


def test_canonical_live_model_default_fallback():
original_default = LlmAgent._default_live_model
LlmAgent.set_default_live_model('gemini-2.0-flash')
Expand Down