kaggle-environments 1.20.1__py3-none-any.whl → 1.21.0__py3-none-any.whl

kaggle_environments/envs/werewolf/game/records.py

@@ -0,0 +1,334 @@
+import json
+from abc import ABC
+from datetime import datetime
+from enum import IntEnum
+from typing import Dict, List, Optional
+from zoneinfo import ZoneInfo
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_serializer
+
+from .base import BaseAction, BaseEvent, PlayerID
+from .consts import DetailedPhase, EventName, ObsKeys, PerceivedThreatLevel, Phase, PhaseDivider, RoleConst, Team
+
+
+def get_utc_now():
+    return str(datetime.now(ZoneInfo("UTC")))
+
+
+class DataAccessLevel(IntEnum):
+    PUBLIC = 0
+    PERSONAL = 1
+
+
+class DataEntry(BaseModel, ABC):
+    """Abstract base class for all data entry types."""
+
+    pass
+
+
+class ActionDataMixin(BaseModel):
+    """
+    A mixin for action-related DataEntry models.
+    Includes the actor performing the action and their private reasoning.
+    """
+
+    actor_id: PlayerID
+    reasoning: Optional[str] = Field(
+        default=None, description="Private reasoning for moderator analysis.", access=DataAccessLevel.PERSONAL
+    )
+    perceived_threat_level: Optional[PerceivedThreatLevel] = Field(
+        default=PerceivedThreatLevel.SAFE, access=DataAccessLevel.PERSONAL
+    )
+    action: Optional[BaseAction] = Field(default=None, access=DataAccessLevel.PERSONAL)
+
+
+class VisibleRawData(BaseModel):
+    data_type: str
+    json_str: str
+
+
+class PlayerEventView(BaseModel):
+    day: int
+    phase: Phase
+    detailed_phase: DetailedPhase
+    event_name: EventName
+    description: str
+    data: Optional[dict | DataEntry] = None
+    source: str
+    created_at: str
+
+    @model_serializer
+    def serialize(self) -> dict:
+        if isinstance(self.data, DataEntry):
+            data = self.data.model_dump()
+        else:
+            data = self.data
+        return dict(
+            day=self.day,
+            phase=self.phase,
+            detailed_phase=self.detailed_phase,
+            event_name=self.event_name,
+            description=self.description,
+            data=data,
+            source=self.source,
+            created_at=self.created_at,
+        )
+
+
+class Event(BaseEvent):
+    day: int  # Day number, 0 for initial night
+    phase: Phase
+    detailed_phase: DetailedPhase
+    event_name: EventName
+    description: str
+    public: bool = False
+    visible_to: List[str] = Field(default_factory=list)
+    data: Optional[dict | DataEntry] = None
+    source: str
+    created_at: str = Field(default_factory=get_utc_now)
+    visible_in_ui: bool = True
+    """Determine if visible to game viewer in UI. Has no effect to game engine flow."""
+
+    @field_serializer("data")
+    def serialize_data(self, data):
+        if data is None:
+            return None
+        if isinstance(data, dict):
+            return data
+        if isinstance(data, BaseModel):
+            return data.model_dump()
+        return None
+
+    def serialize(self):
+        # TODO: this is purely constructed for compatibility with html renderer. Need to refactor werewolf.js to handle
+        #    a direct model_dump of Event
+        data_dict = self.model_dump()
+        return VisibleRawData(data_type=self.data.__class__.__name__, json_str=json.dumps(data_dict)).model_dump()
+
+    def view_by_access(self, user_level: DataAccessLevel) -> PlayerEventView:
+        if isinstance(self.data, ActionDataMixin):
+            fields_to_include = set()
+            fields_to_exclude = set()
+            for name, info in self.data.__class__.model_fields.items():
+                if info.json_schema_extra:
+                    if user_level >= info.json_schema_extra.get("access", DataAccessLevel.PUBLIC):
+                        fields_to_include.add(name)
+                    else:
+                        fields_to_exclude.add(name)
+                else:
+                    fields_to_include.add(name)
+            data = self.data.model_dump(include=fields_to_include, exclude=fields_to_exclude)
+        else:
+            data = self.data
+        out = PlayerEventView(
+            day=self.day,
+            phase=self.phase,
+            detailed_phase=self.detailed_phase,
+            event_name=self.event_name,
+            description=self.description,
+            data=data,
+            source=self.source,
+            created_at=self.created_at,
+        )
+        return out
+
+
+# --- Game State and Setup Data Entries ---
+class GameStartDataEntry(DataEntry):
+    player_ids: List[PlayerID]
+    number_of_players: int
+    role_counts: Dict[RoleConst, int]
+    team_member_counts: Dict[Team, int]
+    day_discussion_protocol_name: str
+    day_discussion_display_name: str
+    day_discussion_protocol_rule: str
+    night_werewolf_discussion_protocol_name: str
+    night_werewolf_discussion_display_name: str
+    night_werewolf_discussion_protocol_rule: str
+    day_voting_protocol_name: str
+    day_voting_display_name: str
+    day_voting_protocol_rule: str
+
+
+class GameStartRoleDataEntry(DataEntry):
+    player_id: PlayerID
+    team: Team
+    role: RoleConst
+    rule_of_role: str
+
+
+class SetNewPhaseDataEntry(DataEntry):
+    new_detailed_phase: DetailedPhase
+
+
+class PhaseDividerDataEntry(DataEntry):
+    divider_type: PhaseDivider
+
+
+# --- Request for Action Data Entries ---
+class RequestForActionDataEntry(DataEntry):
+    action_json_schema: str
+
+
+class RequestDoctorSaveDataEntry(RequestForActionDataEntry):
+    valid_candidates: List[PlayerID]
+
+
+class RequestSeerRevealDataEntry(RequestForActionDataEntry):
+    valid_candidates: List[PlayerID]
+
+
+class RequestWerewolfVotingDataEntry(RequestForActionDataEntry):
+    valid_targets: List[PlayerID]
+    alive_werewolve_player_ids: List[PlayerID]
+    voting_protocol_name: str
+    voting_protocol_rule: str
+
+
+class RequestVillagerToSpeakDataEntry(RequestForActionDataEntry):
+    pass
+
+
+# --- Action and Result Data Entries ---
+class SeerInspectResultDataEntry(DataEntry):
+    actor_id: PlayerID
+    target_id: PlayerID
+    role: Optional[RoleConst]
+    team: Optional[Team]
+
+
+class TargetedActionDataEntry(ActionDataMixin, DataEntry):
+    target_id: PlayerID
+
+
+class SeerInspectActionDataEntry(TargetedActionDataEntry):
+    """This records the Seer's choice of target to inspect."""
+
+
+class DoctorHealActionDataEntry(TargetedActionDataEntry):
+    """This records the Doctor's choice of target to heal."""
+
+
+class WerewolfNightVoteDataEntry(TargetedActionDataEntry):
+    """Records a werewolf's vote, including private reasoning."""
+
+
+class DayExileVoteDataEntry(TargetedActionDataEntry):
+    """Records a player's vote to exile, including private reasoning."""
+
+
+class DoctorSaveDataEntry(DataEntry):
+    """This records that a player was successfully saved by a doctor."""
+
+    saved_player_id: PlayerID
+
+
+class VoteOrderDataEntry(DataEntry):
+    vote_order_of_player_ids: List[PlayerID]
+
+
+class WerewolfNightEliminationElectedDataEntry(DataEntry):
+    """This record the elected elimination target by werewolves."""
+
+    elected_target_player_id: PlayerID
+
+
+class WerewolfNightEliminationDataEntry(DataEntry):
+    """This record the one eventually got eliminated by werewolves without doctor safe."""
+
+    eliminated_player_id: PlayerID
+    eliminated_player_role_name: Optional[RoleConst] = None
+    eliminated_player_team_name: Optional[Team] = None
+
+
+class DayExileElectedDataEntry(DataEntry):
+    elected_player_id: PlayerID
+    elected_player_role_name: Optional[RoleConst] = None
+    elected_player_team_name: Optional[Team] = None
+
+
+class DiscussionOrderDataEntry(DataEntry):
+    chat_order_of_player_ids: List[PlayerID]
+
+
+class ChatDataEntry(ActionDataMixin, DataEntry):
+    """Records a chat message from a player, including private reasoning."""
+
+    # actor_id and reasoning are inherited from ActionDataMixin
+    message: str
+    mentioned_player_ids: List[PlayerID] = Field(default_factory=list)
+
+
+class BidDataEntry(ActionDataMixin, DataEntry):
+    bid_amount: int
+
+
+class BidResultDataEntry(DataEntry):
+    winner_player_ids: List[PlayerID]
+    bid_overview: Dict[PlayerID, int]
+    mentioned_players_in_previous_turn: List[PlayerID] = []
+
+
+# --- Game End and Observation Models (Unchanged) ---
+class GameEndResultsDataEntry(DataEntry):
+    model_config = ConfigDict(use_enum_values=True)
+
+    winner_team: Team
+    winner_ids: List[PlayerID]
+    loser_ids: List[PlayerID]
+    scores: Dict[str, int | float]
+    reason: str
+    last_day: int
+    last_phase: Phase
+    survivors_until_last_round_and_role: Dict[PlayerID, RoleConst]
+    all_players_and_role: Dict[PlayerID, RoleConst]
+    elimination_info: List[Dict]
+    """list each player's elimination status, see GameState.get_elimination_info"""
+
+    all_players: List[Dict]
+    """provide the info dump for each player"""
+
+
+class WerewolfObservationModel(BaseModel):
+    player_id: PlayerID
+    role: RoleConst
+    team: Team
+    is_alive: bool
+    day: int
+    detailed_phase: DetailedPhase
+    all_player_ids: List[PlayerID]
+    player_thumbnails: Dict[PlayerID, str] = {}
+    alive_players: List[PlayerID]
+    revealed_players: Dict[PlayerID, RoleConst | Team | None] = {}
+    new_visible_announcements: List[str]
+    new_player_event_views: List[PlayerEventView]
+    game_state_phase: Phase
+
+    def get_human_readable(self) -> str:
+        # This is a placeholder implementation. A real implementation would format this nicely.
+        return json.dumps(self.model_dump(), indent=2)
+
+
+def set_raw_observation(kaggle_player_state, raw_obs: WerewolfObservationModel):
+    """Persist raw observations for players in kaggle's player state
+
+    Args:
+        kaggle_player_state: Kaggle's interpreter state is a list of player state. This arg is one player state item.
+        raw_obs: the raw observation for a player extracted from game engine.
+
+    Note: using raw_obs.model_dump_json() will greatly increase rendering speed (due to kaggle environment's use
+        of deepcopy for serialization) at the expense of harder to parse JSON rendering, since we are getting a json
+        string instead of human-readable dump. We choose raw_obs.model_dump() for clarity.
+    """
+    kaggle_player_state.observation[ObsKeys.RAW_OBSERVATION] = raw_obs.model_dump()
+
+
+def get_raw_observation(kaggle_observation) -> WerewolfObservationModel:
+    """
+
+    Args:
+        kaggle_observation:
+
+    Returns: a dict of WerewolfObservationModel dump
+    """
+    return WerewolfObservationModel(**kaggle_observation[ObsKeys.RAW_OBSERVATION])

