mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: incorrect logic in LlmRequest.append_tools and make BaseTool to call it
PiperOrigin-RevId: 792518526
This commit is contained in:
Google Team Member
committed by
Copybara-Service
Copybara-Service
parent
4b2a8e3186
commit
98a5730115
@@ -24,24 +24,6 @@ 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
|
||||
|
||||
@@ -99,26 +81,15 @@ class LlmRequest(BaseModel):
|
||||
return
|
||||
declarations = []
|
||||
for tool in tools:
|
||||
declaration = tool._get_declaration()
|
||||
if isinstance(tool, BaseTool):
|
||||
declaration = tool._get_declaration()
|
||||
else:
|
||||
declaration = tool.get_declaration()
|
||||
if declaration:
|
||||
declarations.append(declaration)
|
||||
self.tools_dict[tool.name] = tool
|
||||
if 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))
|
||||
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.
|
||||
|
||||
@@ -125,8 +125,30 @@ class BaseTool(ABC):
|
||||
tool_context: The context of the tool.
|
||||
llm_request: The outgoing LLM request, mutable this method.
|
||||
"""
|
||||
# Use the consolidated logic in LlmRequest.append_tools
|
||||
llm_request.append_tools([self])
|
||||
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])
|
||||
)
|
||||
|
||||
@property
|
||||
def _api_variant(self) -> GoogleLLMVariant:
|
||||
@@ -210,3 +232,20 @@ 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,
|
||||
)
|
||||
|
||||
@@ -1,310 +0,0 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user