fix: updating BaseAgent.clone() and LlmAgent.clone() to properly clone fields that are lists

PiperOrigin-RevId: 797855214
This commit is contained in:
Google Team Member
2025-08-21 11:52:31 -07:00
committed by Copybara-Service
parent 167182be01
commit 29bb75f975
2 changed files with 199 additions and 0 deletions
+12
View File
@@ -181,6 +181,18 @@ class BaseAgent(BaseModel):
cloned_agent = self.model_copy(update=update)
# If any field is stored as list and not provided in the update, need to
# shallow copy it for the cloned agent to avoid sharing the same list object
# with the original agent.
for field_name in cloned_agent.__class__.model_fields:
if field_name == 'sub_agents':
continue
if update is not None and field_name in update:
continue
field = getattr(cloned_agent, field_name)
if isinstance(field, list):
setattr(cloned_agent, field_name, field.copy())
if update is None or 'sub_agents' not in update:
# If `sub_agents` is not provided in the update, need to recursively clone
# the sub-agents to avoid sharing the sub-agents with the original agent.
+187
View File
@@ -14,6 +14,10 @@
"""Testings for the clone functionality of agents."""
from typing import Any
from typing import cast
from typing import Iterable
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.parallel_agent import ParallelAgent
@@ -374,6 +378,189 @@ def test_clone_with_sub_agents_update():
assert original.sub_agents[1].name == "original_sub2"
def _check_lists_contain_same_contents(*lists: Iterable[list[Any]]) -> None:
"""Assert that all provided lists contain the same elements."""
if lists:
first_list = lists[0]
assert all(len(lst) == len(first_list) for lst in lists)
for idx, elem in enumerate(first_list):
assert all(lst[idx] is elem for lst in lists)
def test_clone_shallow_copies_lists():
"""Test that cloning shallow copies fields stored as lists."""
# Define the list fields
before_agent_callback = [lambda *args, **kwargs: None]
after_agent_callback = [lambda *args, **kwargs: None]
before_model_callback = [lambda *args, **kwargs: None]
after_model_callback = [lambda *args, **kwargs: None]
before_tool_callback = [lambda *args, **kwargs: None]
after_tool_callback = [lambda *args, **kwargs: None]
tools = [lambda *args, **kwargs: None]
# Create the original agent with list fields
original = LlmAgent(
name="original_agent",
description="Original agent",
before_agent_callback=before_agent_callback,
after_agent_callback=after_agent_callback,
before_model_callback=before_model_callback,
after_model_callback=after_model_callback,
before_tool_callback=before_tool_callback,
after_tool_callback=after_tool_callback,
tools=tools,
)
# Clone the agent
cloned = original.clone()
# Verify the lists are copied
assert original.before_agent_callback is not cloned.before_agent_callback
assert original.after_agent_callback is not cloned.after_agent_callback
assert original.before_model_callback is not cloned.before_model_callback
assert original.after_model_callback is not cloned.after_model_callback
assert original.before_tool_callback is not cloned.before_tool_callback
assert original.after_tool_callback is not cloned.after_tool_callback
assert original.tools is not cloned.tools
# Verify the list copies are shallow
_check_lists_contain_same_contents(
before_agent_callback,
original.before_agent_callback,
cloned.before_agent_callback,
)
_check_lists_contain_same_contents(
after_agent_callback,
original.after_agent_callback,
cloned.after_agent_callback,
)
_check_lists_contain_same_contents(
before_model_callback,
original.before_model_callback,
cloned.before_model_callback,
)
_check_lists_contain_same_contents(
after_model_callback,
original.after_model_callback,
cloned.after_model_callback,
)
_check_lists_contain_same_contents(
before_tool_callback,
original.before_tool_callback,
cloned.before_tool_callback,
)
_check_lists_contain_same_contents(
after_tool_callback,
original.after_tool_callback,
cloned.after_tool_callback,
)
_check_lists_contain_same_contents(tools, original.tools, cloned.tools)
def test_clone_shallow_copies_lists_with_sub_agents():
"""Test that cloning recursively shallow copies fields stored as lists."""
# Define the list fields for the sub-agent
before_agent_callback = [lambda *args, **kwargs: None]
after_agent_callback = [lambda *args, **kwargs: None]
before_model_callback = [lambda *args, **kwargs: None]
after_model_callback = [lambda *args, **kwargs: None]
before_tool_callback = [lambda *args, **kwargs: None]
after_tool_callback = [lambda *args, **kwargs: None]
tools = [lambda *args, **kwargs: None]
# Create the original sub-agent with list fields and the top-level agent
sub_agents = [
LlmAgent(
name="sub_agent",
description="Sub agent",
before_agent_callback=before_agent_callback,
after_agent_callback=after_agent_callback,
before_model_callback=before_model_callback,
after_model_callback=after_model_callback,
before_tool_callback=before_tool_callback,
after_tool_callback=after_tool_callback,
tools=tools,
)
]
original = LlmAgent(
name="original_agent",
description="Original agent",
sub_agents=sub_agents,
)
# Clone the top-level agent
cloned = original.clone()
# Verify the sub_agents list is copied for the top-level agent
assert original.sub_agents is not cloned.sub_agents
# Retrieve the sub-agent for the original and cloned top-level agent
original_sub_agent = cast(LlmAgent, original.sub_agents[0])
cloned_sub_agent = cast(LlmAgent, cloned.sub_agents[0])
# Verify the lists are copied for the sub-agent
assert (
original_sub_agent.before_agent_callback
is not cloned_sub_agent.before_agent_callback
)
assert (
original_sub_agent.after_agent_callback
is not cloned_sub_agent.after_agent_callback
)
assert (
original_sub_agent.before_model_callback
is not cloned_sub_agent.before_model_callback
)
assert (
original_sub_agent.after_model_callback
is not cloned_sub_agent.after_model_callback
)
assert (
original_sub_agent.before_tool_callback
is not cloned_sub_agent.before_tool_callback
)
assert (
original_sub_agent.after_tool_callback
is not cloned_sub_agent.after_tool_callback
)
assert original_sub_agent.tools is not cloned_sub_agent.tools
# Verify the list copies are shallow for the sub-agent
_check_lists_contain_same_contents(
before_agent_callback,
original_sub_agent.before_agent_callback,
cloned_sub_agent.before_agent_callback,
)
_check_lists_contain_same_contents(
after_agent_callback,
original_sub_agent.after_agent_callback,
cloned_sub_agent.after_agent_callback,
)
_check_lists_contain_same_contents(
before_model_callback,
original_sub_agent.before_model_callback,
cloned_sub_agent.before_model_callback,
)
_check_lists_contain_same_contents(
after_model_callback,
original_sub_agent.after_model_callback,
cloned_sub_agent.after_model_callback,
)
_check_lists_contain_same_contents(
before_tool_callback,
original_sub_agent.before_tool_callback,
cloned_sub_agent.before_tool_callback,
)
_check_lists_contain_same_contents(
after_tool_callback,
original_sub_agent.after_tool_callback,
cloned_sub_agent.after_tool_callback,
)
_check_lists_contain_same_contents(
tools, original_sub_agent.tools, cloned_sub_agent.tools
)
if __name__ == "__main__":
# Run a specific test for debugging
test_three_level_nested_agent()