validation_agent.py

# run_validation_agent.py
import asyncio
from openai import AsyncOpenAI
from schemas import ValidationResult

# ------------------------
# GPT validation via OpenAI SDK (Direct)
# ------------------------
async def validate_gpt_key(api_key: str, model: str = "gpt-4o") -> ValidationResult:
    """
    Validates a GPT API key using OpenAI SDK directly.
    """
    try:
        # Create OpenAI client
        client = AsyncOpenAI(api_key=api_key)
        
        # Simple completion test
        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are an API key validator."},
                {"role": "user", "content": "ping"}
            ],
            max_tokens=10,
            temperature=0
        )
        
        # Check if we got a valid response
        if response.choices and response.choices[0].message.content:
            return ValidationResult(is_valid=True)
        
        return ValidationResult(
            is_valid=False, 
            error="No valid response from OpenAI API"
        )

    except Exception as e:
        error_msg = str(e)
        # Provide more specific error messages
        if "invalid_api_key" in error_msg.lower():
            error_msg = "Invalid API key"
        elif "rate_limit" in error_msg.lower():
            error_msg = "Rate limit exceeded"
        elif "insufficient_quota" in error_msg.lower():
            error_msg = "Insufficient quota/credits"
        
        return ValidationResult(
            is_valid=False, 
            error=f"GPT validation failed: {error_msg}"
        )


# ------------------------
# Gemini (Google Generative AI) validation
# ------------------------
async def validate_gemini_key(api_key: str, model: str = "gemini-2.0-flash-exp") -> ValidationResult:
    """
    Validates a Gemini API key.
    """
    try:
        import google.generativeai as genai
        
        # Configure with API key
        genai.configure(api_key=api_key)
        
        # Try to generate content
        model_instance = genai.GenerativeModel(model)
        response = await model_instance.generate_content_async("ping")
        
        if response and response.text:
            return ValidationResult(is_valid=True)
        
        return ValidationResult(
            is_valid=False, 
            error="No valid response from Gemini API"
        )
        
    except Exception as e:
        error_msg = str(e)
        if "API_KEY_INVALID" in error_msg or "invalid" in error_msg.lower():
            error_msg = "Invalid API key"
        
        return ValidationResult(
            is_valid=False, 
            error=f"Gemini validation failed: {error_msg}"
        )


# ------------------------
# Grok (xAI) validation
# ------------------------
async def validate_grok_key(api_key: str, model: str = "grok-beta") -> ValidationResult:
    """
    Validates a Grok API key using OpenAI-compatible endpoint.
    """
    try:
        # Grok uses OpenAI-compatible API
        client = AsyncOpenAI(
            api_key=api_key,
            base_url="https://api.x.ai/v1"
        )
        
        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are an API key validator."},
                {"role": "user", "content": "ping"}
            ],
            max_tokens=10,
            temperature=0
        )
        
        if response.choices and response.choices[0].message.content:
            return ValidationResult(is_valid=True)
        
        return ValidationResult(
            is_valid=False, 
            error="No valid response from Grok API"
        )
        
    except Exception as e:
        error_msg = str(e)
        if "invalid_api_key" in error_msg.lower():
            error_msg = "Invalid API key"
        
        return ValidationResult(
            is_valid=False, 
            error=f"Grok validation failed: {error_msg}"
        )


# ------------------------
# Main runner
# ------------------------
async def run_validation_agent(model: str, api_key: str) -> ValidationResult:
    """
    Validates API keys for GPT, Gemini, or Grok.
    Routes to appropriate validator based on model name.
    """
    model_lower = model.lower()
    
    try:
        if "gpt" in model_lower or "o1" in model_lower:
            return await validate_gpt_key(api_key, model)
        elif "gemini" in model_lower:
            return await validate_gemini_key(api_key, model)
        elif "grok" in model_lower:
            return await validate_grok_key(api_key, model)
        else:
            return ValidationResult(
                is_valid=False, 
                error=f"Unsupported model: {model}"
            )
    except Exception as e:
        return ValidationResult(
            is_valid=False, 
            error=f"Validation error: {str(e)}"
        )


# ------------------------
# Optional test runner
# ------------------------
if __name__ == "__main__":
    import os
    from dotenv import load_dotenv
    load_dotenv()

    async def main():
        tests = [
            ("gpt-4o", "OPENAI_API_KEY"),
            ("gemini-2.0-flash-exp", "GEMINI_API_KEY"),
            ("grok-beta", "GROK_API_KEY")
        ]
        
        for model, key_env in tests:
            key = os.getenv(key_env, "")
            if not key:
                print(f"⚠️  {model.upper()}: No API key found in env ({key_env})\n")
                continue
                
            print(f"Testing {model.upper()} key...")
            result = await run_validation_agent(model, key)
            
            if result.is_valid:
                print(f"✅ Valid\n")
            else:
                print(f"❌ Invalid | {result.error}\n")

    asyncio.run(main())