config.py

import openai
import anthropic
from typing import Dict, Tuple, Optional
from abc import ABC, abstractmethod

class ModelProvider(ABC):
    """Abstract base class for model providers."""
    
    @abstractmethod
    def setup_client(self, api_key: str) -> Tuple[bool, str]:
        """Setup the API client with the provided key."""
        pass
    
    @abstractmethod
    def generate_response(self, prompt: str, model: str, max_tokens: int = 800) -> str:
        """Generate a response using the specified model."""
        pass
    
    @abstractmethod
    def get_available_models(self) -> list:
        """Return list of available models for this provider."""
        pass

class OpenAIProvider(ModelProvider):
    """OpenAI API provider implementation."""
    
    def __init__(self):
        self.client = None
        self.models = [
            "gpt-3.5-turbo",
            "gpt-4",
            "gpt-4-turbo",
            "gpt-4o",
            "gpt-4o-mini"
        ]
    
    def setup_client(self, api_key: str) -> Tuple[bool, str]:
        """Configure OpenAI client with provided API key."""
        if not api_key.strip():
            return False, "OpenAI API key cannot be empty"
        
        try:
            self.client = openai.OpenAI(api_key=api_key.strip())
            # Test the connection
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": "Hello"}],
                max_tokens=5
            )
            return True, "OpenAI client configured successfully"
        except Exception as e:
            return False, f"Failed to configure OpenAI client: {str(e)}"
    
    def generate_response(self, prompt: str, model: str, max_tokens: int = 800) -> str:
        """Generate response using OpenAI models."""
        if self.client is None:
            raise Exception("OpenAI client not configured")
        
        response = self.client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a precise mathematician who always provides clear, step-by-step solutions and selects the correct answer from given options."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=max_tokens,
            temperature=0.0
        )
        
        return response.choices[0].message.content
    
    def get_available_models(self) -> list:
        """Return available OpenAI models."""
        return self.models

class ClaudeProvider(ModelProvider):
    """Anthropic Claude API provider implementation."""
    
    def __init__(self):
        self.client = None
        self.models = [
            "claude-3-haiku-20240307",
            "claude-3-sonnet-20240229", 
            "claude-3-opus-20240229",
            "claude-3-5-sonnet-20241022",
            "claude-3-5-haiku-20241022",
            "claude-sonnet-4-20250514",
            "claude-opus-4-20250514",
            "claude-opus-4-1-20250805"
        ]
    
    def setup_client(self, api_key: str) -> Tuple[bool, str]:
        """Configure Anthropic client with provided API key."""
        if not api_key.strip():
            return False, "Anthropic API key cannot be empty"
        
        try:
            self.client = anthropic.Anthropic(api_key=api_key.strip())
            # Test the connection
            response = self.client.messages.create(
                model="claude-3-haiku-20240307",
                max_tokens=5,
                messages=[{"role": "user", "content": "Hello"}]
            )
            return True, "Claude client configured successfully"
        except Exception as e:
            return False, f"Failed to configure Claude client: {str(e)}"
    
    def generate_response(self, prompt: str, model: str, max_tokens: int = 800) -> str:
        """Generate response using Claude models."""
        if self.client is None:
            raise Exception("Claude client not configured")
        
        # Add system prompt for Claude
        system_prompt = "You are a precise mathematician who always provides clear, step-by-step solutions and selects the correct answer from given options."
        
        response = self.client.messages.create(
            model=model,
            max_tokens=max_tokens,
            system=system_prompt,
            messages=[{"role": "user", "content": prompt}]
        )
        
        return response.content[0].text
    
    def get_available_models(self) -> list:
        """Return available Claude models."""
        return self.models

class ModelManager:
    """Manages multiple model providers and routing."""
    
    def __init__(self):
        self.providers = {
            "openai": OpenAIProvider(),
            "claude": ClaudeProvider()
        }
        self.configured_providers = set()
    
    def setup_provider(self, provider_name: str, api_key: str) -> Tuple[bool, str]:
        """Setup a specific provider with API key."""
        if provider_name not in self.providers:
            return False, f"Unknown provider: {provider_name}"
        
        success, message = self.providers[provider_name].setup_client(api_key)
        
        if success:
            self.configured_providers.add(provider_name)
        else:
            self.configured_providers.discard(provider_name)
        
        return success, message
    
    def get_provider_for_model(self, model: str) -> Optional[str]:
        """Determine which provider handles the given model."""
        for provider_name, provider in self.providers.items():
            if model in provider.get_available_models():
                return provider_name
        return None
    
    def generate_response(self, prompt: str, model: str, max_tokens: int = 800) -> str:
        """Generate response using the appropriate provider for the model."""
        provider_name = self.get_provider_for_model(model)
        
        if not provider_name:
            raise Exception(f"No provider found for model: {model}")
        
        if provider_name not in self.configured_providers:
            raise Exception(f"Provider {provider_name} not configured")
        
        return self.providers[provider_name].generate_response(prompt, model, max_tokens)
    
    def get_all_models(self) -> Dict[str, list]:
        """Get all available models grouped by provider."""
        return {
            provider_name: provider.get_available_models() 
            for provider_name, provider in self.providers.items()
        }
    
    def get_flat_model_list(self) -> list:
        """Get a flat list of all available models."""
        models = []
        for provider in self.providers.values():
            models.extend(provider.get_available_models())
        return models
    
    def is_provider_configured(self, provider_name: str) -> bool:
        """Check if a provider is configured."""
        return provider_name in self.configured_providers
    
    def get_configured_providers(self) -> list:
        """Get list of configured providers."""
        return list(self.configured_providers)