Skip to content

Commit e162bb8

Browse files
Akshat8510wuliang229
authored andcommitted
feat: Allow thinking_config in generate_content_config
Merge #4117 **Overview** This PR implements the feature request in #4108 to allow `thinking_config` to be set directly within `generate_content_config`, bringing the Python SDK in line with the Go implementation. **Changes** - **llm_agent.py**: Relaxed the validation logic in `validate_generate_content_config` to remove the `ValueError` for `thinking_config`. - **Precedence Warning**: Added an override of `model_post_init` in `LlmAgent` to issue a `UserWarning` if both a `planner` and a manual `thinking_config` are provided. - **built_in_planner.py**: Updated `apply_thinking_config` to log an `INFO` message when the planner overwrites an existing configuration on the `LlmRequest`. **Testing** Verified with a reproduction script covering: 1. Successful initialization of an agent with direct `thinking_config`. 2. Validation of `UserWarning` during initialization when conflicting configurations are present. 3. Confirmation of logger output when the planner performs an overwrite. Closes: #4108 Tagging @invictus2010 for visibility. Co-authored-by: Liang Wu <wuliang@google.com> COPYBARA_INTEGRATE_REVIEW=#4117 from Akshat8510:feat/allow-thinking-config-4108 5deeb89 PiperOrigin-RevId: 856821447
1 parent 19315fe commit e162bb8

File tree

3 files changed

+84
-11
lines changed

3 files changed

+84
-11
lines changed

