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 API key for the specified provider.
        
        Args:
            provider: The name of the API provider (e.g., 'openai', 'jina', 'serper')
            
        Returns:
            The API key as a string
            
        Raises:
            ValueError: If the API key is not found
        """
        # First check environment variables (higher priority)
        env_var_name = f"{provider.upper()}_API_KEY"
        
        # Special case for Jina AI which uses JINA_API_KEY
        if provider.lower() == 'jina':
            env_var_name = "JINA_API_KEY"
        
        # Special case for Groq which might use GROQ_API_KEY
        if provider.lower() == 'groq':
            env_var_name = "GROQ_API_KEY"
        
        # Special case for OpenRouter which might use OPENROUTER_API_KEY
        if provider.lower() == 'openrouter':
            env_var_name = "OPENROUTER_API_KEY"
        
        # Special case for Google which might use GEMINI_API_KEY
        if provider.lower() == 'google':
            env_var_name = "GEMINI_API_KEY"
        
        api_key = os.environ.get(env_var_name)
        
        # If not in environment, check config file
        if not api_key and self.config_data 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. Set {env_var_name} environment variable or add to 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