test_api_key.py

"""
🔐 Quick Perplexity API Key Test
Run this to verify your API key is working correctly
"""

import os
import requests

def test_perplexity_api_key():
    """Test if Perplexity API key is working"""
    
    print("🔐 Testing Perplexity API Key...")
    print("=" * 40)
    
    # Check if API key exists
    api_key = os.getenv('PERPLEXITY_API_KEY')
    
    if not api_key:
        print("❌ PERPLEXITY_API_KEY environment variable not found")
        print("\n💡 To fix this:")
        print("1. Go to your HuggingFace Space Settings")
        print("2. Add Repository secret:")
        print("   Name: PERPLEXITY_API_KEY")
        print("   Value: your_api_key_here")
        print("3. Restart your Space")
        return False
    
    print(f"✅ API key found: {api_key[:12]}...{api_key[-4:]}")
    
    # Test API connection
    print("\n🧪 Testing API connection...")
    
    try:
        headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
        
        payload = {
            "model": "llama-3.1-sonar-large-128k-online",
            "messages": [
                {
                    "role": "user",
                    "content": "Test API connection - respond with 'API working' if you receive this."
                }
            ],
            "max_tokens": 50,
            "temperature": 0.1
        }
        
        response = requests.post(
            'https://api.perplexity.ai/chat/completions',
            headers=headers,
            json=payload,
            timeout=30
        )
        
        if response.status_code == 200:
            print("✅ API connection successful!")
            
            # Parse response
            try:
                result = response.json()
                content = result['choices'][0]['message']['content']
                print(f"📝 API Response: {content}")
                print("\n🎉 Your Perplexity API integration is working correctly!")
                return True
                
            except Exception as e:
                print(f"⚠️ Response parsing issue: {e}")
                print("✅ But API connection works!")
                return True
                
        elif response.status_code == 401:
            print("❌ API key is invalid or expired")
            print("\n💡 To fix this:")
            print("1. Check your API key at https://www.perplexity.ai/")
            print("2. Generate a new key if needed")
            print("3. Update the PERPLEXITY_API_KEY secret in your Space")
            return False
            
        elif response.status_code == 429:
            print("⚠️ API rate limit reached")
            print("✅ But API key is valid!")
            return True
            
        else:
            print(f"❌ API error: {response.status_code}")
            print(f"Response: {response.text}")
            return False
            
    except requests.exceptions.Timeout:
        print("⏰ API request timed out")
        print("⚠️ This might be a temporary network issue")
        return False
        
    except requests.exceptions.RequestException as e:
        print(f"🔌 Connection error: {e}")
        return False

def test_ai_dataset_studio_integration():
    """Test the full AI Dataset Studio integration"""
    
    print("\n🚀 Testing AI Dataset Studio Integration...")
    print("=" * 40)
    
    try:
        # Try to import our Perplexity client
        from perplexity_client import PerplexityClient, SearchType
        print("✅ Perplexity client module imported")
        
        # Initialize client
        client = PerplexityClient()
        print("✅ Perplexity client initialized")
        
        # Test API key validation
        if client._validate_api_key():
            print("✅ API key validation successful")
        else:
            print("❌ API key validation failed")
            return False
        
        # Test source discovery
        print("\n🔍 Testing source discovery...")
        results = client.discover_sources(
            project_description="Test query for API integration verification",
            search_type=SearchType.GENERAL,
            max_sources=3
        )
        
        if results.sources:
            print(f"✅ Source discovery working! Found {len(results.sources)} sources")
            print(f"   Example: {results.sources[0].title}")
        else:
            print("⚠️ Source discovery returned no results (but API is working)")
        
        print("\n🎉 AI Dataset Studio integration is fully functional!")
        return True
        
    except ImportError as e:
        print(f"❌ Import error: {e}")
        print("💡 Make sure perplexity_client.py is uploaded to your Space")
        return False
        
    except Exception as e:
        print(f"❌ Integration test failed: {e}")
        return False

if __name__ == "__main__":
    print("🧪 Perplexity API Integration Test")
    print("=" * 50)
    
    # Test 1: Basic API key
    api_test = test_perplexity_api_key()
    
    # Test 2: Full integration (only if basic test passes)
    if api_test:
        integration_test = test_ai_dataset_studio_integration()
        
        if api_test and integration_test:
            print("\n" + "=" * 50)
            print("🎉 ALL TESTS PASSED!")
            print("Your AI Dataset Studio is ready to use!")
            print("🚀 Start creating amazing datasets with AI-powered source discovery!")
        else:
            print("\n" + "=" * 50)
            print("⚠️ Some tests failed, but basic API connectivity works")
            print("Check the error messages above for specific issues")
    else:
        print("\n" + "=" * 50)
        print("❌ API key test failed")
        print("Please fix the API key configuration before proceeding")