"""
Test script for tool calling methods to ensure proper handling of different LLM method arguments.
"""
from unittest.mock import Mock
import pytest
from langchain_core.language_models.chat_models import BaseChatModel
from browser_use.agent.service import Agent
class TestToolCallingMethods:
"""Tests for different tool calling methods handling."""
@pytest.fixture
def mock_llm_with_structured_output(self):
"""Create a mock LLM that tracks with_structured_output calls."""
mock_llm = Mock(spec=BaseChatModel)
# Track calls to with_structured_output
structured_output_calls = []
def mock_with_structured_output(schema, include_raw=True, method=None):
structured_output_calls.append({'schema': schema, 'include_raw': include_raw, 'method': method})
# Return a mock that can be invoked
mock_structured = Mock()
mock_structured.invoke = Mock(return_value={'parsed': None, 'raw': Mock(content='test')})
mock_structured.ainvoke = Mock(return_value={'parsed': None, 'raw': Mock(content='test')})
return mock_structured
mock_llm.with_structured_output = mock_with_structured_output
mock_llm.structured_output_calls = structured_output_calls
mock_llm.invoke = Mock(return_value=Mock(content='{"answer": "paris"}'))
return mock_llm
async def test_tools_method_error(self, mock_llm_with_structured_output):
"""Test that 'tools' method causes the expected error."""
# Create agent with 'tools' method
agent = Agent(
task='Test task',
llm=mock_llm_with_structured_output,
tool_calling_method='tools',
)
# The error should occur during initialization when _test_tool_calling_method is called
# Check that with_structured_output was called with 'tools' method
assert len(mock_llm_with_structured_output.structured_output_calls) > 0
# Find the call that used 'tools' method
tools_call = next(
(call for call in mock_llm_with_structured_output.structured_output_calls if call['method'] == 'tools'), None
)
assert tools_call is not None, "Expected with_structured_output to be called with method='tools'"
async def test_function_calling_method_works(self, mock_llm_with_structured_output):
"""Test that 'function_calling' method works correctly."""
# Create agent with 'function_calling' method
agent = Agent(
task='Test task',
llm=mock_llm_with_structured_output,
tool_calling_method='function_calling',
)
# Check that with_structured_output was called with 'function_calling' method
assert len(mock_llm_with_structured_output.structured_output_calls) > 0
# Find the call that used 'function_calling' method
fc_call = next(
(call for call in mock_llm_with_structured_output.structured_output_calls if call['method'] == 'function_calling'),
None,
)
assert fc_call is not None, "Expected with_structured_output to be called with method='function_calling'"
async def test_json_mode_method_works(self, mock_llm_with_structured_output):
"""Test that 'json_mode' method works correctly."""
# Create agent with 'json_mode' method
agent = Agent(
task='Test task',
llm=mock_llm_with_structured_output,
tool_calling_method='json_mode',
)
# For json_mode, it should not call with_structured_output during testing
# because it uses raw mode
# Check the calls
json_mode_calls = [
call for call in mock_llm_with_structured_output.structured_output_calls if call['method'] == 'json_mode'
]
# json_mode is handled specially and doesn't use with_structured_output in _test_tool_calling_method
assert len(json_mode_calls) == 0
async def test_raw_method_works(self, mock_llm_with_structured_output):
"""Test that 'raw' method works correctly."""
# Create agent with 'raw' method
agent = Agent(
task='Test task',
llm=mock_llm_with_structured_output,
tool_calling_method='raw',
)
# For raw mode, it should not call with_structured_output during testing
# Check the calls
raw_calls = [call for call in mock_llm_with_structured_output.structured_output_calls if call['method'] == 'raw']
# raw is handled specially and doesn't use with_structured_output in _test_tool_calling_method
assert len(raw_calls) == 0
async def test_auto_method_selection(self, mock_llm_with_structured_output):
"""Test that 'auto' method selects appropriate method based on LLM."""
# Mock the agent to simulate Azure OpenAI with GPT-4
agent = Agent(
task='Test task',
llm=mock_llm_with_structured_output,
tool_calling_method='auto',
)
# Monkey patch to simulate Azure OpenAI
agent.chat_model_library = 'AzureChatOpenAI'
agent.chat_model_name = 'gpt-4-1106-preview'
# Call _get_tool_calling_method_for_model to see what it returns
method = agent._get_tool_calling_method_for_model()
# For Azure OpenAI with GPT-4, it should return 'tools'
assert method == 'tools'
