mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Allow thinking_config in generate_content_config
Merge https://github.com/google/adk-python/pull/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=https://github.com/google/adk-python/pull/4117 from Akshat8510:feat/allow-thinking-config-4108 5deeb893799379c681d6822dc4a1e42f86d3ed01 PiperOrigin-RevId: 856821447
This commit is contained in:
committed by
Copybara-Service
parent
19315fe557
commit
e162bb8832
@@ -285,7 +285,7 @@ class LlmAgent(BaseAgent):
|
|||||||
"""The additional content generation configurations.
|
"""The additional content generation configurations.
|
||||||
|
|
||||||
NOTE: not all fields are usable, e.g. tools must be configured via `tools`,
|
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
|
For example: use this config to adjust model temperature, configure safety
|
||||||
settings, etc.
|
settings, etc.
|
||||||
@@ -849,8 +849,6 @@ class LlmAgent(BaseAgent):
|
|||||||
) -> types.GenerateContentConfig:
|
) -> types.GenerateContentConfig:
|
||||||
if not generate_content_config:
|
if not generate_content_config:
|
||||||
return types.GenerateContentConfig()
|
return types.GenerateContentConfig()
|
||||||
if generate_content_config.thinking_config:
|
|
||||||
raise ValueError('Thinking config should be set via LlmAgent.planner.')
|
|
||||||
if generate_content_config.tools:
|
if generate_content_config.tools:
|
||||||
raise ValueError('All tools must be set via LlmAgent.tools.')
|
raise ValueError('All tools must be set via LlmAgent.tools.')
|
||||||
if generate_content_config.system_instruction:
|
if generate_content_config.system_instruction:
|
||||||
@@ -863,6 +861,23 @@ class LlmAgent(BaseAgent):
|
|||||||
)
|
)
|
||||||
return 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=3,
|
||||||
|
)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@experimental
|
@experimental
|
||||||
def _resolve_tools(
|
def _resolve_tools(
|
||||||
|
|||||||
@@ -12,6 +12,9 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
from typing import List
|
from typing import List
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
@@ -23,6 +26,8 @@ from ..agents.readonly_context import ReadonlyContext
|
|||||||
from ..models.llm_request import LlmRequest
|
from ..models.llm_request import LlmRequest
|
||||||
from .base_planner import BasePlanner
|
from .base_planner import BasePlanner
|
||||||
|
|
||||||
|
logger = logging.getLogger('google_adk.' + __name__)
|
||||||
|
|
||||||
|
|
||||||
class BuiltInPlanner(BasePlanner):
|
class BuiltInPlanner(BasePlanner):
|
||||||
"""The built-in planner that uses model's built-in thinking features.
|
"""The built-in planner that uses model's built-in thinking features.
|
||||||
@@ -57,6 +62,11 @@ class BuiltInPlanner(BasePlanner):
|
|||||||
"""
|
"""
|
||||||
if self.thinking_config:
|
if self.thinking_config:
|
||||||
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
||||||
|
if llm_request.config.thinking_config:
|
||||||
|
logger.debug(
|
||||||
|
'Overwriting `thinking_config` from `generate_content_config` with '
|
||||||
|
'the one provided by the `BuiltInPlanner`.'
|
||||||
|
)
|
||||||
llm_request.config.thinking_config = self.thinking_config
|
llm_request.config.thinking_config = self.thinking_config
|
||||||
|
|
||||||
@override
|
@override
|
||||||
|
|||||||
@@ -14,6 +14,7 @@
|
|||||||
|
|
||||||
"""Unit tests for canonical_xxx fields in LlmAgent."""
|
"""Unit tests for canonical_xxx fields in LlmAgent."""
|
||||||
|
|
||||||
|
import logging
|
||||||
from typing import Any
|
from typing import Any
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from unittest import mock
|
from unittest import mock
|
||||||
@@ -27,6 +28,7 @@ from google.adk.models.google_llm import Gemini
|
|||||||
from google.adk.models.lite_llm import LiteLlm
|
from google.adk.models.lite_llm import LiteLlm
|
||||||
from google.adk.models.llm_request import LlmRequest
|
from google.adk.models.llm_request import LlmRequest
|
||||||
from google.adk.models.registry import LLMRegistry
|
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.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 google_search
|
||||||
from google.adk.tools.google_search_tool import GoogleSearchTool
|
from google.adk.tools.google_search_tool import GoogleSearchTool
|
||||||
@@ -234,17 +236,35 @@ def test_before_model_callback():
|
|||||||
assert agent.before_model_callback is not None
|
assert agent.before_model_callback is not None
|
||||||
|
|
||||||
|
|
||||||
def test_validate_generate_content_config_thinking_config_throw():
|
def test_validate_generate_content_config_thinking_config_allow():
|
||||||
with pytest.raises(ValueError):
|
"""Tests that thinking_config is now allowed directly in the agent init."""
|
||||||
_ = LlmAgent(
|
agent = LlmAgent(
|
||||||
name='test_agent',
|
name='test_agent',
|
||||||
generate_content_config=types.GenerateContentConfig(
|
generate_content_config=types.GenerateContentConfig(
|
||||||
thinking_config=types.ThinkingConfig()
|
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():
|
def test_validate_generate_content_config_tools_throw():
|
||||||
|
"""Tests that tools cannot be set directly in config."""
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
_ = LlmAgent(
|
_ = LlmAgent(
|
||||||
name='test_agent',
|
name='test_agent',
|
||||||
@@ -255,6 +275,7 @@ def test_validate_generate_content_config_tools_throw():
|
|||||||
|
|
||||||
|
|
||||||
def test_validate_generate_content_config_system_instruction_throw():
|
def test_validate_generate_content_config_system_instruction_throw():
|
||||||
|
"""Tests that system instructions cannot be set directly in config."""
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
_ = LlmAgent(
|
_ = LlmAgent(
|
||||||
name='test_agent',
|
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():
|
def test_validate_generate_content_config_response_schema_throw():
|
||||||
|
"""Tests that response schema cannot be set directly in config."""
|
||||||
|
|
||||||
class Schema(BaseModel):
|
class Schema(BaseModel):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@@ -471,3 +494,28 @@ def test_agent_with_litellm_string_model(model_name):
|
|||||||
agent = LlmAgent(name='test_agent', model=model_name)
|
agent = LlmAgent(name='test_agent', model=model_name)
|
||||||
assert isinstance(agent.canonical_model, LiteLlm)
|
assert isinstance(agent.canonical_model, LiteLlm)
|
||||||
assert agent.canonical_model.model == model_name
|
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
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user