kaggle_environments/envs/werewolf/game/roles.py

@@ -0,0 +1,326 @@
+import json
+import logging
+from collections import Counter, defaultdict, deque
+from functools import partial
+from typing import Deque, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_validator
+
+from .actions import HealAction, InspectAction
+from .base import BaseModerator, BasePlayer, BaseRole, EventHandler, PlayerID, on_event
+from .consts import EventName, Phase, RevealLevel, RoleConst, Team
+from .records import (
+    Event,
+    PlayerEventView,
+    RequestDoctorSaveDataEntry,
+    RequestSeerRevealDataEntry,
+    SeerInspectResultDataEntry,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Role(BaseRole):
+    model_config = ConfigDict(use_enum_values=True)
+
+    name: RoleConst = Field(..., frozen=True)
+    team: Team
+    night_priority: int = 100  # lower number acts earlier
+    descriptions: str
+
+
+class Werewolf(Role):
+    name: RoleConst = RoleConst.WEREWOLF
+    team: Team = Team.WEREWOLVES
+    night_priority: int = 2
+    descriptions: str = "Each night, collaborates with fellow werewolves to vote on eliminating one player."
+
+
+class Villager(Role):
+    name: RoleConst = RoleConst.VILLAGER
+    team: Team = Team.VILLAGERS
+    descriptions: str = "No special abilities. Participates in the daily vote to eliminate a suspected werewolf."
+
+
+class DoctorDescription:
+    ALLOW_SELF_SAVE = "Each night, may protect one player from a werewolf attack. Doctor is allowed to save themselves during night time."
+    NO_SELF_SAVE = "Each night, may protect one player from a werewolf attack. Doctor is NOT allowed to save themselves during night time."
+    NO_CONSECUTIVE_SAVE = " Doctor is NOT allowed to save the same player on consecutive nights."
+
+
+class DoctorStateKey:
+    LAST_SAVED_DAY = "last_saved_day"
+    LAST_SAVED_PLAYER_ID = "last_saved_player_id"
+
+
+class Doctor(Role):
+    name: RoleConst = RoleConst.DOCTOR
+    team: Team = Team.VILLAGERS
+    allow_self_save: bool = False
+    allow_consecutive_saves: bool = True
+    descriptions: str = ""
+
+    @model_validator(mode="after")
+    def set_descriptions_default(self) -> "Doctor":
+        if self.descriptions == "":
+            if self.allow_self_save:
+                self.descriptions = DoctorDescription.ALLOW_SELF_SAVE
+            else:
+                self.descriptions = DoctorDescription.NO_SELF_SAVE
+            if not self.allow_consecutive_saves:
+                self.descriptions += DoctorDescription.NO_CONSECUTIVE_SAVE
+        return self
+
+    @on_event(EventName.NIGHT_START)
+    def on_night_starts(self, me: BasePlayer, moderator: BaseModerator, event: Event):
+        if me.alive:
+            current_day = moderator.state.day_count
+            last_saved_day = me.get_role_state(DoctorStateKey.LAST_SAVED_DAY, default=-1)
+            last_saved_player_id = me.get_role_state(DoctorStateKey.LAST_SAVED_PLAYER_ID)
+
+            # Reset consecutive save memory if a night was skipped
+            if not self.allow_consecutive_saves and last_saved_day != -1 and current_day > last_saved_day + 1:
+                me.set_role_state(DoctorStateKey.LAST_SAVED_PLAYER_ID, None)
+                last_saved_player_id = None
+
+            valid_candidates = [p.id for p in moderator.state.alive_players()]
+
+            if not self.allow_self_save:
+                valid_candidates = [p_id for p_id in valid_candidates if p_id != me.id]
+
+            prompt = "Wake up Doctor. Who would you like to save? "
+            if not self.allow_consecutive_saves and last_saved_player_id:
+                valid_candidates = [p_id for p_id in valid_candidates if p_id != last_saved_player_id]
+                prompt += f'You cannot save the same player on consecutive nights. Player "{last_saved_player_id}" is not a valid target this night. '
+
+            data_entry = RequestDoctorSaveDataEntry(
+                valid_candidates=valid_candidates, action_json_schema=json.dumps(HealAction.schema_for_player())
+            )
+            prompt += f"The options are {data_entry.valid_candidates}."
+
+            moderator.request_action(
+                action_cls=HealAction,
+                player_id=me.id,
+                prompt=prompt,
+                data=data_entry,
+                event_name=EventName.HEAL_REQUEST,
+            )
+
+    @on_event(EventName.HEAL_ACTION)
+    def on_heal_action(self, me: BasePlayer, moderator: BaseModerator, event: Event):
+        if not me.alive or event.data.actor_id != me.id:
+            return
+
+        action = event.data.action
+        if isinstance(action, HealAction):
+            if not self.allow_self_save and action.target_id == me.id:
+                moderator.state.push_event(
+                    description=f'Player "{me.id}", doctor is not allowed to self save. '
+                    f"Your target is {action.target_id}, which is your own id.",
+                    event_name=EventName.ERROR,
+                    public=False,
+                    visible_to=[me.id],
+                )
+                return
+
+            if not self.allow_consecutive_saves and action.target_id == me.get_role_state(
+                DoctorStateKey.LAST_SAVED_PLAYER_ID
+            ):
+                moderator.state.push_event(
+                    description=f'Player "{me.id}", you cannot save the same player on consecutive nights. '
+                    f'Your target "{action.target_id}" was also saved last night.',
+                    event_name=EventName.ERROR,
+                    public=False,
+                    visible_to=[me.id],
+                )
+                return
+
+            moderator.record_night_save(me.id, action.target_id)
+            me.set_role_state(DoctorStateKey.LAST_SAVED_PLAYER_ID, action.target_id)
+            me.set_role_state(DoctorStateKey.LAST_SAVED_DAY, moderator.state.day_count)
+
+
+class SeerDescription:
+    REVEAL_ROLE = "Each night, may inspect one player to learn their true role."
+    REVEAL_TEAM = "Each night, may inspect one player's team but not their role."
+
+
+class Seer(Role):
+    name: RoleConst = RoleConst.SEER
+    team: Team = Team.VILLAGERS
+    descriptions: str = ""
+    reveal_level: RevealLevel = RevealLevel.ROLE
+
+    @field_validator("reveal_level")
+    @classmethod
+    def validate_reveal_level(cls, v):
+        if v == RevealLevel.NO_REVEAL:
+            raise ValueError(f"Setting reveal_level of Seer as {v}. Seer will become useless.")
+        return v
+
+    @model_validator(mode="after")
+    def set_descriptions_default(self) -> "Seer":
+        if self.descriptions == "":
+            if self.reveal_level == RevealLevel.ROLE:
+                self.descriptions = SeerDescription.REVEAL_ROLE
+            elif self.reveal_level == RevealLevel.TEAM:
+                self.descriptions = SeerDescription.REVEAL_TEAM
+            else:
+                raise ValueError(f"reveal_level {self.reveal_level} not supported.")
+        return self
+
+    @on_event(EventName.NIGHT_START)
+    def on_night_starts(self, me: BasePlayer, moderator: BaseModerator, event: Event):
+        if me.alive:
+            data_entry = RequestSeerRevealDataEntry(
+                valid_candidates=[p.id for p in moderator.state.alive_players() if p != me],
+                action_json_schema=json.dumps(InspectAction.schema_for_player()),
+            )
+            moderator.request_action(
+                action_cls=InspectAction,
+                player_id=me.id,
+                prompt=f"Wake up Seer. Who would you like to see their true {self.reveal_level}? "
+                f"The options are {data_entry.valid_candidates}.",
+                data=data_entry,
+                event_name=EventName.INSPECT_REQUEST,
+            )
+
+    @on_event(EventName.INSPECT_ACTION)
+    def on_inspect_action(self, me: BasePlayer, moderator: BaseModerator, event: Event):
+        action = event.data.action
+        if not me.alive or action.actor_id != me.id:
+            return
+        actor_id = me.id
+        target_player = moderator.state.get_player_by_id(action.target_id)
+        if target_player:  # Ensure target exists
+            role = None
+            team = None
+            reveal_text = ""
+            if self.reveal_level == RevealLevel.ROLE:
+                role = target_player.role.name
+                team = target_player.role.team
+                reveal_text = f'Their role is a "{target_player.role.name}" in team "{target_player.role.team.value}".'
+            elif self.reveal_level == RevealLevel.TEAM:
+                team = target_player.role.team
+                reveal_text = f"Their team is {team}."
+
+            data = SeerInspectResultDataEntry(actor_id=actor_id, target_id=action.target_id, role=role, team=team)
+            moderator.state.push_event(
+                description=f'Player "{actor_id}", you inspected {target_player.id}. ' + reveal_text,
+                event_name=EventName.INSPECT_RESULT,
+                public=False,
+                visible_to=[actor_id],
+                data=data,
+            )
+        else:
+            moderator.state.push_event(
+                description=f'Player "{actor_id}", you inspected player "{action.target_id}",'
+                f" but this player could not be found.",
+                event_name=EventName.ERROR,
+                public=False,
+                visible_to=[actor_id],
+            )
+
+
+class LLM(BaseModel):
+    model_name: str
+    properties: Dict = {}
+
+
+class Agent(BaseModel):
+    id: PlayerID
+    """The unique name of the player."""
+
+    agent_id: str
+    """Id of the agent. Might not be unique (many players might be using the same underlying agent)."""
+
+    display_name: str = ""
+    """Agent name shown in the UI and only visible to spectator but not the players. e.g. Pete (base_harness-gemini-2.5-pro)
+    base_harness-gemini-2.5-pro is the display_name while Pete is the id. It maybe different from agent_id, 
+    e.g. base_harness_v2-gemini-2.5-pro-0506, to reduce the cognitive load of the spectators.
+    """
+
+    role: RoleConst
+    role_params: Dict = Field(default_factory=dict)
+    """Parameters to the Role constructor"""
+
+    thumbnail: Optional[str] = ""
+    agent_harness_name: str = "basic_llm"
+    llms: List[LLM] = []
+
+    def get_agent_name(self):
+        return f"{self.agent_harness_name}({', '.join([llm.model_name for llm in self.llms])})"
+
+
+class Player(BasePlayer):
+    model_config = ConfigDict(use_enum_values=True)
+
+    id: PlayerID
+    """The unique name of the player."""
+
+    agent: Agent
+    role: BaseRole
+    alive: bool = True
+    eliminated_during_day: int = -1
+    """game starts at night 0, then day 1, night 1, day 2, ..."""
+
+    eliminated_during_phase: Optional[Phase] = None
+
+    _message_queue: Deque[PlayerEventView] = PrivateAttr(default_factory=deque)
+    _role_state: Dict = PrivateAttr(default_factory=dict)
+
+    def set_role_state(self, key, value):
+        self._role_state[key] = value
+
+    def get_role_state(self, key, default=None):
+        return self._role_state.get(key, default)
+
+    def get_event_handlers(self, moderator: BaseModerator) -> Dict[EventName, List[EventHandler]]:
+        handlers = defaultdict(list)
+        for event_type, handler in self.role.get_event_handlers().items():
+            event_handler = partial(handler, self, moderator)
+            handlers[event_type].append(event_handler)
+        return handlers
+
+    def update(self, entry: PlayerEventView):
+        self._message_queue.append(entry)
+
+    def consume_messages(self) -> List[PlayerEventView]:
+        messages = list(self._message_queue)
+        self._message_queue.clear()
+        return messages
+
+    def eliminate(self, day: int, phase: Phase):
+        self.alive = False
+        self.eliminated_during_day = day
+        self.eliminated_during_phase = phase.value
+
+    def report_elimination(self):
+        return {
+            "player_id": self.id,
+            "eliminated_during_day": self.eliminated_during_day,
+            "eliminated_during_phase": self.eliminated_during_phase,
+        }
+
+
+ROLE_CLASS_MAP = {
+    RoleConst.WEREWOLF.value: Werewolf,
+    RoleConst.DOCTOR.value: Doctor,
+    RoleConst.SEER.value: Seer,
+    RoleConst.VILLAGER.value: Villager,
+}
+
+
+def create_players_from_agents_config(agents_config: List[Dict]) -> List[Player]:
+    # check all agents have unique ids
+    agent_ids = [agent_config["id"] for agent_config in agents_config]
+    if len(agent_ids) != len(set(agent_ids)):
+        counts = Counter(agent_ids)
+        duplicates = [item for item, count in counts.items() if count > 1 and item is not None]
+        if duplicates:
+            raise ValueError(f"Duplicate agent ids found: {', '.join(duplicates)}")
+    agents = [Agent(**agent_config) for agent_config in agents_config]
+    players = [
+        Player(id=agent.id, agent=agent, role=ROLE_CLASS_MAP[agent.role](**agent.role_params)) for agent in agents
+    ]
+    return players