src/config.py

import yaml
import os
from typing import List, Dict, Any

def load_config() -> Dict[str, Any]:
    # Default configuration
    config = {
        'queue_repo': 'stacklok/requests',
        'eval_requests_path': './eval-queue',
        'eval_results_path': './eval-results',
        'allowed_weight_types': ['Safetensors', 'PyTorch', 'GGUF', 'Other'],
        'default_revision': 'main',
        'log_level': 'INFO',
        'evaluation_wait_time': 60,
        'is_local': False  # Default to non-local (Hugging Face) environment
    }

    # Load from config.yaml if it exists
    if os.path.exists('config.yaml'):
        with open('config.yaml', 'r') as file:
            yaml_config = yaml.safe_load(file)
            config.update(yaml_config)

    # Override with environment variables if they exist
    env_vars = {
        'API_TOKEN': 'api_token',
        'QUEUE_REPO': 'queue_repo',
        'EVAL_REQUESTS_PATH': 'eval_requests_path',
        'EVAL_RESULTS_PATH': 'eval_results_path',
        'LOG_LEVEL': 'log_level',
        'EVALUATION_WAIT_TIME': 'evaluation_wait_time',
        'IS_LOCAL': 'is_local'
    }

    for env_var, config_key in env_vars.items():
        if os.environ.get(env_var):
            config[config_key] = os.environ[env_var]

    # Convert IS_LOCAL to boolean
    is_local = config.get('is_local', False)
    if isinstance(is_local, str):
        config['is_local'] = is_local.lower() == 'true'
    else:
        config['is_local'] = bool(is_local)

    # Ensure API_TOKEN is set only for local runs
    if config['is_local']:
        if 'api_token' not in config or not config['api_token']:
            raise ValueError("API_TOKEN must be set in environment variables or config.yaml for local runs")
    else:
        # Remove API_TOKEN for non-local runs
        config.pop('api_token', None)

    return config

config = load_config()

# API and Token configurations (only for local runs)
API_TOKEN: str = config.get('api_token', '')
QUEUE_REPO: str = config['queue_repo']

# File paths
EVAL_REQUESTS_PATH: str = config['eval_requests_path']
EVAL_RESULTS_PATH: str = config['eval_results_path']

# Model configurations
ALLOWED_WEIGHT_TYPES: List[str] = config['allowed_weight_types']
DEFAULT_REVISION: str = config['default_revision']

# Logging configuration
LOG_LEVEL: str = config['log_level']

# Evaluation configurations
EVALUATION_WAIT_TIME: int = config['evaluation_wait_time']

# Environment type
IS_LOCAL: bool = config['is_local']