enhanced_llm_code_generator.py

#!/usr/bin/env python3
"""
Enhanced LLM Code Generator with Rate Limiting and Backoff Strategy
Implements exponential backoff, circuit breaker patterns, and parallel execution
"""

import asyncio
import aiohttp
import time
import random
from typing import Dict, Optional, Tuple, List
import json
import os
from dataclasses import dataclass
import logging

# Load environment variables from .env file
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # dotenv not available
except Exception:
    pass  # ignore loading errors

logger = logging.getLogger(__name__)

@dataclass
class APICallResult:
    """Result of an LLM API call"""
    success: bool
    content: str
    time_ms: int
    model_used: str
    error_message: Optional[str] = None
    retry_count: int = 0

class CircuitBreaker:
    """Circuit breaker for API reliability"""
    
    def __init__(self, failure_threshold: int = 5, recovery_timeout: int = 60):
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.failure_count = 0
        self.last_failure_time = None
        self.state = "CLOSED"  # CLOSED, OPEN, HALF_OPEN
        
    def can_execute(self) -> bool:
        """Check if API call can be executed"""
        if self.state == "CLOSED":
            return True
        elif self.state == "OPEN":
            if time.time() - self.last_failure_time > self.recovery_timeout:
                self.state = "HALF_OPEN"
                return True
            return False
        else:  # HALF_OPEN
            return True
            
    def record_success(self):
        """Record successful API call"""
        self.failure_count = 0
        self.state = "CLOSED"
        
    def record_failure(self):
        """Record failed API call"""
        self.failure_count += 1
        self.last_failure_time = time.time()
        
        if self.failure_count >= self.failure_threshold:
            self.state = "OPEN"
            logger.warning(f"Circuit breaker opened after {self.failure_count} failures")

