116 lines
5.2 KiB
Python
116 lines
5.2 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from ra_aid.llm import create_llm_client
|
|
|
|
|
|
class TestDisableThinking:
|
|
"""Test suite for the disable_thinking configuration option."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_id, config, model_config, expected_thinking_param, description",
|
|
[
|
|
# Test case 1: Claude 3.7 model with thinking enabled (default)
|
|
(
|
|
"claude_thinking_enabled",
|
|
{
|
|
"provider": "anthropic",
|
|
"model": "claude-3-7-sonnet-20250219",
|
|
},
|
|
{"supports_thinking": True},
|
|
{"thinking": {"type": "enabled", "budget_tokens": 12000}},
|
|
"Claude 3.7 should have thinking enabled by default",
|
|
),
|
|
# Test case 2: Claude 3.7 model with thinking explicitly disabled
|
|
(
|
|
"claude_thinking_disabled",
|
|
{
|
|
"provider": "anthropic",
|
|
"model": "claude-3-7-sonnet-20250219",
|
|
"disable_thinking": True,
|
|
},
|
|
{"supports_thinking": True},
|
|
{},
|
|
"Claude 3.7 with disable_thinking=True should not have thinking param",
|
|
),
|
|
# Test case 3: Non-thinking model should not have thinking param
|
|
(
|
|
"non_thinking_model",
|
|
{
|
|
"provider": "anthropic",
|
|
"model": "claude-3-5-sonnet-20240620",
|
|
},
|
|
{"supports_thinking": False},
|
|
{},
|
|
"Non-thinking model should not have thinking param",
|
|
),
|
|
# Test case 4: Non-Claude model should not have thinking param
|
|
(
|
|
"non_claude_model",
|
|
{
|
|
"provider": "openai",
|
|
"model": "gpt-4",
|
|
},
|
|
{},
|
|
{},
|
|
"Non-Claude model should not have thinking param",
|
|
),
|
|
# Test case 5: Claude 3.7 model with skip_sonnet37_workaround enabled
|
|
(
|
|
"skip_sonnet37_workaround",
|
|
{
|
|
"provider": "anthropic",
|
|
"model": "claude-3-7-sonnet-20250219",
|
|
"skip_sonnet37_workaround": True,
|
|
},
|
|
{"supports_thinking": True},
|
|
{"thinking": {"type": "enabled", "budget_tokens": 12000}},
|
|
"Claude 3.7 with skip_sonnet37_workaround=True should still have thinking param",
|
|
),
|
|
],
|
|
)
|
|
def test_disable_thinking_option(
|
|
self, test_id, config, model_config, expected_thinking_param, description
|
|
):
|
|
"""Test that the disable_thinking option correctly controls thinking mode."""
|
|
# Mock the necessary dependencies
|
|
with patch("ra_aid.llm.ChatAnthropic") as mock_anthropic:
|
|
with patch("ra_aid.llm.ChatOpenAI") as mock_openai:
|
|
with patch("ra_aid.llm.models_params") as mock_models_params:
|
|
# Set up the mock to return the specified model_config
|
|
mock_models_params.get.return_value = {config["model"]: model_config}
|
|
|
|
# Set up the mock for get_provider_config
|
|
with patch("ra_aid.llm.get_provider_config") as mock_get_provider_config:
|
|
# Include disable_thinking in the provider config if it's in the test config
|
|
provider_config = {
|
|
"api_key": "test-key",
|
|
"base_url": None,
|
|
}
|
|
if "disable_thinking" in config:
|
|
provider_config["disable_thinking"] = config["disable_thinking"]
|
|
|
|
mock_get_provider_config.return_value = provider_config
|
|
|
|
# Call the function without passing disable_thinking directly
|
|
create_llm_client(
|
|
config["provider"],
|
|
config["model"],
|
|
temperature=None,
|
|
is_expert=False
|
|
)
|
|
|
|
# Check if the correct parameters were passed
|
|
if config["provider"] == "anthropic":
|
|
# Get the kwargs passed to ChatAnthropic
|
|
_, kwargs = mock_anthropic.call_args
|
|
|
|
# Check if thinking param was included or not
|
|
if expected_thinking_param:
|
|
assert "thinking" in kwargs, f"Test {test_id} failed: {description}"
|
|
assert kwargs["thinking"] == expected_thinking_param["thinking"], \
|
|
f"Test {test_id} failed: Thinking param doesn't match expected value"
|
|
else:
|
|
assert "thinking" not in kwargs, \
|
|
f"Test {test_id} failed: Thinking param should not be present"
|