server/tasks/base.py

"""Base task class."""

import random
from typing import Dict, Optional

from server.models import TaskDifficulty


class BaseTask:
    NAME = "Base Task"
    DESCRIPTION = "Base task"
    DIFFICULTY = TaskDifficulty.EASY
    AVAILABLE_SECRETS: list = []
    SCENARIOS: list = []

    def load_scenario(self, scenario_id: Optional[str] = None) -> Dict:
        if not self.SCENARIOS:
            raise ValueError(f"Task {self.__class__.__name__} has no scenarios defined")
        if scenario_id:
            for scenario in self.SCENARIOS:
                if scenario["id"] == scenario_id:
                    return scenario
            raise ValueError(f"Unknown scenario: {scenario_id}")
        return random.choice(self.SCENARIOS)