diff --git a/src/google/adk/models/llm_request.py b/src/google/adk/models/llm_request.py index 39fddef4..b83fd1d9 100644 --- a/src/google/adk/models/llm_request.py +++ b/src/google/adk/models/llm_request.py @@ -24,6 +24,24 @@ from pydantic import Field from ..tools.base_tool import BaseTool +def _find_tool_with_function_declarations( + llm_request: LlmRequest, +) -> Optional[types.Tool]: + """Find an existing Tool with function_declarations in the LlmRequest.""" + # TODO: add individual tool with declaration and merge in google_llm.py + if not llm_request.config or not llm_request.config.tools: + return None + + return next( + ( + tool + for tool in llm_request.config.tools + if isinstance(tool, types.Tool) and tool.function_declarations + ), + None, + ) + + class LlmRequest(BaseModel): """LLM request class that allows passing in tools, output schema and system @@ -81,15 +99,26 @@ class LlmRequest(BaseModel): return declarations = [] for tool in tools: - if isinstance(tool, BaseTool): - declaration = tool._get_declaration() - else: - declaration = tool.get_declaration() + declaration = tool._get_declaration() if declaration: declarations.append(declaration) self.tools_dict[tool.name] = tool if declarations: - self.config.tools.append(types.Tool(function_declarations=declarations)) + if self.config.tools is None: + self.config.tools = [] + + # Find existing tool with function_declarations and append to it + if tool_with_function_declarations := _find_tool_with_function_declarations( + self + ): + if tool_with_function_declarations.function_declarations is None: + tool_with_function_declarations.function_declarations = [] + tool_with_function_declarations.function_declarations.extend( + declarations + ) + else: + # No existing tool with function_declarations, create new one + self.config.tools.append(types.Tool(function_declarations=declarations)) def set_output_schema(self, base_model: type[BaseModel]) -> None: """Sets the output schema for the request. diff --git a/src/google/adk/tools/base_tool.py b/src/google/adk/tools/base_tool.py index 21f721fb..90c57539 100644 --- a/src/google/adk/tools/base_tool.py +++ b/src/google/adk/tools/base_tool.py @@ -125,30 +125,8 @@ class BaseTool(ABC): tool_context: The context of the tool. llm_request: The outgoing LLM request, mutable this method. """ - if (function_declaration := self._get_declaration()) is None: - return - - llm_request.tools_dict[self.name] = self - if tool_with_function_declarations := _find_tool_with_function_declarations( - llm_request - ): - if tool_with_function_declarations.function_declarations is None: - tool_with_function_declarations.function_declarations = [] - tool_with_function_declarations.function_declarations.append( - function_declaration - ) - else: - llm_request.config = ( - types.GenerateContentConfig() - if not llm_request.config - else llm_request.config - ) - llm_request.config.tools = ( - [] if not llm_request.config.tools else llm_request.config.tools - ) - llm_request.config.tools.append( - types.Tool(function_declarations=[function_declaration]) - ) + # Use the consolidated logic in LlmRequest.append_tools + llm_request.append_tools([self]) @property def _api_variant(self) -> GoogleLLMVariant: @@ -232,20 +210,3 @@ class BaseTool(ABC): else: logger.warning("Unsupported parsing for argument: %s.", param_name) return cls(**kwargs) - - -def _find_tool_with_function_declarations( - llm_request: LlmRequest, -) -> Optional[types.Tool]: - # TODO: add individual tool with declaration and merge in google_llm.py - if not llm_request.config or not llm_request.config.tools: - return None - - return next( - ( - tool - for tool in llm_request.config.tools - if isinstance(tool, types.Tool) and tool.function_declarations - ), - None, - ) diff --git a/tests/unittests/models/test_llm_request.py b/tests/unittests/models/test_llm_request.py new file mode 100644 index 00000000..78942296 --- /dev/null +++ b/tests/unittests/models/test_llm_request.py @@ -0,0 +1,310 @@ +# Copyright 2025 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. + +"""Tests for LlmRequest functionality.""" + +import asyncio +from typing import Optional + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.sequential_agent import SequentialAgent +from google.adk.models.llm_request import LlmRequest +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.tool_context import ToolContext +from google.genai import types +import pytest + + +def dummy_tool(query: str) -> str: + """A dummy tool for testing.""" + return f'Searched for: {query}' + + +def test_append_tools_with_none_config_tools(): + """Test that append_tools initializes config.tools when it's None.""" + request = LlmRequest() + + # Initially config.tools should be None + assert request.config.tools is None + + # Create a tool to append + tool = FunctionTool(func=dummy_tool) + + # This should not raise an AttributeError + request.append_tools([tool]) + + # Now config.tools should be initialized and contain the tool + assert request.config.tools is not None + assert len(request.config.tools) == 1 + assert len(request.config.tools[0].function_declarations) == 1 + assert request.config.tools[0].function_declarations[0].name == 'dummy_tool' + + # Tool should also be in tools_dict + assert 'dummy_tool' in request.tools_dict + assert request.tools_dict['dummy_tool'] == tool + + +def test_append_tools_with_existing_tools(): + """Test that append_tools works correctly when config.tools already exists.""" + request = LlmRequest() + + # Pre-initialize config.tools with an existing tool + existing_declaration = types.FunctionDeclaration( + name='existing_tool', description='An existing tool', parameters={} + ) + request.config.tools = [ + types.Tool(function_declarations=[existing_declaration]) + ] + + # Create a new tool to append + tool = FunctionTool(func=dummy_tool) + + # Append the new tool + request.append_tools([tool]) + + # Should still have 1 tool but now with 2 function declarations + assert len(request.config.tools) == 1 + assert len(request.config.tools[0].function_declarations) == 2 + + # Verify both declarations are present + decl_names = { + decl.name for decl in request.config.tools[0].function_declarations + } + assert decl_names == {'existing_tool', 'dummy_tool'} + + +def test_append_tools_empty_list(): + """Test that append_tools handles empty list correctly.""" + request = LlmRequest() + + # This should not modify anything + request.append_tools([]) + + # config.tools should still be None + assert request.config.tools is None + assert len(request.tools_dict) == 0 + + +def test_append_tools_tool_with_no_declaration(): + """Test append_tools with a BaseTool that returns None from _get_declaration.""" + from google.adk.tools.base_tool import BaseTool + + request = LlmRequest() + + # Create a mock tool that inherits from BaseTool and returns None for declaration + class NoDeclarationTool(BaseTool): + + def __init__(self): + super().__init__( + name='no_decl_tool', description='A tool with no declaration' + ) + + def _get_declaration(self): + return None + + tool = NoDeclarationTool() + + # This should not add anything to config.tools but should handle gracefully + request.append_tools([tool]) + + # config.tools should still be None since no declarations were added + assert request.config.tools is None + # tools_dict should be empty since no valid declaration + assert len(request.tools_dict) == 0 + + +def test_append_tools_consolidates_declarations_in_single_tool(): + """Test that append_tools puts all function declarations in a single Tool.""" + request = LlmRequest() + + # Create multiple tools + tool1 = FunctionTool(func=dummy_tool) + + def another_tool(param: str) -> str: + return f'Another: {param}' + + def third_tool(value: int) -> int: + return value * 2 + + tool2 = FunctionTool(func=another_tool) + tool3 = FunctionTool(func=third_tool) + + # Append all tools at once + request.append_tools([tool1, tool2, tool3]) + + # Should have exactly 1 Tool with 3 function declarations + assert len(request.config.tools) == 1 + assert len(request.config.tools[0].function_declarations) == 3 + + # Verify all tools are in tools_dict + assert len(request.tools_dict) == 3 + assert 'dummy_tool' in request.tools_dict + assert 'another_tool' in request.tools_dict + assert 'third_tool' in request.tools_dict + + +async def _create_tool_context() -> ToolContext: + """Helper to create a ToolContext for testing.""" + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name='test_app', user_id='test_user' + ) + agent = SequentialAgent(name='test_agent') + invocation_context = InvocationContext( + invocation_id='invocation_id', + agent=agent, + session=session, + session_service=session_service, + ) + return ToolContext(invocation_context) + + +class _MockTool(BaseTool): + """Mock tool for testing process_llm_request behavior.""" + + def __init__(self, name: str): + super().__init__(name=name, description=f'Mock tool {name}') + + def _get_declaration(self) -> Optional[types.FunctionDeclaration]: + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters=types.Schema(type=types.Type.STRING, title='param'), + ) + + +@pytest.mark.asyncio +async def test_process_llm_request_consolidates_declarations_in_single_tool(): + """Test that multiple process_llm_request calls consolidate in single Tool.""" + request = LlmRequest() + tool_context = await _create_tool_context() + + # Create multiple tools + tool1 = _MockTool('tool1') + tool2 = _MockTool('tool2') + tool3 = _MockTool('tool3') + + # Process each tool individually (simulating what happens in real usage) + await tool1.process_llm_request( + tool_context=tool_context, llm_request=request + ) + await tool2.process_llm_request( + tool_context=tool_context, llm_request=request + ) + await tool3.process_llm_request( + tool_context=tool_context, llm_request=request + ) + + # Should have exactly 1 Tool with 3 function declarations + assert len(request.config.tools) == 1 + assert len(request.config.tools[0].function_declarations) == 3 + + # Verify all function declaration names + decl_names = [ + decl.name for decl in request.config.tools[0].function_declarations + ] + assert 'tool1' in decl_names + assert 'tool2' in decl_names + assert 'tool3' in decl_names + + # Verify all tools are in tools_dict + assert len(request.tools_dict) == 3 + assert 'tool1' in request.tools_dict + assert 'tool2' in request.tools_dict + assert 'tool3' in request.tools_dict + + +@pytest.mark.asyncio +async def test_append_tools_and_process_llm_request_consistent_behavior(): + """Test that append_tools and process_llm_request produce same structure.""" + tool_context = await _create_tool_context() + + # Test 1: Using append_tools + request1 = LlmRequest() + tool1 = _MockTool('tool1') + tool2 = _MockTool('tool2') + tool3 = _MockTool('tool3') + request1.append_tools([tool1, tool2, tool3]) + + # Test 2: Using process_llm_request + request2 = LlmRequest() + tool4 = _MockTool('tool1') # Same names for comparison + tool5 = _MockTool('tool2') + tool6 = _MockTool('tool3') + await tool4.process_llm_request( + tool_context=tool_context, llm_request=request2 + ) + await tool5.process_llm_request( + tool_context=tool_context, llm_request=request2 + ) + await tool6.process_llm_request( + tool_context=tool_context, llm_request=request2 + ) + + # Both approaches should produce identical structure + assert len(request1.config.tools) == len(request2.config.tools) == 1 + assert len(request1.config.tools[0].function_declarations) == 3 + assert len(request2.config.tools[0].function_declarations) == 3 + + # Function declaration names should match + decl_names1 = { + decl.name for decl in request1.config.tools[0].function_declarations + } + decl_names2 = { + decl.name for decl in request2.config.tools[0].function_declarations + } + assert decl_names1 == decl_names2 == {'tool1', 'tool2', 'tool3'} + + +def test_multiple_append_tools_calls_consolidate(): + """Test that multiple append_tools calls add to the same Tool.""" + request = LlmRequest() + + # First call to append_tools + tool1 = FunctionTool(func=dummy_tool) + request.append_tools([tool1]) + + # Should have 1 tool with 1 declaration + assert len(request.config.tools) == 1 + assert len(request.config.tools[0].function_declarations) == 1 + assert request.config.tools[0].function_declarations[0].name == 'dummy_tool' + + # Second call to append_tools with different tools + def another_tool(param: str) -> str: + return f'Another: {param}' + + def third_tool(value: int) -> int: + return value * 2 + + tool2 = FunctionTool(func=another_tool) + tool3 = FunctionTool(func=third_tool) + request.append_tools([tool2, tool3]) + + # Should still have 1 tool but now with 3 declarations + assert len(request.config.tools) == 1 + assert len(request.config.tools[0].function_declarations) == 3 + + # Verify all declaration names are present + decl_names = { + decl.name for decl in request.config.tools[0].function_declarations + } + assert decl_names == {'dummy_tool', 'another_tool', 'third_tool'} + + # Verify all tools are in tools_dict + assert len(request.tools_dict) == 3 + assert 'dummy_tool' in request.tools_dict + assert 'another_tool' in request.tools_dict + assert 'third_tool' in request.tools_dict