agents/random_agent.py

"""
Simple random comparison agent.
"""

import random
from typing import Optional, Protocol

from env.models import Action, ActionPayload, Observation, StepResult

ACTION_TYPES = ("execute", "delay", "reallocate")


class SupportsEpisodeEnv(Protocol):
    def reset(self) -> Observation:
        ...

    def step(self, action: Action) -> StepResult:
        ...


def select_random_action(obs: Observation, rng: Optional[random.Random] = None) -> Action:
    rng = rng or random.Random()
    task_ids = [task.id for task in obs.tasks]
    if not task_ids:
        return Action(action_type="delay", task_id=0)

    return Action(
        action_type=rng.choice(ACTION_TYPES),
        task_id=rng.choice(task_ids),
    )


def run_episode(
    env: SupportsEpisodeEnv,
    max_steps: int = 30,
    rng: Optional[random.Random] = None,
) -> list[ActionPayload]:
    rng = rng or random.Random(42)
    obs = env.reset()
    actions_taken: list[ActionPayload] = []

    for _ in range(max_steps):
        if obs.episode_done:
            break
        action = select_random_action(obs, rng=rng)
        actions_taken.append(action.model_dump())
        obs = env.step(action).observation

    return actions_taken