src/google/adk/agents/llm_agent.py

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ class LlmAgent(BaseAgent):
285285
"""The additional content generation configurations.
286286
287287
NOTE: not all fields are usable, e.g. tools must be configured via `tools`,
288-
thinking_config must be configured via `planner` in LlmAgent.
288+
thinking_config can be configured here or via the `planner`. If both are set, the planner's configuration takes precedence.
289289
290290
For example: use this config to adjust model temperature, configure safety
291291
settings, etc.
@@ -849,8 +849,6 @@ def validate_generate_content_config(
849849
) -> types.GenerateContentConfig:
850850
if not generate_content_config:
851851
return types.GenerateContentConfig()
852-
if generate_content_config.thinking_config:
853-
raise ValueError('Thinking config should be set via LlmAgent.planner.')
854852
if generate_content_config.tools:
855853
raise ValueError('All tools must be set via LlmAgent.tools.')
856854
if generate_content_config.system_instruction:
@@ -863,6 +861,23 @@ def validate_generate_content_config(
863861
)
864862
return generate_content_config
865863

864+
@override
865+
def model_post_init(self, __context: Any) -> None:
866+
"""Provides a warning if multiple thinking configurations are found."""
867+
super().model_post_init(__context)
868+
869+
# Note: Using getattr to check both locations for thinking_config
870+
if getattr(
871+
self.generate_content_config, 'thinking_config', None
872+
) and getattr(self.planner, 'thinking_config', None):
873+
warnings.warn(
874+
'Both `thinking_config` in `generate_content_config` and a '
875+
'planner with `thinking_config` are provided. The '
876+
"planner's configuration will take precedence.",
877+
UserWarning,
878+
stacklevel=3,
879+
)
880+
866881
@classmethod
867882
@experimental
868883
def _resolve_tools(

src/google/adk/planners/built_in_planner.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from __future__ import annotations
16+
17+
import logging
1518
from typing import List
1619
from typing import Optional
1720

@@ -23,6 +26,8 @@
2326
from ..models.llm_request import LlmRequest
2427
from .base_planner import BasePlanner
2528

29+
logger = logging.getLogger('google_adk.' + __name__)
30+
2631

2732
class BuiltInPlanner(BasePlanner):
2833
"""The built-in planner that uses model's built-in thinking features.
@@ -57,6 +62,11 @@ def apply_thinking_config(self, llm_request: LlmRequest) -> None:
5762
"""
5863
if self.thinking_config:
5964
llm_request.config = llm_request.config or types.GenerateContentConfig()
65+
if llm_request.config.thinking_config:
66+
logger.debug(
67+
'Overwriting `thinking_config` from `generate_content_config` with '
68+
'the one provided by the `BuiltInPlanner`.'
69+
)
6070
llm_request.config.thinking_config = self.thinking_config
6171

6272
@override

tests/unittests/agents/test_llm_agent_fields.py

Lines changed: 56 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
"""Unit tests for canonical_xxx fields in LlmAgent."""
1616

17+
import logging
1718
from typing import Any
1819
from typing import Optional
1920
from unittest import mock
@@ -27,6 +28,7 @@
2728
from google.adk.models.lite_llm import LiteLlm
2829
from google.adk.models.llm_request import LlmRequest
2930
from google.adk.models.registry import LLMRegistry
31+
from google.adk.planners.built_in_planner import BuiltInPlanner
3032
from google.adk.sessions.in_memory_session_service import InMemorySessionService
3133
from google.adk.tools.google_search_tool import google_search
3234
from google.adk.tools.google_search_tool import GoogleSearchTool
@@ -234,17 +236,35 @@ def _before_model_callback(
234236
assert agent.before_model_callback is not None
235237

236238

237-
def test_validate_generate_content_config_thinking_config_throw():
238-
with pytest.raises(ValueError):
239-
_ = LlmAgent(
240-
name='test_agent',
241-
generate_content_config=types.GenerateContentConfig(
242-
thinking_config=types.ThinkingConfig()
243-
),
244-
)
239+
def test_validate_generate_content_config_thinking_config_allow():
240+
"""Tests that thinking_config is now allowed directly in the agent init."""
241+
agent = LlmAgent(
242+
name='test_agent',
243+
generate_content_config=types.GenerateContentConfig(
244+
thinking_config=types.ThinkingConfig(include_thoughts=True)
245+
),
246+
)
247+
assert agent.generate_content_config.thinking_config.include_thoughts is True
248+
249+
250+
def test_thinking_config_precedence_warning():
251+
"""Tests that a UserWarning is issued when both manual config and planner exist."""
252+
253+
config = types.GenerateContentConfig(
254+
thinking_config=types.ThinkingConfig(include_thoughts=True)
255+
)
256+
planner = BuiltInPlanner(
257+
thinking_config=types.ThinkingConfig(include_thoughts=True)
258+
)
259+
260+
with pytest.warns(
261+
UserWarning, match="planner's configuration will take precedence"
262+
):
263+
LlmAgent(name='test_agent', generate_content_config=config, planner=planner)
245264

246265

247266
def test_validate_generate_content_config_tools_throw():
267+
"""Tests that tools cannot be set directly in config."""
248268
with pytest.raises(ValueError):
249269
_ = LlmAgent(
250270
name='test_agent',
@@ -255,6 +275,7 @@ def test_validate_generate_content_config_tools_throw():
255275

256276

257277
def test_validate_generate_content_config_system_instruction_throw():
278+
"""Tests that system instructions cannot be set directly in config."""
258279
with pytest.raises(ValueError):
259280
_ = LlmAgent(
260281
name='test_agent',
@@ -265,6 +286,8 @@ def test_validate_generate_content_config_system_instruction_throw():
265286

266287

267288
def test_validate_generate_content_config_response_schema_throw():
289+
"""Tests that response schema cannot be set directly in config."""
290+
268291
class Schema(BaseModel):
269292
pass
270293

@@ -471,3 +494,28 @@ def test_agent_with_litellm_string_model(model_name):
471494
agent = LlmAgent(name='test_agent', model=model_name)
472495
assert isinstance(agent.canonical_model, LiteLlm)
473496
assert agent.canonical_model.model == model_name
497+
498+
499+
def test_builtin_planner_overwrite_logging(caplog):
500+
"""Tests that the planner logs an DEBUG message when overwriting a config."""
501+
502+
planner = BuiltInPlanner(
503+
thinking_config=types.ThinkingConfig(include_thoughts=True)
504+
)
505+
506+
# Create a request that already has a thinking_config
507+
req = LlmRequest(
508+
contents=[],
509+
config=types.GenerateContentConfig(
510+
thinking_config=types.ThinkingConfig(include_thoughts=True)
511+
),
512+
)
513+
514+
with caplog.at_level(
515+
logging.DEBUG, logger='google_adk.google.adk.planners.built_in_planner'
516+
):
517+
planner.apply_thinking_config(req)
518+
assert (
519+
'Overwriting `thinking_config` from `generate_content_config`'
520+
in caplog.text
521+
)

0 commit comments

Comments
 (0)