video-image-production-build/tests/agents/test_model_config.py

from __future__ import annotations

import os
import unittest
from pathlib import Path
from unittest.mock import patch

from production_build_agents.agents.model_config import (
    create_configured_chat_model,
)
from production_build_agents.agents.executor.agent import (
    create_real_executor_model,
)
from production_build_agents.agents.planner.agent import (
    create_real_planner_model,
)
from production_build_agents.agents.validator.task_agent import (
    create_real_validator_model,
)


class _ConfigurationError(RuntimeError):
    pass


class ModelConfigTest(unittest.TestCase):
    @patch("production_build_agents.agents.model_config.init_chat_model")
    @patch("production_build_agents.agents.model_config.load_dotenv")
    def test_uses_root_env_and_role_fallback(
        self,
        load_dotenv,
        init_chat_model,
    ) -> None:
        configured_model = object()
        init_chat_model.return_value = configured_model
        environment = {
            "PLANNER_MODEL": "planner-model",
            "PLANNER_MODEL_PROVIDER": "openai",
            "PLANNER_API_KEY": "test-key",
            "PLANNER_BASE_URL": "https://example.test/v1",
            "PLANNER_ENABLE_THINKING": "true",
        }
        with patch.dict(os.environ, environment, clear=True):
            actual = create_configured_chat_model(
                roles=("EXECUTOR", "PLANNER"),
                error_type=_ConfigurationError,
                missing_message="missing",
            )

        project_root = Path(__file__).resolve().parents[2]
        load_dotenv.assert_called_once_with(
            project_root / ".env",
            override=False,
        )
        init_chat_model.assert_called_once_with(
            "planner-model",
            model_provider="openai",
            temperature=0,
            timeout=180.0,
            max_retries=1,
            api_key="test-key",
            base_url="https://example.test/v1",
            extra_body={"enable_thinking": True},
        )
        self.assertIs(actual, configured_model)

    @patch("production_build_agents.agents.model_config.load_dotenv")
    def test_missing_model_uses_role_error(self, load_dotenv) -> None:
        with patch.dict(os.environ, {}, clear=True):
            with self.assertRaisesRegex(_ConfigurationError, "missing model"):
                create_configured_chat_model(
                    roles=("VALIDATOR", "EXECUTOR", "PLANNER"),
                    error_type=_ConfigurationError,
                    missing_message="missing model",
                )
        load_dotenv.assert_called_once()

    @patch("production_build_agents.agents.model_config.init_chat_model")
    @patch(
        "production_build_agents.observability.context_tracking."
        "observed_http_clients"
    )
    def test_full_observation_adds_transport_client_without_changing_retries(
        self,
        observed_http_clients,
        init_chat_model,
    ) -> None:
        transport_client = object()
        observed_http_clients.return_value = {
            "http_client": transport_client,
        }
        with patch.dict(
            os.environ,
            {
                "PLANNER_MODEL": "planner-model",
                "PLANNER_MODEL_PROVIDER": "openai",
                "PLANNER_MAX_RETRIES": "3",
            },
            clear=True,
        ):
            create_configured_chat_model(
                roles=("PLANNER",),
                error_type=_ConfigurationError,
                missing_message="missing",
            )

        observed_http_clients.assert_called_once_with("openai")
        self.assertEqual(
            init_chat_model.call_args.kwargs["max_retries"],
            3,
        )
        self.assertIs(
            init_chat_model.call_args.kwargs["http_client"],
            transport_client,
        )
        self.assertEqual(
            init_chat_model.call_args.kwargs["http_socket_options"],
            (),
        )

    @patch("production_build_agents.agents.model_config.init_chat_model")
    def test_each_agent_loads_its_model_from_root_env(
        self,
        init_chat_model,
    ) -> None:
        project_root = Path(__file__).resolve().parents[2]
        cases = (
            ("PLANNER", create_real_planner_model),
            ("EXECUTOR", create_real_executor_model),
            ("VALIDATOR", create_real_validator_model),
        )
        for role, factory in cases:
            with self.subTest(role=role), patch.dict(
                os.environ,
                {f"{role}_MODEL": f"{role.lower()}-model"},
                clear=True,
            ), patch(
                "production_build_agents.agents.model_config.load_dotenv"
            ) as load_dotenv:
                factory()
                load_dotenv.assert_called_once_with(
                    project_root / ".env",
                    override=False,
                )
        self.assertEqual(init_chat_model.call_count, 3)