openenv/core/config.py

"""
Configuration module for OpenEnv Email Triage environment parameters.
"""

from pydantic import BaseModel, Field
from typing import Optional, Dict, Any, Tuple
import json

class EnvConfig(BaseModel):
    """
    Configuration class for OpenEnv Email Triage environment.
    """
    # Task difficulty
    task_level: str = Field(default="medium", description="Task level: 'easy', 'medium', or 'hard'")
    
    # Environment dynamics
    num_emails: int = Field(default=20, description="Total number of emails to triage")
    spam_ratio: float = Field(default=0.3, description="Ratio of emails that are spam")
    urgent_ratio: float = Field(default=0.2, description="Ratio of emails that are urgent")
    confounding_ratio: float = Field(default=0.1, description="Ratio of confusing/nuanced emails (medium/hard only)")
    
    # Reward configuration
    reward_scale: float = Field(default=1.0)
    
    # Rendering options
    render_mode: Optional[str] = Field(default=None)
    render_fps: int = Field(default=60)
    screen_size: Tuple[int, int] = Field(default=(1024, 768))
    
    # Logging
    verbose: bool = Field(default=True)
    random_seed: Optional[int] = Field(default=None)
    
    custom_params: Dict[str, Any] = Field(default_factory=dict)
    
    def validate(self) -> bool:
        if self.num_emails <= 0:
            raise ValueError("num_emails must be positive")
        if not (0.0 <= self.spam_ratio <= 1.0):
            raise ValueError("spam_ratio must be between 0 and 1")
        if self.task_level not in ["easy", "medium", "hard"]:
            raise ValueError(f"Unknown task level: {self.task_level}")
        return True
    
    def to_dict(self) -> Dict[str, Any]:
        return self.model_dump()
    
    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> 'EnvConfig':
        return cls(**config_dict)
    
    def save(self, filepath: str) -> None:
        with open(filepath, 'w') as f:
            json.dump(self.to_dict(), f, indent=2)
    
    @staticmethod
    def load(filepath: str) -> 'EnvConfig':
        with open(filepath, 'r') as f:
            config_dict = json.load(f)
        return EnvConfig.from_dict(config_dict)