llm_client.py

#!/usr/bin/env python3
"""
IQKiller Multi-Provider LLM Client
Supports OpenAI GPT-4o-mini (primary) and Anthropic Claude-3.5-Sonnet (fallback)
Enterprise-grade with retries, timeouts, and cost optimization
"""

import asyncio
import logging
import time
from typing import Optional, Dict, Any, List, Tuple, Union
from dataclasses import dataclass
from enum import Enum

# Third-party imports
import openai
import anthropic
from openai import AsyncOpenAI
from anthropic import AsyncAnthropic

# Local imports
from config import get_config, IQKillerConfig

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class LLMProvider(Enum):
    """Supported LLM providers"""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"

@dataclass
class LLMResponse:
    """Standardized LLM response format"""
    content: str
    provider: LLMProvider
    model: str
    usage: Dict[str, Any]
    processing_time: float
    cost_estimate: float

@dataclass
class LLMRequest:
    """Standardized LLM request format"""
    prompt: str
    system_prompt: Optional[str] = None
    temperature: float = 0.1
    max_tokens: int = 2000
    model_override: Optional[str] = None

class LLMClient:
    """Multi-provider LLM client with fallback support"""
    
    def __init__(self, config: Optional[IQKillerConfig] = None):
        """Initialize LLM client with configuration"""
        self.config = config or get_config()
        
        # Initialize clients
        self.openai_client: Optional[AsyncOpenAI] = None
        self.anthropic_client: Optional[AsyncAnthropic] = None
        
        # Model configurations
        self.openai_config = {
            "model": "gpt-4o-mini",
            "temperature": 0.1,
            "max_tokens": 2000,
        }
        
        self.anthropic_config = {
            "model": "claude-3-5-sonnet-20241022",
            "temperature": 0.1,
            "max_tokens": 2000,
        }
        
        # Cost estimates per 1K tokens (approximate)
        self.cost_estimates = {
            "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
            "claude-3-5-sonnet-20241022": {"input": 0.003, "output": 0.015}
        }
        
        # Initialize available clients
        self._init_clients()
    
    def _init_clients(self) -> None:
        """Initialize API clients based on available keys"""
        
        # Initialize OpenAI client
        if self.config.openai_api_key:
            try:
                self.openai_client = AsyncOpenAI(api_key=self.config.openai_api_key)
                logger.info("✅ OpenAI client initialized")
            except Exception as e:
                logger.error(f"❌ Failed to initialize OpenAI client: {e}")
        else:
            logger.warning("⚠️  OpenAI API key not provided")
        
        # Initialize Anthropic client
        if self.config.anthropic_api_key:
            try:
                self.anthropic_client = AsyncAnthropic(api_key=self.config.anthropic_api_key)
                logger.info("✅ Anthropic client initialized")
            except Exception as e:
                logger.error(f"❌ Failed to initialize Anthropic client: {e}")
        else:
            logger.warning("⚠️  Anthropic API key not provided (fallback unavailable)")
    
    def get_available_providers(self) -> List[LLMProvider]:
        """Get list of available providers"""
        providers = []
        
        if self.openai_client:
            providers.append(LLMProvider.OPENAI)
        
        if self.anthropic_client:
            providers.append(LLMProvider.ANTHROPIC)
        
        return providers
    
    def estimate_cost(self, prompt: str, response: str, model: str) -> float:
        """Estimate cost for a request/response pair"""
        
        # Simple token estimation (rough approximation)
        input_tokens = len(prompt.split()) * 1.3  # ~1.3 tokens per word
        output_tokens = len(response.split()) * 1.3
        
        if model in self.cost_estimates:
            cost_config = self.cost_estimates[model]
            total_cost = (
                (input_tokens / 1000) * cost_config["input"] + 
                (output_tokens / 1000) * cost_config["output"]
            )
            return round(total_cost, 6)
        
        return 0.0
    
    async def _call_openai(self, request: LLMRequest) -> LLMResponse:
        """Call OpenAI API"""
        
        if not self.openai_client:
            raise Exception("OpenAI client not available")
        
        start_time = time.time()
        
        # Prepare messages
        messages = []
        if request.system_prompt:
            messages.append({"role": "system", "content": request.system_prompt})
        messages.append({"role": "user", "content": request.prompt})
        
        # Get model
        model = request.model_override or self.openai_config["model"]
        
        try:
            response = await self.openai_client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=request.temperature,
                max_tokens=request.max_tokens,
                timeout=self.config.request_timeout
            )
            
            processing_time = time.time() - start_time
            content = response.choices[0].message.content or ""
            
            # Extract usage info
            usage = {
                "prompt_tokens": response.usage.prompt_tokens if response.usage else 0,
                "completion_tokens": response.usage.completion_tokens if response.usage else 0,
                "total_tokens": response.usage.total_tokens if response.usage else 0
            }
            
            # Estimate cost
            cost = self.estimate_cost(request.prompt, content, model)
            
            return LLMResponse(
                content=content,
                provider=LLMProvider.OPENAI,
                model=model,
                usage=usage,
                processing_time=processing_time,
                cost_estimate=cost
            )
            
        except Exception as e:
            logger.error(f"❌ OpenAI API call failed: {e}")
            raise
    
    async def _call_anthropic(self, request: LLMRequest) -> LLMResponse:
        """Call Anthropic API"""
        
        if not self.anthropic_client:
            raise Exception("Anthropic client not available")
        
        start_time = time.time()
        
        # Get model
        model = request.model_override or self.anthropic_config["model"]
        
        try:
            # Prepare message for Claude
            message_content = request.prompt
            if request.system_prompt:
                message_content = f"System: {request.system_prompt}\n\nUser: {request.prompt}"
            
            message = await self.anthropic_client.messages.create(
                model=model,
                max_tokens=request.max_tokens,
                temperature=request.temperature,
                messages=[{"role": "user", "content": message_content}],
                timeout=self.config.request_timeout
            )
            
            processing_time = time.time() - start_time
            # Extract content from message
            content = ""
            if hasattr(message, 'content') and message.content:
                # Anthropic returns content as a list of blocks, typically text blocks
                try:
                    content = message.content[0].text if message.content else ""
                except (IndexError, AttributeError):
                    content = str(message.content) if message.content else ""
            
            # Extract usage info
            usage = {
                "prompt_tokens": message.usage.input_tokens if hasattr(message, 'usage') else 0,
                "completion_tokens": message.usage.output_tokens if hasattr(message, 'usage') else 0,
                "total_tokens": (message.usage.input_tokens + message.usage.output_tokens) if hasattr(message, 'usage') else 0
            }
            
            # Estimate cost
            cost = self.estimate_cost(request.prompt, content, model)
            
            return LLMResponse(
                content=content,
                provider=LLMProvider.ANTHROPIC,
                model=model,
                usage=usage,
                processing_time=processing_time,
                cost_estimate=cost
            )
            
        except Exception as e:
            logger.error(f"❌ Anthropic API call failed: {e}")
            raise
    
    async def generate(
        self, 
        request: LLMRequest,
        preferred_provider: Optional[LLMProvider] = None,
        use_fallback: bool = True
    ) -> LLMResponse:
        """Generate response using primary provider with fallback support"""
        
        available_providers = self.get_available_providers()
        
        if not available_providers:
            raise Exception("❌ No LLM providers available. Please check API keys.")
        
        # Determine provider order
        if preferred_provider and preferred_provider in available_providers:
            primary_provider = preferred_provider
            fallback_providers = [p for p in available_providers if p != preferred_provider]
        else:
            # Default order: OpenAI first, then Anthropic
            primary_provider = LLMProvider.OPENAI if LLMProvider.OPENAI in available_providers else available_providers[0]
            fallback_providers = [p for p in available_providers if p != primary_provider]
        
        # Try primary provider
        for attempt in range(self.config.retry_attempts):
            try:
                logger.info(f"🔄 Attempt {attempt + 1}: Trying {primary_provider.value}")
                
                if primary_provider == LLMProvider.OPENAI:
                    return await self._call_openai(request)
                elif primary_provider == LLMProvider.ANTHROPIC:
                    return await self._call_anthropic(request)
                
            except Exception as e:
                logger.warning(f"⚠️  {primary_provider.value} attempt {attempt + 1} failed: {e}")
                
                if attempt < self.config.retry_attempts - 1:
                    await asyncio.sleep(2 ** attempt)  # Exponential backoff
                    continue
                else:
                    logger.error(f"❌ {primary_provider.value} failed after {self.config.retry_attempts} attempts")
                    break
        
        # Try fallback providers if enabled
        if use_fallback and fallback_providers:
            for fallback_provider in fallback_providers:
                logger.info(f"🔄 Trying fallback provider: {fallback_provider.value}")
                
                try:
                    if fallback_provider == LLMProvider.OPENAI:
                        return await self._call_openai(request)
                    elif fallback_provider == LLMProvider.ANTHROPIC:
                        return await self._call_anthropic(request)
                        
                except Exception as e:
                    logger.warning(f"⚠️  Fallback {fallback_provider.value} failed: {e}")
                    continue
        
        # All providers failed
        raise Exception("❌ All LLM providers failed. Please check your API keys and network connection.")
    
    async def generate_simple(
        self, 
        prompt: str, 
        system_prompt: Optional[str] = None,
        temperature: float = 0.1,
        max_tokens: int = 2000
    ) -> str:
        """Simple interface for quick generation"""
        
        request = LLMRequest(
            prompt=prompt,
            system_prompt=system_prompt,
            temperature=temperature,
            max_tokens=max_tokens
        )
        
        response = await self.generate(request)
        return response.content
    
    def get_status(self) -> Dict[str, Any]:
        """Get client status information"""
        
        available_providers = self.get_available_providers()
        
        return {
            "available_providers": [p.value for p in available_providers],
            "primary_provider": "openai" if LLMProvider.OPENAI in available_providers else (
                "anthropic" if LLMProvider.ANTHROPIC in available_providers else "none"
            ),
            "fallback_available": len(available_providers) > 1,
            "openai_available": LLMProvider.OPENAI in available_providers,
            "anthropic_available": LLMProvider.ANTHROPIC in available_providers,
            "config": {
                "request_timeout": self.config.request_timeout,
                "retry_attempts": self.config.retry_attempts,
                "max_analysis_time": self.config.max_analysis_time
            }
        }

