"""Simulation engine for Analog Town state transitions.""" import json import traceback from datetime import datetime from schemas import ( AgentProfile, AgentState, BroadcastEvent, SimulationResult, StateTransition, Town, ) from prompts import SYSTEM_PROMPT, STATE_TRANSITION_PROMPT from model_client import ModelClient class Simulator: """Runs state-transition simulations for all agents in a town.""" def __init__(self, model_client: ModelClient | None = None): self.model_client = model_client or ModelClient() def _get_initial_state(self, agent: AgentProfile) -> AgentState: """Create initial agent state from profile defaults.""" return AgentState( agent_id=agent.id, mood_label="baseline", trust=0.5, anger=0.2, fear=0.3, hope=0.5, curiosity=0.5, social_energy=0.5, current_belief="No strong opinion yet.", active_memory=None, unresolved_tension=None, ) def _clamp(self, value: float, min_val: float = 0.0, max_val: float = 1.0) -> float: """Clamp a value between min and max.""" return max(min_val, min(max_val, value)) def _apply_emotion_deltas( self, old_state: AgentState, deltas: dict[str, float] ) -> dict[str, float]: """Apply emotion deltas to old state, clamping to [0.0, 1.0].""" emotions = {} for key in ["trust", "anger", "fear", "hope", "curiosity", "social_energy"]: old_val = getattr(old_state, key, 0.5) delta = deltas.get(key, 0.0) emotions[key] = self._clamp(old_val + delta) return emotions def _validate_transition( self, transition_data: dict, agent: AgentProfile ) -> tuple[bool, str]: """Validate a raw transition dict before parsing into Pydantic model.""" errors = [] # Check agent_id matches if transition_data.get("agent_id") != agent.id: transition_data["agent_id"] = agent.id # Check required fields required = [ "event_summary_from_agent_view", "noticed_detail", "activated_memory", "value_conflict", "emotion_delta", "updated_state", "internal_monologue", "likely_private_action", "likely_public_action", "uncertainty", ] for field in required: if field not in transition_data: errors.append(f"Missing required field: {field}") # Check emotion values in updated_state updated = transition_data.get("updated_state", {}) for emotion in ["trust", "anger", "fear", "hope", "curiosity", "social_energy"]: val = updated.get(emotion) if val is not None: if not isinstance(val, (int, float)): errors.append(f"Emotion '{emotion}' must be a number, got {type(val)}") elif val < 0.0 or val > 1.0: # Clamp instead of error updated[emotion] = self._clamp(float(val)) # Ensure agent_id in updated_state if "updated_state" in transition_data: transition_data["updated_state"]["agent_id"] = agent.id # Check internal monologue is not empty monologue = transition_data.get("internal_monologue", "") if not monologue or not monologue.strip(): errors.append("Internal monologue is empty") # Ensure safety note if "safety_note" not in transition_data or not transition_data["safety_note"]: transition_data["safety_note"] = ( "This is a fictional perspective rehearsal, not a prediction." ) if errors: return False, "; ".join(errors) return True, "" def run_agent_transition( self, agent: AgentProfile, state: AgentState, event: BroadcastEvent, day: int = 1, ) -> StateTransition: """Run a single agent through the state transition. Args: agent: The agent's profile state: The agent's current state event: The broadcast event day: 1-indexed broadcast day, used to instruct the model to vary follow-up beats Returns: StateTransition with updated state and monologue Raises: RuntimeError: If both generation and repair fail """ user_prompt = STATE_TRANSITION_PROMPT.format( agent_profile=agent.model_dump_json(indent=2), previous_state=state.model_dump_json(indent=2), broadcast_event=event.model_dump_json(indent=2), ) if day > 1: user_prompt = ( f"DAY {day} OF THIS SCENARIO. " f"The agent has already lived through {day - 1} previous broadcast(s). " "The fields current_belief, active_memory, and unresolved_tension in 'Previous State' " "capture where the agent ENDED LAST TIME — they are CONTEXT, not your script. " "Write a FRESH internal_monologue that is clearly different in wording from any prior beat: " "the agent has had time to process, talk to others, sleep on it, or harden their view. " "Make at least one emotion_delta non-zero. Do NOT echo any sentence verbatim from current_belief.\n\n" + user_prompt ) original_temp = getattr(self.model_client, "temperature", 0.3) try: if day > 1: self.model_client.temperature = min(0.85, original_temp + 0.25) transition_data = self.model_client.generate_json( system_prompt=SYSTEM_PROMPT, user_prompt=user_prompt, ) finally: self.model_client.temperature = original_temp # Validate and fix is_valid, error_msg = self._validate_transition(transition_data, agent) if not is_valid: raise RuntimeError(f"Transition validation failed: {error_msg}") # Parse into Pydantic model transition = StateTransition(**transition_data) return transition def run_town_simulation( self, town: Town, event: BroadcastEvent, previous_states: dict[str, AgentState] | None = None, progress_callback=None, day: int = 1, ) -> SimulationResult: """Run simulation for all agents in the town. Args: town: The town with agents event: The broadcast event previous_states: Optional dict of agent_id -> AgentState seeded from prior run progress_callback: Optional callback(agent_name, status, index, total) Returns: SimulationResult with all transitions (failed agents are skipped) """ transitions = [] total = len(town.agents) for i, agent in enumerate(town.agents): agent_name = agent.name try: if progress_callback: progress_callback(agent_name, "processing", i, total) state = ( previous_states.get(agent.id, self._get_initial_state(agent)) if previous_states else self._get_initial_state(agent) ) transition = self.run_agent_transition(agent, state, event, day=day) transitions.append(transition) if progress_callback: progress_callback(agent_name, "complete", i, total) except Exception as e: # Don't crash on individual agent failures print(f"⚠ Agent '{agent_name}' failed: {e}") traceback.print_exc() if progress_callback: progress_callback(agent_name, f"failed: {str(e)[:100]}", i, total) continue return SimulationResult( town_id=town.id, event=event, transitions=transitions, created_at=datetime.now().isoformat(), )