Spaces:
Sleeping
Sleeping
| """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(), | |
| ) | |