# Global client instance
_llm_client: Optional[LLMClient] = None

def get_llm_client() -> LLMClient:
    """Get global LLM client instance (singleton pattern)"""
    global _llm_client
    if _llm_client is None:
        _llm_client = LLMClient()
    return _llm_client

def reload_llm_client() -> LLMClient:
    """Reload LLM client with fresh configuration"""
    global _llm_client
    _llm_client = LLMClient()
    return _llm_client

# Convenience functions for common use cases
async def generate_interview_content(prompt: str, system_prompt: Optional[str] = None) -> str:
    """Generate interview-related content"""
    client = get_llm_client()
    return await client.generate_simple(prompt, system_prompt, temperature=0.1)

async def generate_analysis_content(prompt: str, system_prompt: Optional[str] = None) -> str:
    """Generate analysis content with slightly higher creativity"""
    client = get_llm_client()
    return await client.generate_simple(prompt, system_prompt, temperature=0.2)

async def generate_creative_content(prompt: str, system_prompt: Optional[str] = None) -> str:
    """Generate creative content like salary scenarios"""
    client = get_llm_client()
    return await client.generate_simple(prompt, system_prompt, temperature=0.3)

if __name__ == "__main__":
    async def test_llm_client():
        """Test the LLM client"""
        client = LLMClient()
        
        print("🧪 Testing LLM Client")
        print("=" * 50)
        
        # Print status
        status = client.get_status()
        print("📊 Client Status:")
        for key, value in status.items():
            print(f"  {key}: {value}")
        
        # Test simple generation if providers available
        if status["available_providers"]:
            print("\n🔄 Testing simple generation...")
            try:
                response = await client.generate_simple(
                    "What are the top 3 skills for a software engineer?",
                    "You are a helpful career advisor."
                )
                print(f"✅ Response: {response[:100]}...")
            except Exception as e:
                print(f"❌ Test failed: {e}")
        else:
            print("⚠️  No providers available for testing")
        
        print("=" * 50)
    
    # Run test
    asyncio.run(test_llm_client())