Create ai_resilience.py

ai_resilience.py

@@ -0,0 +1,208 @@
+# File: ai_resilience.py
+# Location: /ai_resilience.py
+# Description: Robust AI API error handling and fallback mechanism
+
+import logging
+from typing import List, Dict, Any, Callable
+import openai
+import anthropic
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    retry_if_exception_type
+)
+
+class AIResilientClient:
+    def __init__(self, 
+                 openai_key: str = None, 
+                 anthropic_key: str = None,
+                 max_retries: int = 3):
+        """
+        Initialize resilient AI client with multiple model support
+        
+        Args:
+            openai_key (str): OpenAI API key
+            anthropic_key (str): Anthropic API key
+            max_retries (int): Maximum retry attempts for API calls
+        """
+        # Configure logging
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(levelname)s - %(message)s'
+        )
+        self.logger = logging.getLogger(__name__)
+        
+        # Initialize API clients
+        self.openai_client = openai.OpenAI(api_key=openai_key) if openai_key else None
+        self.anthropic_client = anthropic.Anthropic(api_key=anthropic_key) if anthropic_key else None
+        
+        # Fallback models
+        self.models = [
+            {
+                'provider': 'openai',
+                'model': 'gpt-4-turbo',
+                'temperature': 0.7,
+                'max_tokens': 4096
+            },
+            {
+                'provider': 'anthropic',
+                'model': 'claude-3-opus-20240229',
+                'temperature': 0.7,
+                'max_tokens': 4096
+            },
+            {
+                'provider': 'openai',
+                'model': 'gpt-3.5-turbo',
+                'temperature': 0.7,
+                'max_tokens': 4096
+            }
+        ]
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type(
+            (openai.APIError, 
+             openai.APIConnectionError, 
+             openai.RateLimitError,
+             anthropic.APIError)
+        )
+    )
+    def generate_content(
+        self, 
+        prompt: str, 
+        context: Dict[str, Any] = None
+    ) -> str:
+        """
+        Generate content with intelligent fallback and error handling
+        
+        Args:
+            prompt (str): Content generation prompt
+            context (Dict, optional): Additional context for generation
+        
+        Returns:
+            str: Generated content
+        """
+        # Attempt generation with each configured model
+        for model_config in self.models:
+            try:
+                return self._generate_with_model(model_config, prompt, context)
+            except Exception as e:
+                self.logger.warning(
+                    f"Failed with {model_config['provider']} {model_config['model']}: {str(e)}"
+                )
+        
+        # If all models fail, raise a critical error
+        error_msg = "All AI models failed. Unable to generate content."
+        self.logger.critical(error_msg)
+        raise RuntimeError(error_msg)
+    
+    def _generate_with_model(
+        self, 
+        model_config: Dict[str, Any], 
+        prompt: str, 
+        context: Dict[str, Any] = None
+    ) -> str:
+        """
+        Generate content with a specific model
+        
+        Args:
+            model_config (Dict): Model configuration
+            prompt (str): Content generation prompt
+            context (Dict, optional): Additional context
+        
+        Returns:
+            str: Generated content
+        """
+        provider = model_config['provider']
+        
+        if provider == 'openai' and self.openai_client:
+            response = self.openai_client.chat.completions.create(
+                model=model_config['model'],
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": prompt}
+                ],
+                temperature=model_config.get('temperature', 0.7),
+                max_tokens=model_config.get('max_tokens', 4096)
+            )
+            return response.choices[0].message.content
+        
+        elif provider == 'anthropic' and self.anthropic_client:
+            response = self.anthropic_client.messages.create(
+                model=model_config['model'],
+                max_tokens=model_config.get('max_tokens', 4096),
+                temperature=model_config.get('temperature', 0.7),
+                messages=[
+                    {"role": "user", "content": prompt}
+                ]
+            )
+            return response.content[0].text
+        
+        raise ValueError(f"No valid client for provider: {provider}")
+    
+    def validate_api_keys(self) -> Dict[str, bool]:
+        """
+        Validate API keys for configured providers
+        
+        Returns:
+            Dict with validation status for each provider
+        """
+        results = {
+            'openai': False,
+            'anthropic': False
+        }
+        
+        # Test OpenAI
+        if self.openai_client:
+            try:
+                self.openai_client.models.list()
+                results['openai'] = True
+            except Exception:
+                pass
+        
+        # Test Anthropic
+        if self.anthropic_client:
+            try:
+                self.anthropic_client.models.list()
+                results['anthropic'] = True
+            except Exception:
+                pass
+        
+        return results
+
+# Comprehensive Test Suite
+def run_resilience_tests():
+    """
+    Run comprehensive tests for AI resilience module
+    """
+    # Initialize with mock API keys (replace with actual keys for real testing)
+    client = AIResilientClient(
+        openai_key=os.getenv('OPENAI_API_KEY'),
+        anthropic_key=os.getenv('ANTHROPIC_API_KEY')
+    )
+    
+    # Test API Key Validation
+    api_status = client.validate_api_keys()
+    print("API Key Validation:", api_status)
+    
+    # Test Content Generation
+    test_prompts = [
+        "Write a short paragraph about artificial intelligence.",
+        "Explain the concept of quantum computing in simple terms."
+    ]
+    
+    for prompt in test_prompts:
+        try:
+            generated_content = client.generate_content(prompt)
+            print(f"\nGenerated Content for Prompt: {prompt}")
+            print(generated_content)
+        except Exception as e:
+            print(f"Content generation failed: {e}")
+    
+    print("\nResilience tests completed.")
+
+# Run tests if script is executed directly
+if __name__ == "__main__":
+    run_resilience_tests()