ira/config/config.py

"""
Configuration management for the intelligent research system.

This module handles loading configuration from files and environment variables,
providing secure access to API keys and model settings.
"""

import os
import yaml
from pathlib import Path
from typing import Dict, Any, Optional
from dotenv import load_dotenv

# Load environment variables from .env file if it exists
load_dotenv()

class Config:
    """Configuration manager for the intelligent research system."""
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize the configuration manager.
        
        Args:
            config_path: Path to the configuration file. If None, will look for
                        config.yaml in the same directory as this file.
        """
        self.config_data = {}
        self.config_path = config_path
        
        if not config_path:
            # Default to config.yaml in the same directory as this file
            self.config_path = Path(__file__).parent / "config.yaml"
        
        self.load_config()
    
    def load_config(self) -> None:
        """Load configuration from file if it exists."""
        try:
            if Path(self.config_path).exists():
                with open(self.config_path, 'r') as f:
                    self.config_data = yaml.safe_load(f)
                print(f"Configuration loaded from {self.config_path}")
            else:
                print(f"Configuration file {self.config_path} not found. Using environment variables only.")
        except Exception as e:
            print(f"Error loading configuration: {e}")
    
    def get_api_key(self, provider: str) -> str:
        """Get the API key for a specific provider.
        
        Args:
            provider: The provider name (e.g., 'openai', 'anthropic', 'google')
            
        Returns:
            The API key for the specified provider
            
        Raises:
            ValueError: If the API key is not found
        """
        provider = provider.lower()
        
        # Map provider names to environment variable names
        provider_env_map = {
            'openai': 'OPENAI_API_KEY',
            'anthropic': 'ANTHROPIC_API_KEY',
            'google': 'GEMINI_API_KEY',
            'gemini': 'GEMINI_API_KEY',
            'vertex_ai': 'GOOGLE_APPLICATION_CREDENTIALS',
            'groq': 'GROQ_API_KEY',
            'openrouter': 'OPENROUTER_API_KEY',
            'serper': 'SERPER_API_KEY',
            'tavily': 'TAVILY_API_KEY',
            'perplexity': 'PERPLEXITY_API_KEY'
        }
        
        # Get the environment variable name for the provider
        env_var = provider_env_map.get(provider)
        if not env_var:
            env_var = f"{provider.upper()}_API_KEY"
        
        # Try to get the API key from environment variables
        api_key = os.environ.get(env_var)
        
        # If not found in environment, check the config file
        if not api_key and 'api_keys' in self.config_data:
            api_key = self.config_data['api_keys'].get(provider)
        
        if not api_key:
            raise ValueError(f"API key for {provider} not found. Please set the {env_var} environment variable or add it to the config file.")
        
        return api_key
    
    def get_model_config(self, model_name: str) -> Dict[str, Any]:
        """
        Get configuration for a specific model.
        
        Args:
            model_name: The name of the model
            
        Returns:
            Dictionary containing model configuration
        """
        if self.config_data and 'models' in self.config_data:
            return self.config_data['models'].get(model_name, {})
        return {}
    
    def get_module_model(self, module_name: str, function_name: str) -> str:
        """
        Get the model assigned to a specific module function.
        
        Args:
            module_name: The name of the module (e.g., 'query_processing')
            function_name: The name of the function (e.g., 'enhance_query')
            
        Returns:
            The name of the model to use, or the default model if not specified
        """
        default = self.config_data.get('default_model', 'gpt-3.5-turbo')
        
        if (self.config_data and 'module_models' in self.config_data and 
            module_name in self.config_data['module_models'] and
            function_name in self.config_data['module_models'][module_name]):
            return self.config_data['module_models'][module_name][function_name]
        
        return default
    
    def get_search_config(self, search_engine: str) -> Dict[str, Any]:
        """
        Get configuration for a specific search engine.
        
        Args:
            search_engine: The name of the search engine
            
        Returns:
            Dictionary containing search engine configuration
        """
        if self.config_data and 'search_engines' in self.config_data:
            return self.config_data['search_engines'].get(search_engine, {})
        return {}
    
    def get_ui_config(self) -> Dict[str, Any]:
        """
        Get UI configuration.
        
        Returns:
            Dictionary containing UI configuration
        """
        if self.config_data and 'ui' in self.config_data:
            return self.config_data['ui']
        return {}


# Create a singleton instance for global use
config = Config()


def get_config() -> Config:
    """
    Get the global configuration instance.
    
    Returns:
        The global Config instance
    """
    return config


def get_api_key(service_name: str) -> Optional[str]:
    """
    Get an API key for a specific service.
    
    Args:
        service_name: Name of the service to get the API key for
        
    Returns:
        API key as a string, or None if not found
    """
    # First check environment variables
    env_var_name = f"{service_name.upper()}_API_KEY"
    api_key = os.environ.get(env_var_name)
    
    # If not found in environment, check config file
    if not api_key:
        cfg = get_config()
        api_key = cfg.config_data.get('api_keys', {}).get(service_name)
    
    return api_key