mirror of
https://github.com/browser-use/browser-use
synced 2026-05-06 17:52:15 +02:00
128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
"""
|
|
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'
|