chore: import upstream snapshot with attribution

tests/unittests/agents/test_llm_agent_fields.py

@@ -0,0 +1,605 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for canonical_xxx fields in LlmAgent."""
+
+import logging
+from typing import Any
+from typing import Optional
+from unittest import mock
+
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.models.anthropic_llm import Claude
+from google.adk.models.google_llm import Gemini
+from google.adk.models.lite_llm import LiteLlm
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.registry import LLMRegistry
+from google.adk.planners.built_in_planner import BuiltInPlanner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.google_search_tool import google_search
+from google.adk.tools.google_search_tool import GoogleSearchTool
+from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
+from google.genai import types
+from pydantic import BaseModel
+import pytest
+
+
+async def _create_readonly_context(
+    agent: LlmAgent, state: Optional[dict[str, Any]] = None
+) -> ReadonlyContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user', state=state
+  )
+  invocation_context = InvocationContext(
+      invocation_id='test_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  return ReadonlyContext(invocation_context)
+
+
+@pytest.mark.parametrize(
+    ('default_model', 'expected_model_name', 'expected_model_type'),
+    [
+        (LlmAgent.DEFAULT_MODEL, LlmAgent.DEFAULT_MODEL, Gemini),
+        ('gemini-2.5-flash', 'gemini-2.5-flash', Gemini),
+    ],
+)
+def test_canonical_model_default_fallback(
+    default_model, expected_model_name, expected_model_type
+):
+  original_default = LlmAgent._default_model
+  LlmAgent.set_default_model(default_model)
+  try:
+    agent = LlmAgent(name='test_agent')
+    assert isinstance(agent.canonical_model, expected_model_type)
+    assert agent.canonical_model.model == expected_model_name
+  finally:
+    LlmAgent.set_default_model(original_default)
+
+
+def test_canonical_model_str():
+  agent = LlmAgent(name='test_agent', model='gemini-pro')
+
+  assert agent.canonical_model.model == 'gemini-pro'
+
+
+def test_canonical_model_llm():
+  llm = LLMRegistry.new_llm('gemini-pro')
+  agent = LlmAgent(name='test_agent', model=llm)
+
+  assert agent.canonical_model == llm
+
+
+def test_canonical_model_inherit():
+  sub_agent = LlmAgent(name='sub_agent')
+  parent_agent = LlmAgent(
+      name='parent_agent', model='gemini-pro', sub_agents=[sub_agent]
+  )
+
+  assert sub_agent.canonical_model == parent_agent.canonical_model
+
+
+def test_canonical_live_model_default_fallback():
+  original_default = LlmAgent._default_live_model
+  LlmAgent.set_default_live_model('gemini-2.0-flash')
+  try:
+    agent = LlmAgent(name='test_agent')
+    assert agent.canonical_live_model.model == 'gemini-2.0-flash'
+  finally:
+    LlmAgent.set_default_live_model(original_default)
+
+
+def test_canonical_live_model_str():
+  agent = LlmAgent(name='test_agent', model='gemini-pro')
+  assert agent.canonical_live_model.model == 'gemini-pro'
+
+
+def test_canonical_live_model_llm():
+  llm = LLMRegistry.new_llm('gemini-pro')
+  agent = LlmAgent(name='test_agent', model=llm)
+  assert agent.canonical_live_model == llm
+
+
+def test_canonical_live_model_inherit():
+  sub_agent = LlmAgent(name='sub_agent')
+  parent_agent = LlmAgent(
+      name='parent_agent', model='gemini-pro', sub_agents=[sub_agent]
+  )
+  assert sub_agent.canonical_live_model == parent_agent.canonical_live_model
+
+
+async def test_canonical_instruction_str():
+  agent = LlmAgent(name='test_agent', instruction='instruction')
+  ctx = await _create_readonly_context(agent)
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_instruction(ctx)
+  )
+  assert canonical_instruction == 'instruction'
+  assert not bypass_state_injection
+
+
+async def test_canonical_instruction():
+  def _instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_instruction(ctx)
+  )
+  assert canonical_instruction == 'instruction: state_value'
+  assert bypass_state_injection
+
+
+async def test_async_canonical_instruction():
+  async def _instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_instruction(ctx)
+  )
+  assert canonical_instruction == 'instruction: state_value'
+  assert bypass_state_injection
+
+
+async def test_canonical_global_instruction_str():
+  agent = LlmAgent(name='test_agent', global_instruction='global instruction')
+  ctx = await _create_readonly_context(agent)
+
+  canonical_instruction, bypass_state_injection = (
+      await agent.canonical_global_instruction(ctx)
+  )
+  assert canonical_instruction == 'global instruction'
+  assert not bypass_state_injection
+
+
+async def test_canonical_global_instruction():
+  def _global_instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'global instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(
+      name='test_agent', global_instruction=_global_instruction_provider
+  )
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+
+  canonical_global_instruction, bypass_state_injection = (
+      await agent.canonical_global_instruction(ctx)
+  )
+  assert canonical_global_instruction == 'global instruction: state_value'
+  assert bypass_state_injection
+
+
+async def test_async_canonical_global_instruction():
+  async def _global_instruction_provider(ctx: ReadonlyContext) -> str:
+    return f'global instruction: {ctx.state["state_var"]}'
+
+  agent = LlmAgent(
+      name='test_agent', global_instruction=_global_instruction_provider
+  )
+  ctx = await _create_readonly_context(
+      agent, state={'state_var': 'state_value'}
+  )
+  canonical_global_instruction, bypass_state_injection = (
+      await agent.canonical_global_instruction(ctx)
+  )
+  assert canonical_global_instruction == 'global instruction: state_value'
+  assert bypass_state_injection
+
+
+def test_output_schema_with_sub_agents_will_not_throw():
+  class Schema(BaseModel):
+    pass
+
+  sub_agent = LlmAgent(
+      name='sub_agent',
+  )
+
+  agent = LlmAgent(
+      name='test_agent',
+      output_schema=Schema,
+      sub_agents=[sub_agent],
+  )
+
+  # Transfer is not disabled
+  assert not agent.disallow_transfer_to_parent
+  assert not agent.disallow_transfer_to_peers
+
+  assert agent.output_schema == Schema
+  assert agent.sub_agents == [sub_agent]
+
+
+def test_output_schema_with_tools_will_not_throw():
+  class Schema(BaseModel):
+    pass
+
+  def _a_tool():
+    pass
+
+  LlmAgent(
+      name='test_agent',
+      output_schema=Schema,
+      tools=[_a_tool],
+  )
+
+
+def test_before_model_callback():
+  def _before_model_callback(
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+  ) -> None:
+    return None
+
+  agent = LlmAgent(
+      name='test_agent', before_model_callback=_before_model_callback
+  )
+
+  # TODO: add more logic assertions later.
+  assert agent.before_model_callback is not None
+
+
+def test_validate_generate_content_config_thinking_config_allow():
+  """Tests that thinking_config is now allowed directly in the agent init."""
+  agent = LlmAgent(
+      name='test_agent',
+      generate_content_config=types.GenerateContentConfig(
+          thinking_config=types.ThinkingConfig(include_thoughts=True)
+      ),
+  )
+  assert agent.generate_content_config.thinking_config.include_thoughts is True
+
+
+def test_thinking_config_precedence_warning():
+  """Tests that a UserWarning is issued when both manual config and planner exist."""
+
+  config = types.GenerateContentConfig(
+      thinking_config=types.ThinkingConfig(include_thoughts=True)
+  )
+  planner = BuiltInPlanner(
+      thinking_config=types.ThinkingConfig(include_thoughts=True)
+  )
+
+  with pytest.warns(
+      UserWarning, match="planner's configuration will take precedence"
+  ):
+    LlmAgent(name='test_agent', generate_content_config=config, planner=planner)
+
+
+def test_validate_generate_content_config_tools_throw():
+  """Tests that tools cannot be set directly in config."""
+  with pytest.raises(ValueError):
+    _ = LlmAgent(
+        name='test_agent',
+        generate_content_config=types.GenerateContentConfig(
+            tools=[types.Tool(function_declarations=[])]
+        ),
+    )
+
+
+def test_validate_generate_content_config_system_instruction_throw():
+  """Tests that system instructions cannot be set directly in config."""
+  with pytest.raises(ValueError):
+    _ = LlmAgent(
+        name='test_agent',
+        generate_content_config=types.GenerateContentConfig(
+            system_instruction='system instruction'
+        ),
+    )
+
+
+def test_validate_generate_content_config_response_schema_throw():
+  """Tests that response schema cannot be set directly in config."""
+
+  class Schema(BaseModel):
+    pass
+
+  with pytest.raises(ValueError):
+    _ = LlmAgent(
+        name='test_agent',
+        generate_content_config=types.GenerateContentConfig(
+            response_schema=Schema
+        ),
+    )
+
+
+def test_allow_transfer_by_default():
+  sub_agent = LlmAgent(name='sub_agent')
+  agent = LlmAgent(name='test_agent', sub_agents=[sub_agent])
+
+  assert not agent.disallow_transfer_to_parent
+  assert not agent.disallow_transfer_to_peers
+
+
+# TODO(b/448114567): Remove TestCanonicalTools once the workaround
+# is no longer needed.
+class TestCanonicalTools:
+  """Unit tests for canonical_tools in LlmAgent."""
+
+  @staticmethod
+  def _my_tool(sides: int) -> int:
+    return sides
+
+  async def test_handle_google_search_with_other_tools(self):
+    """Test that google_search is wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            GoogleSearchTool(bypass_multi_tools_limit=True),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'google_search_agent'
+    assert tools[1].__class__.__name__ == 'GoogleSearchAgentTool'
+
+  async def test_handle_google_search_with_other_tools_no_bypass(self):
+    """Test that google_search is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            GoogleSearchTool(bypass_multi_tools_limit=False),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'google_search'
+    assert tools[1].__class__.__name__ == 'GoogleSearchTool'
+
+  async def test_handle_google_search_only(self):
+    """Test that google_search is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            google_search,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == 'google_search'
+    assert tools[0].__class__.__name__ == 'GoogleSearchTool'
+
+  async def test_function_tool_only(self):
+    """Test that function tool is not affected."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+
+  @mock.patch(
+      'google.auth.default',
+      mock.MagicMock(return_value=('credentials', 'project')),
+  )
+  async def test_handle_vais_with_other_tools(self):
+    """Test that VertexAiSearchTool is replaced with Discovery Engine Search."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            VertexAiSearchTool(
+                data_store_id='test_data_store_id',
+                bypass_multi_tools_limit=True,
+            ),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'discovery_engine_search'
+    assert tools[1].__class__.__name__ == 'DiscoveryEngineSearchTool'
+
+  async def test_handle_vais_with_other_tools_no_bypass(self):
+    """Test that VertexAiSearchTool is not replaced."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            VertexAiSearchTool(
+                data_store_id='test_data_store_id',
+                bypass_multi_tools_limit=False,
+            ),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'
+    assert tools[1].name == 'vertex_ai_search'
+    assert tools[1].__class__.__name__ == 'VertexAiSearchTool'
+
+  async def test_handle_vais_only(self):
+    """Test that VertexAiSearchTool is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            VertexAiSearchTool(data_store_id='test_data_store_id'),
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == 'vertex_ai_search'
+    assert tools[0].__class__.__name__ == 'VertexAiSearchTool'
+
+  async def test_multiple_tools_resolution(self):
+    """Test that multiple tools are resolved correctly."""
+
+    def _tool_1():
+      pass
+
+    def _tool_2():
+      pass
+
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[_tool_1, _tool_2],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_tool_1'
+    assert tools[1].name == '_tool_2'
+
+  async def test_canonical_tools_graceful_degradation_on_toolset_error(self):
+    """Test that canonical_tools returns tools from working toolsets when one fails."""
+    from google.adk.tools.base_tool import BaseTool
+    from google.adk.tools.base_toolset import BaseToolset
+
+    class FailingToolset(BaseToolset):
+
+      async def get_tools(self, readonly_context=None):
+        raise ConnectionError('MCP server unavailable')
+
+    class WorkingToolset(BaseToolset):
+
+      async def get_tools(self, readonly_context=None):
+        tool = mock.MagicMock(spec=BaseTool)
+        tool.name = 'working_tool'
+        tool._get_declaration = mock.MagicMock(return_value=None)
+        return [tool]
+
+    def _regular_tool():
+      pass
+
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[_regular_tool, FailingToolset(), WorkingToolset()],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    # Should have the regular tool + working toolset tool, but not crash
+    assert len(tools) == 2
+    assert tools[0].name == '_regular_tool'
+    assert tools[1].name == 'working_tool'
+
+
+# Tests for multi-provider model support via string model names
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'gemini-2.5-flash',
+        'gemini-2.5-pro',
+    ],
+)
+def test_agent_with_gemini_string_model(model_name):
+  """Test that Agent accepts Gemini model strings and resolves to Gemini."""
+  agent = LlmAgent(name='test_agent', model=model_name)
+  assert isinstance(agent.canonical_model, Gemini)
+  assert agent.canonical_model.model == model_name
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'claude-3-5-sonnet-v2@20241022',
+        'claude-sonnet-4@20250514',
+    ],
+)
+def test_agent_with_claude_string_model(model_name):
+  """Test that Agent accepts Claude model strings and resolves to Claude."""
+  agent = LlmAgent(name='test_agent', model=model_name)
+  assert isinstance(agent.canonical_model, Claude)
+  assert agent.canonical_model.model == model_name
+
+
+@pytest.mark.parametrize(
+    'model_name',
+    [
+        'openai/gpt-4o',
+        'groq/llama3-70b-8192',
+        'anthropic/claude-3-opus-20240229',
+    ],
+)
+def test_agent_with_litellm_string_model(model_name):
+  """Test that Agent accepts LiteLLM provider strings."""
+  agent = LlmAgent(name='test_agent', model=model_name)
+  assert isinstance(agent.canonical_model, LiteLlm)
+  assert agent.canonical_model.model == model_name
+
+
+def test_builtin_planner_overwrite_logging(caplog):
+  """Tests that the planner logs an DEBUG message when overwriting a config."""
+
+  planner = BuiltInPlanner(
+      thinking_config=types.ThinkingConfig(include_thoughts=True)
+  )
+
+  # Create a request that already has a thinking_config
+  req = LlmRequest(
+      contents=[],
+      config=types.GenerateContentConfig(
+          thinking_config=types.ThinkingConfig(include_thoughts=True)
+      ),
+  )
+
+  with caplog.at_level(
+      logging.DEBUG, logger='google_adk.google.adk.planners.built_in_planner'
+  ):
+    planner.apply_thinking_config(req)
+  assert (
+      'Overwriting `thinking_config` from `generate_content_config`'
+      in caplog.text
+  )