crewAI/tests/test_gemini_html_template.py

"""Test Gemini models with HTML templates."""

from unittest.mock import MagicMock, patch

import pytest

from crewai import Agent, Task
from crewai.llm import LLM


def test_gemini_empty_response_handling():
    """Test that empty responses from Gemini models are handled correctly."""
    # Create a mock LLM instance
    llm = LLM(model="gemini/gemini-pro")
    
    # Create a mock response with empty content
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message = MagicMock()
    mock_response.choices[0].message.content = ""
    
    # Mock litellm.completion to return our mock response
    with patch('litellm.completion', return_value=mock_response):
        # Call the non-streaming response handler directly
        result = llm._handle_non_streaming_response({"model": "gemini/gemini-pro"})
        
        # Verify that our fix works - empty string should be replaced with placeholder
        assert "Response processed successfully" in result
        assert "HTML template" in result


def test_openrouter_gemini_empty_response_handling():
    """Test that empty responses from OpenRouter with Gemini models are handled correctly."""
    # Create a mock LLM instance with OpenRouter base URL
    llm = LLM(
        model="openrouter/google/gemini-pro", 
        base_url="https://openrouter.ai/api/v1"
    )
    
    # Create a mock response with empty content
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message = MagicMock()
    mock_response.choices[0].message.content = ""
    
    # Mock litellm.completion to return our mock response
    with patch('litellm.completion', return_value=mock_response):
        # Call the non-streaming response handler directly
        result = llm._handle_non_streaming_response({"model": "openrouter/google/gemini-pro"})
        
        # Verify that our fix works - empty string should be replaced with placeholder
        assert "Response processed successfully" in result
        assert "HTML template" in result


def test_gemini_none_response_handling():
    """Test that None responses are properly handled."""
    llm = LLM(model="gemini/gemini-pro")
    
    # Create a mock response with None content
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message = MagicMock()
    mock_response.choices[0].message.content = None
    
    # Mock litellm.completion to return our mock response
    with patch('litellm.completion', return_value=mock_response):
        # Call the non-streaming response handler directly
        # None content should be converted to empty string and then handled
        result = llm._handle_non_streaming_response({"model": "gemini/gemini-pro"})
        
        # Verify that our fix works - None should be converted to empty string
        # and then handled as an empty string for Gemini models
        assert "Response processed successfully" in result
        assert "HTML template" in result


@pytest.mark.parametrize("model_name,base_url", [
    ("gemini/gemini-pro", None),
    ("gemini-pro", None),
    ("google/gemini-pro", None),
    ("openrouter/google/gemini-pro", "https://openrouter.ai/api/v1"),
    ("openrouter/gemini-pro", "https://openrouter.ai/api/v1"),
])
def test_various_gemini_configurations(model_name, base_url):
    """Test different Gemini model configurations with the _is_gemini_model helper."""
    # Create a mock LLM instance with the specified model and base URL
    llm = LLM(model=model_name, base_url=base_url)
    
    # Verify that _is_gemini_model correctly identifies all these configurations
    assert llm._is_gemini_model() is True
    
    # Create a mock response with empty content
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message = MagicMock()
    mock_response.choices[0].message.content = ""
    
    # Mock litellm.completion to return our mock response
    with patch('litellm.completion', return_value=mock_response):
        # Call the non-streaming response handler directly
        result = llm._handle_non_streaming_response({"model": model_name})
        
        # Verify that our fix works for all Gemini configurations
        assert "Response processed successfully" in result
        assert "HTML template" in result


def test_non_gemini_model():
    """Test that non-Gemini models don't get special handling for empty responses."""
    # Create a mock LLM instance with a non-Gemini model
    llm = LLM(model="gpt-4")
    
    # Verify that _is_gemini_model correctly identifies this as not a Gemini model
    assert llm._is_gemini_model() is False
    
    # Create a mock response with empty content
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message = MagicMock()
    mock_response.choices[0].message.content = ""
    
    # Mock litellm.completion to return our mock response
    with patch('litellm.completion', return_value=mock_response):
        # Call the non-streaming response handler directly
        result = llm._handle_non_streaming_response({"model": "gpt-4"})
        
        # Verify that non-Gemini models just return the empty string
        assert result == ""