feat: allow thinking_config in generate_content_config (#4108)

src/google/adk/agents/llm_agent.py

@@ -285,7 +285,7 @@ class LlmAgent(BaseAgent):
  """The additional content generation configurations.
 
  NOTE: not all fields are usable, e.g. tools must be configured via `tools`,
- thinking_config must be configured via `planner` in LlmAgent.
+ thinking_config can be configured here or via the `planner`. If both are set, the planner's configuration takes precedence.
 
  For example: use this config to adjust model temperature, configure safety
  settings, etc.
@@ -849,8 +849,6 @@ def validate_generate_content_config(
  ) -> types.GenerateContentConfig:
  if not generate_content_config:
  return types.GenerateContentConfig()
- if generate_content_config.thinking_config:
- raise ValueError('Thinking config should be set via LlmAgent.planner.')
  if generate_content_config.tools:
  raise ValueError('All tools must be set via LlmAgent.tools.')
  if generate_content_config.system_instruction:
@@ -863,6 +861,23 @@ def validate_generate_content_config(
  )
  return generate_content_config
 
+ @override
+ def model_post_init(self, __context: Any) -> None:
+ """Provides a warning if multiple thinking configurations are found."""
+ super().model_post_init(__context)
+
+ # Note: Using getattr to check both locations for thinking_config
+ if getattr(
+ self.generate_content_config, 'thinking_config', None
+ ) and getattr(self.planner, 'thinking_config', None):
+ warnings.warn(
+ 'Both `thinking_config` in `generate_content_config` and a '
+ 'planner with `thinking_config` are provided. The '
+ "planner's configuration will take precedence.",
+ UserWarning,
+ stacklevel=5,
+ )
+
  @classmethod
  @experimental
  def _resolve_tools(

src/google/adk/planners/built_in_planner.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import logging
 from typing import List
 from typing import Optional
 
@@ -23,6 +26,8 @@
 from ..models.llm_request import LlmRequest
 from .base_planner import BasePlanner
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class BuiltInPlanner(BasePlanner):
  """The built-in planner that uses model's built-in thinking features.
@@ -56,7 +61,11 @@ def apply_thinking_config(self, llm_request: LlmRequest) -> None:
  llm_request: The LLM request to apply the thinking config to.
  """
  if self.thinking_config:
- llm_request.config = llm_request.config or types.GenerateContentConfig()
+ if llm_request.config.thinking_config:
+ _LOGGER.info(
+ 'Overwriting `thinking_config` from `generate_content_config` with '
+ 'the one provided by the `BuiltInPlanner`.'
+ )
  llm_request.config.thinking_config = self.thinking_config
 
  @override

tests/unittests/agents/test_llm_agent_fields.py

@@ -14,6 +14,7 @@
 
 """Unit tests for canonical_xxx fields in LlmAgent."""
 
+import logging
 from typing import Any
 from typing import Optional
 from unittest import mock
@@ -27,6 +28,7 @@
 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
@@ -234,17 +236,35 @@ def _before_model_callback(
  assert agent.before_model_callback is not None
 
 
-def test_validate_generate_content_config_thinking_config_throw():
- with pytest.raises(ValueError):
- _ = LlmAgent(
- name='test_agent',
- generate_content_config=types.GenerateContentConfig(
- thinking_config=types.ThinkingConfig()
- ),
- )
+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',
@@ -255,6 +275,7 @@ def test_validate_generate_content_config_tools_throw():
 
 
 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',
@@ -265,6 +286,8 @@ def test_validate_generate_content_config_system_instruction_throw():
 
 
 def test_validate_generate_content_config_response_schema_throw():
+ """Tests that response schema cannot be set directly in config."""
+
  class Schema(BaseModel):
  pass
 
@@ -471,3 +494,27 @@ def test_agent_with_litellm_string_model(model_name):
  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 INFO 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.INFO):
+ planner.apply_thinking_config(req)
+
+ assert (
+ 'Overwriting `thinking_config` from `generate_content_config`'
+ in caplog.text
+ )