llm_code_generator.py

#!/usr/bin/env python3
"""
LLM Code Generator for API Research Study
Generates Python code using multiple LLM providers
"""

import asyncio
import aiohttp
import json
import os
import time
from typing import Dict, Optional, Tuple
import logging

logger = logging.getLogger(__name__)

class LLMCodeGenerator:
    """Generates API integration code using multiple LLM providers"""
    
    def __init__(self):
        self.session: Optional[aiohttp.ClientSession] = None
        self.api_keys = self.load_api_keys()
        
    def load_api_keys(self) -> Dict[str, str]:
        """Load LLM API keys from environment"""
        keys = {}
        
        # Load from .env file
        try:
            with open('.env', 'r') as f:
                for line in f:
                    if '=' in line and not line.startswith('#'):
                        key, value = line.strip().split('=', 1)
                        if key in ['ANTHROPIC_API_KEY', 'OPENAI_API_KEY', 'GOOGLE_API_KEY']:
                            keys[key] = value
        except Exception as e:
            logger.error(f"❌ Failed to load API keys: {e}")
            
        logger.info(f"🔑 Loaded {len(keys)} LLM API keys")
        return keys
        
    async def initialize(self):
        """Initialize HTTP session"""
        self.session = aiohttp.ClientSession()
        logger.info("✅ LLM code generator initialized")
        
    async def generate_code(self, api_name: str, documentation: str, auth_type: str, env_key: str) -> Tuple[bool, str, int, str]:
        """
        Generate Python code for API integration
        
        Returns:
            Tuple of (success, generated_code, generation_time_ms, model_used)
        """
        
        # Try each LLM model
        models = ['claude', 'gpt', 'gemini']
        
        for model in models:
            try:
                success, code, time_ms = await self.generate_with_model(
                    model, api_name, documentation, auth_type, env_key
                )
                if success:
                    return True, code, time_ms, model
            except Exception as e:
                logger.warning(f"⚠️ {model} failed: {e}")
                continue
                
        return False, "", 0, "none"
        
    async def generate_with_model(self, model: str, api_name: str, documentation: str, auth_type: str, env_key: str) -> Tuple[bool, str, int]:
        """Generate code using specific LLM model"""
        
        start_time = time.time()
        
        prompt = self.create_standardized_prompt(api_name, documentation, auth_type, env_key)
        
        if model == 'claude':
            success, code = await self.call_claude(prompt)
        elif model == 'gpt':
            success, code = await self.call_openai(prompt)
        elif model == 'gemini':
            success, code = await self.call_gemini(prompt)
        else:
            return False, "", 0
            
        generation_time = int((time.time() - start_time) * 1000)
        
        if success:
            logger.info(f"✅ {model} generated {len(code)} characters for {api_name} ({generation_time}ms)")
        else:
            logger.warning(f"⚠️ {model} failed to generate code for {api_name}")
            
        return success, code, generation_time
        
    def create_standardized_prompt(self, api_name: str, documentation: str, auth_type: str, env_key: str) -> str:
        """Create standardized prompt for fair comparison across LLMs"""
        
        auth_instructions = {
            'query': f"Use the API key from environment variable {env_key} as a query parameter",
            'header': f"Use the API key from environment variable {env_key} in the request headers",
            'path': f"Use the API key from environment variable {env_key} in the URL path",
            'bearer': f"Use the API key from environment variable {env_key} as a Bearer token in Authorization header",
            'basic': f"Use the credentials from environment variables for HTTP Basic authentication"
        }
        
        prompt = f"""
You are tasked with creating Python code to integrate with the {api_name} API based on the following documentation.

DOCUMENTATION:
{documentation}

REQUIREMENTS:
1. Create a complete Python script that makes a simple API call to {api_name}
2. {auth_instructions.get(auth_type, 'Use appropriate authentication method')}
3. Include proper error handling and response parsing
4. Use the requests library for HTTP calls
5. Include necessary imports
6. Add comments explaining the code
7. Make the code production-ready with proper exception handling
8. Test with a simple, safe API endpoint (like getting current data or basic info)

AUTHENTICATION TYPE: {auth_type}
ENVIRONMENT VARIABLE: {env_key}

Please provide ONLY the Python code without any explanations or markdown formatting.
The code should be ready to run as-is.
"""
        
        return prompt.strip()
        
    async def call_claude(self, prompt: str) -> Tuple[bool, str]:
        """Call Anthropic Claude API"""
        
        if 'ANTHROPIC_API_KEY' not in self.api_keys:
            return False, "No Anthropic API key"
            
        headers = {
            'Content-Type': 'application/json',
            'x-api-key': self.api_keys['ANTHROPIC_API_KEY'],
            'anthropic-version': '2023-06-01'
        }
        
        data = {
            'model': 'claude-3-5-sonnet-20241022',
            'max_tokens': 1000,
            'messages': [
                {
                    'role': 'user',
                    'content': prompt
                }
            ]
        }
        
        try:
            async with self.session.post(
                'https://api.anthropic.com/v1/messages',
                headers=headers,
                json=data,
                timeout=30
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    code = result['content'][0]['text'].strip()
                    return True, code
                else:
                    error_text = await response.text()
                    logger.error(f"Claude API error {response.status}: {error_text}")
                    return False, f"API Error: {response.status}"
                    
        except Exception as e:
            logger.error(f"Claude request failed: {e}")
            return False, f"Request failed: {str(e)}"
            
    async def call_openai(self, prompt: str) -> Tuple[bool, str]:
        """Call OpenAI GPT API"""
        
        if 'OPENAI_API_KEY' not in self.api_keys:
            return False, "No OpenAI API key"
            
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_keys["OPENAI_API_KEY"]}'
        }
        
        data = {
            'model': 'gpt-4',
            'messages': [
                {
                    'role': 'user',
                    'content': prompt
                }
            ],
            'max_tokens': 1000,
            'temperature': 0.1
        }
        
        try:
            async with self.session.post(
                'https://api.openai.com/v1/chat/completions',
                headers=headers,
                json=data,
                timeout=30
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    code = result['choices'][0]['message']['content'].strip()
                    return True, code
                else:
                    error_text = await response.text()
                    logger.error(f"OpenAI API error {response.status}: {error_text}")
                    return False, f"API Error: {response.status}"
                    
        except Exception as e:
            logger.error(f"OpenAI request failed: {e}")
            return False, f"Request failed: {str(e)}"
            
    async def call_gemini(self, prompt: str) -> Tuple[bool, str]:
        """Call Google Gemini API"""
        
        if 'GOOGLE_API_KEY' not in self.api_keys:
            return False, "No Google API key"
            
        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key={self.api_keys['GOOGLE_API_KEY']}"
        
        data = {
            'contents': [
                {
                    'parts': [
                        {
                            'text': prompt
                        }
                    ]
                }
            ],
            'generationConfig': {
                'temperature': 0.1,
                'maxOutputTokens': 1000
            }
        }
        
        try:
            async with self.session.post(
                url,
                json=data,
                timeout=30
            ) as response:
                if response.status == 200:
                    result = await response.json()
                    code = result['candidates'][0]['content']['parts'][0]['text'].strip()
                    return True, code
                else:
                    error_text = await response.text()
                    logger.error(f"Gemini API error {response.status}: {error_text}")
                    return False, f"API Error: {response.status}"
                    
        except Exception as e:
            logger.error(f"Gemini request failed: {e}")
            return False, f"Request failed: {str(e)}"
            
    async def cleanup(self):
        """Clean up HTTP session"""
        if self.session:
            await self.session.close()
            logger.info("🧹 LLM code generator cleaned up")

# Test function
async def test_code_generation():
    """Test code generation with sample documentation"""
    generator = LLMCodeGenerator()
    
    try:
        await generator.initialize()
        
        sample_doc = """
        # Weather API Documentation
        
        Get current weather data for any location.
        
        ## Authentication
        Use your API key as a query parameter: ?appid=YOUR_API_KEY
        
        ## Endpoint
        GET https://api.openweathermap.org/data/2.5/weather?q={city}&appid={API_KEY}
        
        ## Example Response
        {
            "weather": [{"main": "Clear", "description": "clear sky"}],
            "main": {"temp": 280.32, "pressure": 1012},
            "name": "London"
        }
        """
        
        success, code, time_ms, model = await generator.generate_code(
            "OpenWeatherMap", sample_doc, "query", "OPENWEATHER_API_KEY"
        )
        
        print(f"Success: {success}")
        print(f"Model: {model}")
        print(f"Time: {time_ms}ms")
        print(f"Generated code:\n{code}")
        
    finally:
        await generator.cleanup()

if __name__ == "__main__":
    asyncio.run(test_code_generation())