agents/rule_based/equal_split.py

"""Equal-split baseline: even share of the discretionary pool."""

from collections.abc import Iterable

from agents.base import PolicyAdapter
from agents.rule_based.discretionary import discretionary_pool
from agents.rule_based.voting import first_vote_target
from schemas.actions import (
    Action,
    ActionType,
    DebateAction,
    ProposeBudgetAction,
    VoteAction,
    VoteChoice,
)
from schemas.departments import DEFAULT_DEPARTMENT_NAMES
from schemas.observations import Observation


class EqualSplitAdapter(PolicyAdapter):
    def __init__(self, department_count: int = len(DEFAULT_DEPARTMENT_NAMES)) -> None:
        if department_count <= 0:
            raise ValueError("department_count must be positive.")
        self.department_count = department_count

    def act(
        self,
        observation: Observation,
        valid_actions: Iterable[str],
        agent_id: str,
    ) -> Action:
        valid_action_set = set(valid_actions)
        if ActionType.PROPOSE_BUDGET.value in valid_action_set:
            pool = discretionary_pool(observation)
            share = pool / self.department_count
            return ProposeBudgetAction(
                type=ActionType.PROPOSE_BUDGET,
                department=observation.own_department.name,
                amount=share,
                justification="Request an equal share of the discretionary treasury pool.",
            )

        vote_target = first_vote_target(observation.proposals, agent_id)
        if ActionType.VOTE.value in valid_action_set and vote_target is not None:
            return VoteAction(
                type=ActionType.VOTE,
                proposal_id=vote_target.proposal_id,
                vote=VoteChoice.YES,
            )

        if ActionType.DEBATE.value in valid_action_set:
            return DebateAction(
                type=ActionType.DEBATE,
                message=f"{agent_id} supports equal treasury distribution.",
            )

        return DebateAction(
            type=ActionType.DEBATE,
            message=f"{agent_id} (equal split) idle.",
        )