crewAI/tests/test_aiml_api_integration.py

"""Tests for AI/ML API integration with CrewAI."""

from unittest.mock import patch

from crewai.llm import LLM
from crewai.utilities.llm_utils import create_llm


class TestAIMLAPIIntegration:
    """Test suite for AI/ML API provider integration."""

    def test_aiml_api_model_context_windows(self):
        """Test that AI/ML API models have correct context window sizes."""
        test_cases = [
            ("openai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo", 131072),
            ("openai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", 131072),
            ("openai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", 131072),
            ("openai/anthropic/claude-3-5-sonnet-20241022", 200000),
            ("openai/anthropic/claude-3-5-haiku-20241022", 200000),
            ("openai/mistralai/Mistral-7B-Instruct-v0.3", 32768),
            ("openai/Qwen/Qwen2.5-72B-Instruct-Turbo", 131072),
            ("openai/deepseek-ai/DeepSeek-V2.5", 131072),
        ]
        
        for model_name, expected_context_size in test_cases:
            llm = LLM(model=model_name)
            expected_usable_size = int(expected_context_size * 0.85)
            actual_context_size = llm.get_context_window_size()
            assert actual_context_size == expected_usable_size, (
                f"Model {model_name} should have context window size {expected_usable_size}, "
                f"but got {actual_context_size}"
            )

    def test_aiml_api_provider_detection(self):
        """Test that AI/ML API models are correctly identified as openai provider."""
        llm = LLM(model="openai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo")
        provider = llm._get_custom_llm_provider()
        assert provider == "openai", f"Expected provider 'openai', but got '{provider}'"

    def test_aiml_api_model_instantiation(self):
        """Test that AI/ML API models can be instantiated correctly."""
        model_names = [
            "openai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
            "openai/anthropic/claude-3-5-sonnet-20241022",
            "openai/mistralai/Mixtral-8x7B-Instruct-v0.1",
            "openai/Qwen/Qwen2.5-72B-Instruct-Turbo",
        ]
        
        for model_name in model_names:
            llm = LLM(model=model_name)
            assert llm.model == model_name
            assert llm._get_custom_llm_provider() == "openai"
            assert llm.get_context_window_size() > 0

    @patch('crewai.llm.litellm.utils.supports_function_calling')
    def test_aiml_api_function_calling_support(self, mock_supports_function_calling):
        """Test function calling support detection for AI/ML API models."""
        mock_supports_function_calling.return_value = True
        
        llm = LLM(model="openai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo")
        supports_fc = llm.supports_function_calling()
        
        assert supports_fc is True
        mock_supports_function_calling.assert_called_once_with(
            "openai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
            custom_llm_provider="openai"
        )

    def test_aiml_api_with_create_llm(self):
        """Test that AI/ML API models work with create_llm utility."""
        model_name = "openai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"
        llm = create_llm(model_name)
        
        assert isinstance(llm, LLM)
        assert llm.model == model_name
        assert llm._get_custom_llm_provider() == "openai"

    def test_aiml_api_model_validation(self):
        """Test that AI/ML API models pass validation checks."""
        llm = LLM(model="openai/anthropic/claude-3-5-sonnet-20241022")
        
        llm._validate_call_params()
        
        llm_with_format = LLM(
            model="openai/anthropic/claude-3-5-sonnet-20241022",
            response_format={"type": "json_object"}
        )
        try:
            llm_with_format._validate_call_params()
        except ValueError as e:
            assert "does not support response_format" in str(e)

    def test_aiml_api_context_window_bounds(self):
        """Test that AI/ML API model context windows are within valid bounds."""
        from crewai.llm import LLM_CONTEXT_WINDOW_SIZES
        
        aiml_models = {k: v for k, v in LLM_CONTEXT_WINDOW_SIZES.items() 
                      if k.startswith("openai/")}
        
        MIN_CONTEXT = 1024
        MAX_CONTEXT = 2097152
        
        for model_name, context_size in aiml_models.items():
            assert MIN_CONTEXT <= context_size <= MAX_CONTEXT, (
                f"Model {model_name} context window {context_size} is outside "
                f"valid bounds [{MIN_CONTEXT}, {MAX_CONTEXT}]"
            )

    def test_aiml_api_model_prefixes(self):
        """Test that all AI/ML API models use the correct openai/ prefix."""
        from crewai.llm import LLM_CONTEXT_WINDOW_SIZES
        
        aiml_models = [k for k in LLM_CONTEXT_WINDOW_SIZES.keys() 
                      if k.startswith("openai/")]
        
        assert len(aiml_models) > 0, "No AI/ML API models found in context window sizes"
        
        for model_name in aiml_models:
            assert model_name.startswith("openai/"), (
                f"AI/ML API model {model_name} should start with 'openai/' prefix"
            )
            parts = model_name.split("/")
            assert len(parts) >= 3, (
                f"AI/ML API model {model_name} should have format 'openai/provider/model'"
            )


class TestAIMLAPIExamples:
    """Test examples of using AI/ML API with CrewAI components."""

    def test_aiml_api_with_agent_example(self):
        """Test example usage of AI/ML API with CrewAI Agent."""
        from crewai import Agent
        
        llm = LLM(model="openai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo")
        
        agent = Agent(
            role="AI Assistant",
            goal="Help users with their questions",
            backstory="You are a helpful AI assistant powered by Llama 3.1",
            llm=llm,
        )
        
        assert agent.llm.model == "openai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"
        assert agent.llm._get_custom_llm_provider() == "openai"

    def test_aiml_api_different_model_types(self):
        """Test different types of models available through AI/ML API."""
        model_types = {
            "llama": "openai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
            "claude": "openai/anthropic/claude-3-5-sonnet-20241022",
            "mistral": "openai/mistralai/Mixtral-8x7B-Instruct-v0.1",
            "qwen": "openai/Qwen/Qwen2.5-72B-Instruct-Turbo",
            "deepseek": "openai/deepseek-ai/DeepSeek-V2.5",
        }
        
        for model_type, model_name in model_types.items():
            llm = LLM(model=model_name)
            assert llm.model == model_name
            assert llm._get_custom_llm_provider() == "openai"
            assert llm.get_context_window_size() > 0, (
                f"{model_type} model should have positive context window"
            )