class EnhancedLLMCodeGenerator:
    """Enhanced LLM code generator with reliability improvements"""
    
    def __init__(self):
        self.api_keys = self.load_api_keys()
        self.session: Optional[aiohttp.ClientSession] = None
        self.circuit_breakers = {
            'claude': CircuitBreaker(failure_threshold=3, recovery_timeout=30),
            'gpt': CircuitBreaker(failure_threshold=3, recovery_timeout=30),
            'gemini': CircuitBreaker(failure_threshold=3, recovery_timeout=30)
        }
        self.rate_limiters = {
            'claude': {'last_call': 0, 'min_interval': 1.0},
            'gpt': {'last_call': 0, 'min_interval': 1.0},
            'gemini': {'last_call': 0, 'min_interval': 1.0}
        }
        
    def load_api_keys(self) -> Dict[str, str]:
        """Load API keys from environment variables"""
        keys = {}
        
        # Claude API key
        claude_key = os.getenv('ANTHROPIC_API_KEY')
        if claude_key:
            keys['claude'] = claude_key
            
        # OpenAI API key
        openai_key = os.getenv('OPENAI_API_KEY')
        if openai_key:
            keys['gpt'] = openai_key
            
        # Google API key
        google_key = os.getenv('GOOGLE_API_KEY')
        if google_key:
            keys['gemini'] = google_key
            
        logger.info(f"Loaded {len(keys)} LLM API keys")
        return keys
        
    async def initialize(self):
        """Initialize HTTP session and components"""
        self.session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=60),
            connector=aiohttp.TCPConnector(limit=10, limit_per_host=3)
        )
        logger.info("Enhanced LLM code generator initialized")
        
    async def cleanup(self):
        """Clean up resources"""
        if self.session:
            await self.session.close()
        logger.info("Enhanced LLM code generator cleaned up")
        
    async def apply_rate_limiting(self, model: str):
        """Apply rate limiting for API calls"""
        rate_limiter = self.rate_limiters[model]
        time_since_last = time.time() - rate_limiter['last_call']
        
        if time_since_last < rate_limiter['min_interval']:
            sleep_time = rate_limiter['min_interval'] - time_since_last
            await asyncio.sleep(sleep_time)
            
        rate_limiter['last_call'] = time.time()
        
    async def exponential_backoff_retry(self, model: str, api_call_func, max_retries: int = 4) -> APICallResult:
        """Execute API call with exponential backoff retry"""
        
        for attempt in range(max_retries + 1):
            # Check circuit breaker
            if not self.circuit_breakers[model].can_execute():
                return APICallResult(
                    success=False,
                    content="",
                    time_ms=0,
                    model_used=model,
                    error_message="Circuit breaker open",
                    retry_count=attempt
                )
                
            # Apply rate limiting
            await self.apply_rate_limiting(model)
            
            try:
                start_time = time.time()
                result = await api_call_func()
                execution_time = int((time.time() - start_time) * 1000)
                
                # Success
                self.circuit_breakers[model].record_success()
                return APICallResult(
                    success=True,
                    content=result,
                    time_ms=execution_time,
                    model_used=model,
                    retry_count=attempt
                )
                
            except Exception as e:
                error_msg = str(e)
                self.circuit_breakers[model].record_failure()
                
                # Check if this is a retryable error
                if attempt < max_retries and self.is_retryable_error(error_msg):
                    # Exponential backoff: 1s, 2s, 4s, 8s
                    backoff_time = (2 ** attempt) + random.uniform(0, 1)
                    logger.warning(f"{model} API call failed (attempt {attempt + 1}), retrying in {backoff_time:.1f}s: {error_msg}")
                    await asyncio.sleep(backoff_time)
                    continue
                else:
                    # Non-retryable error or max retries reached
                    execution_time = int((time.time() - start_time) * 1000)
                    return APICallResult(
                        success=False,
                        content="",
                        time_ms=execution_time,
                        model_used=model,
                        error_message=error_msg,
                        retry_count=attempt
                    )
                    
        # Should never reach here
        return APICallResult(
            success=False,
            content="",
            time_ms=0,
            model_used=model,
            error_message="Max retries exceeded",
            retry_count=max_retries
        )
        
    def is_retryable_error(self, error_msg: str) -> bool:
        """Determine if an error is retryable"""
        retryable_patterns = [
            "overloaded",
            "rate limit",
            "timeout",
            "503",
            "502",
            "529",
            "connection",
            "network"
        ]
        
        error_lower = error_msg.lower()
        return any(pattern in error_lower for pattern in retryable_patterns)
        
    async def generate_code_parallel(self, api_name: str, documentation: str, auth_type: str, env_key: str) -> List[APICallResult]:
        """Generate code using all available LLMs in parallel"""
        
        # Create tasks for all available models
        tasks = []
        available_models = []
        
        for model in ['claude', 'gpt', 'gemini']:
            if model in self.api_keys:
                available_models.append(model)
                task = self.generate_with_single_model(model, api_name, documentation, auth_type, env_key)
                tasks.append(task)
                
        if not tasks:
            logger.error("No LLM API keys available")
            return []
            
        logger.info(f"Starting parallel code generation for {api_name} using {len(available_models)} models: {available_models}")
        
        # Execute all models in parallel
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Process results
        api_results = []
        for i, result in enumerate(results):
            model = available_models[i]
            
            if isinstance(result, Exception):
                api_results.append(APICallResult(
                    success=False,
                    content="",
                    time_ms=0,
                    model_used=model,
                    error_message=f"Task exception: {str(result)}",
                    retry_count=0
                ))
            else:
                api_results.append(result)
                
        # Log summary
        successful = len([r for r in api_results if r.success])
        logger.info(f"Parallel generation completed for {api_name}: {successful}/{len(api_results)} models successful")
        
        return api_results
        
    async def generate_with_single_model(self, model: str, api_name: str, documentation: str, auth_type: str, env_key: str) -> APICallResult:
        """Generate code with a single LLM model"""
        
        prompt = self.create_code_generation_prompt(api_name, documentation, auth_type, env_key)
        
        if model == 'claude':
            api_call = lambda: self.call_claude_api(prompt)
        elif model == 'gpt':
            api_call = lambda: self.call_openai_api(prompt)
        elif model == 'gemini':
            api_call = lambda: self.call_gemini_api(prompt)
        else:
            return APICallResult(
                success=False,
                content="",
                time_ms=0,
                model_used=model,
                error_message=f"Unknown model: {model}",
                retry_count=0
            )
            
        return await self.exponential_backoff_retry(model, api_call)
        
    def create_code_generation_prompt(self, api_name: str, documentation: str, auth_type: str, env_key: str) -> str:
        """Create optimized prompt for code generation"""
        
        prompt = f"""Generate Python code to integrate with the {api_name} API based on the following documentation.

DOCUMENTATION:
{documentation}

REQUIREMENTS:
1. Use the environment variable '{env_key}' for authentication
2. Authentication type: {auth_type}
3. Include proper error handling with try/catch blocks
4. Add meaningful comments explaining the code
5. Use the requests library for HTTP calls
6. Include proper imports at the top
7. Make a simple test API call to verify the integration works
8. Handle common HTTP errors (401, 403, 404, 429, 500)

IMPORTANT: Generate ONLY the Python code without any markdown formatting or explanations. The code should be ready to execute.

Python code:"""

        return prompt

    async def call_claude_api(self, prompt: str) -> str:
        """Call Claude API with improved error handling"""

        headers = {
            'Content-Type': 'application/json',
            'x-api-key': self.api_keys['claude'],
            'anthropic-version': '2023-06-01'
        }

        data = {
            'model': 'claude-3-5-sonnet-20241022',
            'max_tokens': 1500,
            'messages': [
                {
                    'role': 'user',
                    'content': prompt
                }
            ]
        }

        async with self.session.post(
            'https://api.anthropic.com/v1/messages',
            headers=headers,
            json=data
        ) as response:
            if response.status == 200:
                result = await response.json()
                return result['content'][0]['text']
            else:
                error_text = await response.text()
                raise Exception(f"Claude API error {response.status}: {error_text}")

    async def call_openai_api(self, prompt: str) -> str:
        """Call OpenAI API with improved error handling"""

        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_keys["gpt"]}'
        }

        data = {
            'model': 'gpt-4',
            'messages': [
                {
                    'role': 'user',
                    'content': prompt
                }
            ],
            'max_tokens': 1500,
            'temperature': 0.1
        }

        async with self.session.post(
            'https://api.openai.com/v1/chat/completions',
            headers=headers,
            json=data
        ) as response:
            if response.status == 200:
                result = await response.json()
                return result['choices'][0]['message']['content']
            else:
                error_text = await response.text()
                raise Exception(f"OpenAI API error {response.status}: {error_text}")

    async def call_gemini_api(self, prompt: str) -> str:
        """Call Google Gemini API with improved error handling"""

        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key={self.api_keys['gemini']}"

        data = {
            'contents': [
                {
                    'parts': [
                        {
                            'text': prompt
                        }
                    ]
                }
            ],
            'generationConfig': {
                'temperature': 0.1,
                'maxOutputTokens': 1500
            }
        }

        async with self.session.post(url, json=data) as response:
            if response.status == 200:
                result = await response.json()
                return result['candidates'][0]['content']['parts'][0]['text']
            else:
                error_text = await response.text()
                raise Exception(f"Gemini API error {response.status}: {error_text}")

# Test function for standalone usage
async def test_enhanced_generator():
    """Test enhanced LLM code generator"""
    generator = EnhancedLLMCodeGenerator()

    try:
        await generator.initialize()

        # Test parallel generation
        sample_documentation = """
        OpenWeatherMap API provides weather data.
        Authentication: Add your API key as 'appid' parameter.
        Endpoint: https://api.openweathermap.org/data/2.5/weather
        Parameters: q (city name), appid (API key)
        """

        results = await generator.generate_code_parallel(
            "OpenWeatherMap",
            sample_documentation,
            "query",
            "OPENWEATHER_API_KEY"
        )

        print(f"\nParallel Generation Results:")
        print("=" * 50)

        for result in results:
            print(f"\nModel: {result.model_used}")
            print(f"Success: {result.success}")
            print(f"Time: {result.time_ms}ms")
            print(f"Retries: {result.retry_count}")
            if result.success:
                print(f"Code length: {len(result.content)}")
                print(f"Code preview: {result.content[:200]}...")
            else:
                print(f"Error: {result.error_message}")

    finally:
        await generator.cleanup()

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