test_api_models.py

"""
Test script for Model 4 (Perplexity) and Model 5 (Gemini) APIs.
"""
import os
from pathlib import Path
from dotenv import load_dotenv
from perplexity_service import get_perplexity_service
from gemini_service import get_gemini_service

# Load .env file
env_path = Path(__file__).parent / '.env'
if env_path.exists():
    load_dotenv(env_path)
    print(f"✓ Loaded .env file from {env_path}")
else:
    print(f"⚠ .env file not found at {env_path}, using environment variables only")

# Test data
TEST_CONVERSATION = "Tối nay anh có lịch đột xuất. ||| Thế mai được không?"
TEST_TRIGGER = "neutral"
TEST_MOVE = "escalate"

def test_perplexity():
    """Test Perplexity API (Model 4)."""
    print("=" * 60)
    print("Testing Model 4 - Perplexity API")
    print("=" * 60)
    
    try:
        service = get_perplexity_service()
        print(f"✓ Perplexity service initialized")
        print(f"  Model: {service.model}")
        print(f"  API Key: {'Set' if service.api_key else 'Not set'}")
        
        print("\nGenerating reply...")
        reply = service.generate_reply(
            conversation=TEST_CONVERSATION,
            trigger=TEST_TRIGGER,
            move=TEST_MOVE,
        )
        
        print(f"\n✓ Success!")
        print(f"  Input conversation: {TEST_CONVERSATION}")
        print(f"  Trigger: {TEST_TRIGGER}")
        print(f"  Move: {TEST_MOVE}")
        print(f"  Generated reply: {reply}")
        print(f"  Reply length: {len(reply.split())} words")
        
        return True, reply
        
    except Exception as e:
        print(f"\n✗ Error: {str(e)}")
        return False, str(e)


def test_gemini():
    """Test Gemini API (Model 5)."""
    print("\n" + "=" * 60)
    print("Testing Model 5 - Gemini API")
    print("=" * 60)
    
    try:
        # Try default model first
        model_name = "gemini-2.0-flash"
        service = get_gemini_service(model_name=model_name)
        print(f"✓ Gemini service initialized")
        print(f"  Model: {service.model_name}")
        print(f"  API Key: {'Set' if service.api_key else 'Not set'}")
        
        print("\nGenerating reply...")
        formatted_conversation = f"Male: {TEST_CONVERSATION.split('|||')[0].strip()} ||| Female: {TEST_CONVERSATION.split('|||')[1].strip()}"
        reply = service.generate_reply(
            conversation=formatted_conversation,
            trigger=TEST_TRIGGER,
            move=TEST_MOVE,
            max_output_tokens=500,  # Use higher limit for testing
        )
        
        print(f"\n✓ Success!")
        print(f"  Input conversation: {formatted_conversation}")
        print(f"  Trigger: {TEST_TRIGGER}")
        print(f"  Move: {TEST_MOVE}")
        print(f"  Generated reply: {reply}")
        print(f"  Reply length: {len(reply.split())} words")
        
        return True, reply
        
    except Exception as e:
        print(f"\n✗ Error: {str(e)}")
        return False, str(e)


def main():
    """Run all tests."""
    print("\n" + "=" * 60)
    print("API Endpoint Test Script")
    print("=" * 60)
    print(f"\nTest Data:")
    print(f"  Conversation: {TEST_CONVERSATION}")
    print(f"  Trigger: {TEST_TRIGGER}")
    print(f"  Move: {TEST_MOVE}")
    print()
    
    # Check environment variables
    print("Environment Variables:")
    perplexity_key = os.getenv("PERPLEXITY_API_KEY")
    google_key = os.getenv("GOOGLE_API_KEY")
    print(f"  PERPLEXITY_API_KEY: {'✓ Set' if perplexity_key else '✗ Not set'}")
    print(f"  GOOGLE_API_KEY: {'✓ Set' if google_key else '✗ Not set'}")
    print()
    
    # Test Perplexity
    perplexity_success, perplexity_result = test_perplexity()
    
    # Test Gemini
    gemini_success, gemini_result = test_gemini()
    
    # Summary
    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)
    print(f"  Model 4 (Perplexity): {'✓ PASSED' if perplexity_success else '✗ FAILED'}")
    print(f"  Model 5 (Gemini): {'✓ PASSED' if gemini_success else '✗ FAILED'}")
    print()
    
    if perplexity_success and gemini_success:
        print("✓ All tests passed!")
        return 0
    else:
        print("✗ Some tests failed. Check errors above.")
        return 1


if __name__ == "__main__":
    exit(main())