test_models.py

import unittest
from unittest.mock import ANY, MagicMock, patch

from aider.models import (
    ANTHROPIC_BETA_HEADER,
    Model,
    ModelInfoManager,
    register_models,
    sanity_check_model,
    sanity_check_models,
)


class TestModels(unittest.TestCase):
    def setUp(self):
        """Reset MODEL_SETTINGS before each test"""
        from aider.models import MODEL_SETTINGS

        self._original_settings = MODEL_SETTINGS.copy()

    def tearDown(self):
        """Restore original MODEL_SETTINGS after each test"""
        from aider.models import MODEL_SETTINGS

        MODEL_SETTINGS.clear()
        MODEL_SETTINGS.extend(self._original_settings)

    def test_get_model_info_nonexistent(self):
        manager = ModelInfoManager()
        info = manager.get_model_info("non-existent-model")
        self.assertEqual(info, {})

    def test_max_context_tokens(self):
        model = Model("gpt-3.5-turbo")
        self.assertEqual(model.info["max_input_tokens"], 16385)

        model = Model("gpt-3.5-turbo-16k")
        self.assertEqual(model.info["max_input_tokens"], 16385)

        model = Model("gpt-3.5-turbo-1106")
        self.assertEqual(model.info["max_input_tokens"], 16385)

        model = Model("gpt-4")
        self.assertEqual(model.info["max_input_tokens"], 8 * 1024)

        model = Model("gpt-4-0613")
        self.assertEqual(model.info["max_input_tokens"], 8 * 1024)

    @patch("os.environ")
    def test_sanity_check_model_all_set(self, mock_environ):
        mock_environ.get.return_value = "dummy_value"
        mock_io = MagicMock()
        model = MagicMock()
        model.name = "test-model"
        model.missing_keys = ["API_KEY1", "API_KEY2"]
        model.keys_in_environment = True
        model.info = {"some": "info"}

        sanity_check_model(mock_io, model)

        mock_io.tool_output.assert_called()
        calls = mock_io.tool_output.call_args_list
        self.assertIn("- API_KEY1: Set", str(calls))
        self.assertIn("- API_KEY2: Set", str(calls))

    @patch("os.environ")
    def test_sanity_check_model_not_set(self, mock_environ):
        mock_environ.get.return_value = ""
        mock_io = MagicMock()
        model = MagicMock()
        model.name = "test-model"
        model.missing_keys = ["API_KEY1", "API_KEY2"]
        model.keys_in_environment = True
        model.info = {"some": "info"}

        sanity_check_model(mock_io, model)

        mock_io.tool_output.assert_called()
        calls = mock_io.tool_output.call_args_list
        self.assertIn("- API_KEY1: Not set", str(calls))
        self.assertIn("- API_KEY2: Not set", str(calls))

    def test_sanity_check_models_bogus_editor(self):
        mock_io = MagicMock()
        main_model = Model("gpt-4")
        main_model.editor_model = Model("bogus-model")

        result = sanity_check_models(mock_io, main_model)

        self.assertTrue(
            result
        )  # Should return True because there's a problem with the editor model
        mock_io.tool_warning.assert_called_with(ANY)  # Ensure a warning was issued

        warning_messages = [
            warning_call.args[0] for warning_call in mock_io.tool_warning.call_args_list
        ]
        print("Warning messages:", warning_messages)  # Add this line

        self.assertGreaterEqual(mock_io.tool_warning.call_count, 1)  # Expect two warnings
        self.assertTrue(
            any("bogus-model" in msg for msg in warning_messages)
        )  # Check that one of the warnings mentions the bogus model

    @patch("aider.models.check_for_dependencies")
    def test_sanity_check_model_calls_check_dependencies(self, mock_check_deps):
        """Test that sanity_check_model calls check_for_dependencies"""
        mock_io = MagicMock()
        model = MagicMock()
        model.name = "test-model"
        model.missing_keys = []
        model.keys_in_environment = True
        model.info = {"some": "info"}

        sanity_check_model(mock_io, model)

        # Verify check_for_dependencies was called with the model name
        mock_check_deps.assert_called_once_with(mock_io, "test-model")

    def test_model_aliases(self):
        # Test common aliases
        model = Model("4")
        self.assertEqual(model.name, "gpt-4-0613")

        model = Model("4o")
        self.assertEqual(model.name, "gpt-4o")

        model = Model("35turbo")
        self.assertEqual(model.name, "gpt-3.5-turbo")

        model = Model("35-turbo")
        self.assertEqual(model.name, "gpt-3.5-turbo")

        model = Model("3")
        self.assertEqual(model.name, "gpt-3.5-turbo")

        model = Model("sonnet")
        self.assertEqual(model.name, "claude-sonnet-4-6")

        model = Model("haiku")
        self.assertEqual(model.name, "claude-haiku-4-5")

        model = Model("opus")
        self.assertEqual(model.name, "claude-opus-4-7")

        # Test non-alias passes through unchanged
        model = Model("gpt-4")
        self.assertEqual(model.name, "gpt-4")

    def test_o1_use_temp_false(self):
        # Test GitHub Copilot models
        model = Model("github/o1-mini")
        self.assertEqual(model.name, "github/o1-mini")
        self.assertEqual(model.use_temperature, False)

        model = Model("github/o1-preview")
        self.assertEqual(model.name, "github/o1-preview")
        self.assertEqual(model.use_temperature, False)

    def test_parse_token_value(self):
        # Create a model instance to test the parse_token_value method
        model = Model("gpt-4")

        # Test integer inputs
        self.assertEqual(model.parse_token_value(8096), 8096)
        self.assertEqual(model.parse_token_value(1000), 1000)

        # Test string inputs
        self.assertEqual(model.parse_token_value("8096"), 8096)

        # Test k/K suffix (kilobytes)
        self.assertEqual(model.parse_token_value("8k"), 8 * 1024)
        self.assertEqual(model.parse_token_value("8K"), 8 * 1024)
        self.assertEqual(model.parse_token_value("10.5k"), 10.5 * 1024)
        self.assertEqual(model.parse_token_value("0.5K"), 0.5 * 1024)

        # Test m/M suffix (megabytes)
        self.assertEqual(model.parse_token_value("1m"), 1 * 1024 * 1024)
        self.assertEqual(model.parse_token_value("1M"), 1 * 1024 * 1024)
        self.assertEqual(model.parse_token_value("0.5M"), 0.5 * 1024 * 1024)

        # Test with spaces
        self.assertEqual(model.parse_token_value(" 8k "), 8 * 1024)

        # Test conversion from other types
        self.assertEqual(model.parse_token_value(8.0), 8)

    def test_set_thinking_tokens(self):
        # Test that set_thinking_tokens correctly sets the tokens with different formats
        model = Model("gpt-4")

        # Test with integer
        model.set_thinking_tokens(8096)
        self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 8096)
        self.assertFalse(model.use_temperature)

        # Test with string
        model.set_thinking_tokens("10k")
        self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 10 * 1024)

        # Test with decimal value
        model.set_thinking_tokens("0.5M")
        self.assertEqual(model.extra_params["thinking"]["budget_tokens"], 0.5 * 1024 * 1024)

    @patch("aider.models.check_pip_install_extra")
    def test_check_for_dependencies_bedrock(self, mock_check_pip):
        """Test that check_for_dependencies calls check_pip_install_extra for Bedrock models"""
        from aider.io import InputOutput

        io = InputOutput()

        # Test with a Bedrock model
        from aider.models import check_for_dependencies

        check_for_dependencies(io, "bedrock/anthropic.claude-3-sonnet-20240229-v1:0")

        # Verify check_pip_install_extra was called with correct arguments
        mock_check_pip.assert_called_once_with(
            io, "boto3", "AWS Bedrock models require the boto3 package.", ["boto3"]
        )

    @patch("aider.models.check_pip_install_extra")
    def test_check_for_dependencies_vertex_ai(self, mock_check_pip):
        """Test that check_for_dependencies calls check_pip_install_extra for Vertex AI models"""
        from aider.io import InputOutput

        io = InputOutput()

        # Test with a Vertex AI model
        from aider.models import check_for_dependencies

        check_for_dependencies(io, "vertex_ai/gemini-1.5-pro")

        # Verify check_pip_install_extra was called with correct arguments
        mock_check_pip.assert_called_once_with(
            io,
            "google.cloud.aiplatform",
            "Google Vertex AI models require the google-cloud-aiplatform package.",
            ["google-cloud-aiplatform"],
        )

    @patch("aider.models.check_pip_install_extra")
    def test_check_for_dependencies_other_model(self, mock_check_pip):
        """Test that check_for_dependencies doesn't call check_pip_install_extra for other models"""
        from aider.io import InputOutput

        io = InputOutput()

        # Test with a non-Bedrock, non-Vertex AI model
        from aider.models import check_for_dependencies

        check_for_dependencies(io, "gpt-4")

        # Verify check_pip_install_extra was not called
        mock_check_pip.assert_not_called()

    def test_get_repo_map_tokens(self):
        # Test default case (no max_input_tokens in info)
        model = Model("gpt-4")
        model.info = {}
        self.assertEqual(model.get_repo_map_tokens(), 1024)

        # Test minimum boundary (max_input_tokens < 8192)
        model.info = {"max_input_tokens": 4096}
        self.assertEqual(model.get_repo_map_tokens(), 1024)

        # Test middle range (max_input_tokens = 16384)
        model.info = {"max_input_tokens": 16384}
        self.assertEqual(model.get_repo_map_tokens(), 2048)

        # Test maximum boundary (max_input_tokens > 32768)
        model.info = {"max_input_tokens": 65536}
        self.assertEqual(model.get_repo_map_tokens(), 4096)

        # Test exact boundary values
        model.info = {"max_input_tokens": 8192}
        self.assertEqual(model.get_repo_map_tokens(), 1024)

        model.info = {"max_input_tokens": 32768}
        self.assertEqual(model.get_repo_map_tokens(), 4096)

    def test_configure_model_settings(self):
        # Test o3-mini case
        model = Model("something/o3-mini")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertFalse(model.use_temperature)

        # Test o1-mini case
        model = Model("something/o1-mini")
        self.assertTrue(model.use_repo_map)
        self.assertFalse(model.use_temperature)
        self.assertFalse(model.use_system_prompt)

        # Test o1-preview case
        model = Model("something/o1-preview")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertFalse(model.use_temperature)
        self.assertFalse(model.use_system_prompt)

        # Test o1 case
        model = Model("something/o1")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertFalse(model.use_temperature)
        self.assertFalse(model.streaming)

        # Test deepseek v3 case
        model = Model("deepseek-v3")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertEqual(model.reminder, "sys")
        self.assertTrue(model.examples_as_sys_msg)

        # Test deepseek reasoner case
        model = Model("deepseek-r1")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertTrue(model.examples_as_sys_msg)
        self.assertFalse(model.use_temperature)
        self.assertEqual(model.reasoning_tag, "think")

        # Test provider/deepseek-r1 case
        model = Model("someprovider/deepseek-r1")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertTrue(model.examples_as_sys_msg)
        self.assertFalse(model.use_temperature)
        self.assertEqual(model.reasoning_tag, "think")

        # Test provider/deepseek-v3 case
        model = Model("anotherprovider/deepseek-v3")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertEqual(model.reminder, "sys")
        self.assertTrue(model.examples_as_sys_msg)

        # Test llama3 70b case
        model = Model("llama3-70b")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertTrue(model.send_undo_reply)
        self.assertTrue(model.examples_as_sys_msg)

        # Test gpt-4 case
        model = Model("gpt-4")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertTrue(model.send_undo_reply)

        # Test gpt-3.5 case
        model = Model("gpt-3.5")
        self.assertEqual(model.reminder, "sys")

        # Test 3.5-sonnet case
        model = Model("claude-3.5-sonnet")
        self.assertEqual(model.edit_format, "diff")
        self.assertTrue(model.use_repo_map)
        self.assertTrue(model.examples_as_sys_msg)
        self.assertEqual(model.reminder, "user")

        # Test o1- prefix case
        model = Model("o1-something")
        self.assertFalse(model.use_system_prompt)
        self.assertFalse(model.use_temperature)

        # Test qwen case
        model = Model("qwen-coder-2.5-32b")
        self.assertEqual(model.edit_format, "diff")
        self.assertEqual(model.editor_edit_format, "editor-diff")
        self.assertTrue(model.use_repo_map)

    def test_aider_extra_model_settings(self):
        import tempfile

        import yaml

        # Create temporary YAML file with test settings
        test_settings = [
            {
                "name": "aider/extra_params",
                "extra_params": {
                    "extra_headers": {"Foo": "bar"},
                    "some_param": "some value",
                },
            },
        ]

        # Write to a regular file instead of NamedTemporaryFile
        # for better cross-platform compatibility
        tmp = tempfile.mktemp(suffix=".yml")
        try:
            with open(tmp, "w") as f:
                yaml.dump(test_settings, f)

            # Register the test settings
            register_models([tmp])

            # Test that defaults are applied when no exact match
            model = Model("claude-3-5-sonnet-20240620")
            # Test that both the override and existing headers are present
            model = Model("claude-3-5-sonnet-20240620")
            self.assertEqual(model.extra_params["extra_headers"]["Foo"], "bar")
            self.assertEqual(
                model.extra_params["extra_headers"]["anthropic-beta"],
                ANTHROPIC_BETA_HEADER,
            )
            self.assertEqual(model.extra_params["some_param"], "some value")
            self.assertEqual(model.extra_params["max_tokens"], 8192)

            # Test that exact match overrides defaults but not overrides
            model = Model("gpt-4")
            self.assertEqual(model.extra_params["extra_headers"]["Foo"], "bar")
            self.assertEqual(model.extra_params["some_param"], "some value")
        finally:
            # Clean up the temporary file
            import os

            try:
                os.unlink(tmp)
            except OSError:
                pass

    @patch("aider.models.litellm.completion")
    @patch.object(Model, "token_count")
    def test_ollama_num_ctx_set_when_missing(self, mock_token_count, mock_completion):
        mock_token_count.return_value = 1000

        model = Model("ollama/llama3")
        messages = [{"role": "user", "content": "Hello"}]

        model.send_completion(messages, functions=None, stream=False)

        # Verify num_ctx was calculated and added to call
        expected_ctx = int(1000 * 1.25) + 8192  # 9442
        mock_completion.assert_called_once_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0,
            num_ctx=expected_ctx,
            timeout=600,
        )

    @patch("aider.models.litellm.completion")
    def test_ollama_uses_existing_num_ctx(self, mock_completion):
        model = Model("ollama/llama3")
        model.extra_params = {"num_ctx": 4096}

        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)

        # Should use provided num_ctx from extra_params
        mock_completion.assert_called_once_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0,
            num_ctx=4096,
            timeout=600,
        )

    @patch("aider.models.litellm.completion")
    def test_non_ollama_no_num_ctx(self, mock_completion):
        model = Model("gpt-4")
        messages = [{"role": "user", "content": "Hello"}]

        model.send_completion(messages, functions=None, stream=False)

        # Regular models shouldn't get num_ctx
        mock_completion.assert_called_once_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0,
            timeout=600,
        )
        self.assertNotIn("num_ctx", mock_completion.call_args.kwargs)

    def test_use_temperature_settings(self):
        # Test use_temperature=True (default) uses temperature=0
        model = Model("gpt-4")
        self.assertTrue(model.use_temperature)
        self.assertEqual(model.use_temperature, True)

        # Test use_temperature=False doesn't pass temperature
        model = Model("github/o1-mini")
        self.assertFalse(model.use_temperature)

        # Test use_temperature as float value
        model = Model("gpt-4")
        model.use_temperature = 0.7
        self.assertEqual(model.use_temperature, 0.7)

    @patch("aider.models.litellm.completion")
    def test_request_timeout_default(self, mock_completion):
        # Test default timeout is used when not specified in extra_params
        model = Model("gpt-4")
        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)
        mock_completion.assert_called_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0,
            timeout=600,  # Default timeout
        )

    @patch("aider.models.litellm.completion")
    def test_request_timeout_from_extra_params(self, mock_completion):
        # Test timeout from extra_params overrides default
        model = Model("gpt-4")
        model.extra_params = {"timeout": 300}  # 5 minutes
        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)
        mock_completion.assert_called_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0,
            timeout=300,  # From extra_params
        )

    @patch("aider.models.litellm.completion")
    def test_use_temperature_in_send_completion(self, mock_completion):
        # Test use_temperature=True sends temperature=0
        model = Model("gpt-4")
        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)
        mock_completion.assert_called_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0,
            timeout=600,
        )

        # Test use_temperature=False doesn't send temperature
        model = Model("github/o1-mini")
        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)
        self.assertNotIn("temperature", mock_completion.call_args.kwargs)

        # Test use_temperature as float sends that value
        model = Model("gpt-4")
        model.use_temperature = 0.7
        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)
        mock_completion.assert_called_with(
            model=model.name,
            messages=messages,
            stream=False,
            temperature=0.7,
            timeout=600,
        )

    def test_gpt_5_5_model_settings(self):
        base_models = [
            "gpt-5.5",
            "gpt-5.5-chat-latest",
            "openai/gpt-5.5",
            "openai/gpt-5.5-chat-latest",
            "azure/gpt-5.5",
            "azure/gpt-5.5-chat-latest",
            "openrouter/openai/gpt-5.5",
            "openrouter/openai/gpt-5.5-chat-latest",
        ]

        for name in base_models:
            with self.subTest(name=name):
                model = Model(name)
                self.assertEqual(model.edit_format, "diff")
                self.assertTrue(model.use_repo_map)
                self.assertFalse(model.use_temperature)
                self.assertIn("reasoning_effort", model.accepts_settings)

        self.assertTrue(Model("gpt-5.5").overeager)

        pro_models = {
            "gpt-5.5-pro": "gpt-5.5",
            "openai/gpt-5.5-pro": "openai/gpt-5.5",
            "azure/gpt-5.5-pro": "azure/gpt-5.5",
            "openrouter/openai/gpt-5.5-pro": "openrouter/openai/gpt-5.5",
        }

        for name, editor_name in pro_models.items():
            with self.subTest(name=name):
                model = Model(name)
                self.assertFalse(model.streaming)
                self.assertEqual(model.edit_format, "diff")
                self.assertTrue(model.use_repo_map)
                self.assertFalse(model.use_temperature)
                self.assertEqual(model.editor_model_name, editor_name)
                self.assertEqual(model.editor_model.name, editor_name)
                self.assertIn("reasoning_effort", model.accepts_settings)

    def test_parse_custom_headers_key_value_and_json(self):
        from aider.args import parse_custom_headers

        self.assertEqual(parse_custom_headers(["X-Api-Key=secret"]), {"X-Api-Key": "secret"})
        self.assertEqual(
            parse_custom_headers(["X-Tenant: acme", "X-Trace=abc"]),
            {"X-Tenant": "acme", "X-Trace": "abc"},
        )
        self.assertEqual(
            parse_custom_headers(['{"Authorization": "Bearer tok", "X-Org": "42"}']),
            {"Authorization": "Bearer tok", "X-Org": "42"},
        )
        self.assertEqual(
            parse_custom_headers(["X-Foo=bar", '{"X-Bar": "baz"}']),
            {"X-Foo": "bar", "X-Bar": "baz"},
        )

    @patch("aider.models.litellm.completion")
    def test_custom_headers_applied_when_initializing_model_client(self, mock_completion):
        headers = {"X-Custom-Auth": "token-123", "X-Org-Id": "org-9"}
        model = Model("gpt-4", extra_headers=headers, weak_model=False)

        self.assertEqual(model.extra_params["extra_headers"]["X-Custom-Auth"], "token-123")
        self.assertEqual(model.extra_params["extra_headers"]["X-Org-Id"], "org-9")

        messages = [{"role": "user", "content": "Hello"}]
        model.send_completion(messages, functions=None, stream=False)

        mock_completion.assert_called_once()
        kwargs = mock_completion.call_args.kwargs
        self.assertEqual(kwargs["extra_headers"]["X-Custom-Auth"], "token-123")
        self.assertEqual(kwargs["extra_headers"]["X-Org-Id"], "org-9")
        self.assertEqual(kwargs["model"], model.name)

    @patch("aider.models.litellm.completion")
    def test_custom_headers_merge_with_existing_extra_headers(self, mock_completion):
        model = Model(
            "claude-3-5-sonnet-20240620",
            extra_headers={"X-Gateway": "byok"},
            weak_model=False,
        )
        extra_headers = model.extra_params["extra_headers"]
        self.assertEqual(extra_headers["X-Gateway"], "byok")
        self.assertEqual(extra_headers["anthropic-beta"], ANTHROPIC_BETA_HEADER)

        model.send_completion([{"role": "user", "content": "Hi"}], functions=None, stream=False)
        sent = mock_completion.call_args.kwargs["extra_headers"]
        self.assertEqual(sent["X-Gateway"], "byok")
        self.assertEqual(sent["anthropic-beta"], ANTHROPIC_BETA_HEADER)


if __name__ == "__main__":
    unittest.main()