"""Export simulation traces to local files and Hugging Face Hub.""" import json import os import re import shutil import tempfile from datetime import datetime from pathlib import Path from schemas import SimulationResult, Town, BroadcastEvent # PII patterns to sanitize from exports PII_PATTERNS = [ (r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '[EMAIL_REMOVED]'), (r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b', '[PHONE_REMOVED]'), (r'\b\d{1,5}\s+[\w\s]+(?:Street|St|Avenue|Ave|Road|Rd|Boulevard|Blvd|Drive|Dr|Lane|Ln|Way|Court|Ct)\b', '[ADDRESS_REMOVED]'), ] DATASET_CARD_TEMPLATE = """# Analog Town Trace Dataset ## Description This dataset contains fictional agent state-transition traces generated by the Analog Town simulator. **Town:** {town_name} **Event:** {event_title} **Agents:** {agent_count} **Generated:** {created_at} ## Intended Use This dataset is intended for demonstrating small-model perspective rehearsal, fictional character simulation, and inspectable state-evolution prompting. ## Limitations The traces are fictional. They should not be interpreted as predictions of real people. ## Ethical Considerations Users should avoid uploading real private personal information. Real individuals should be fictionalized or represented only with consent. ## Data Fields - `town_id` — identifier for the simulated town - `event_id` — identifier for the broadcast event - `agent_id` — identifier for the responding agent - `frequency` — the agent's radio frequency - `previous_state` — agent state before the event - `transition` — the state transition details - `updated_state` — agent state after the event - `created_at` — timestamp of generation ## Source Generated by [Analog Town](https://github.com/analog-town) — a shortwave simulator for fictional minds. """ class ExportManager: """Manages export of simulation traces to local files and HF Hub.""" def _sanitize_text(self, text: str) -> str: """Remove PII patterns from text.""" for pattern, replacement in PII_PATTERNS: text = re.sub(pattern, replacement, text, flags=re.IGNORECASE) return text def _sanitize_dict(self, data: dict) -> dict: """Recursively sanitize PII from a dictionary.""" sanitized = {} for key, value in data.items(): if isinstance(value, str): sanitized[key] = self._sanitize_text(value) elif isinstance(value, dict): sanitized[key] = self._sanitize_dict(value) elif isinstance(value, list): sanitized[key] = [ self._sanitize_dict(item) if isinstance(item, dict) else self._sanitize_text(item) if isinstance(item, str) else item for item in value ] else: sanitized[key] = value return sanitized def _make_event_id(self, event: BroadcastEvent) -> str: """Generate a simple event ID from the title.""" return re.sub(r'[^a-z0-9]+', '_', event.title.lower()).strip('_')[:50] def export_local( self, result: SimulationResult, town: Town, output_dir: str = "outputs", ) -> str: """Export simulation traces to local directory. Creates: - town.json - broadcast_event.json - agent_traces.jsonl - README.md Returns the output directory path. """ # Create timestamped output directory timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") export_dir = Path(output_dir) / f"export_{result.town_id}_{timestamp}" export_dir.mkdir(parents=True, exist_ok=True) event_id = self._make_event_id(result.event) # 1. town.json town_data = self._sanitize_dict(town.model_dump()) with open(export_dir / "town.json", "w") as f: json.dump(town_data, f, indent=2, ensure_ascii=False) # 2. broadcast_event.json event_data = self._sanitize_dict(result.event.model_dump()) with open(export_dir / "broadcast_event.json", "w") as f: json.dump(event_data, f, indent=2, ensure_ascii=False) # 3. agent_traces.jsonl # Find agent frequencies from town freq_map = {agent.id: agent.frequency for agent in town.agents} with open(export_dir / "agent_traces.jsonl", "w") as f: for transition in result.transitions: trace_line = { "town_id": result.town_id, "event_id": event_id, "agent_id": transition.agent_id, "frequency": freq_map.get(transition.agent_id, 0.0), "previous_state": {}, # We don't store previous state separately "transition": self._sanitize_dict( transition.model_dump( exclude={"updated_state"} ) ), "updated_state": self._sanitize_dict( transition.updated_state.model_dump() ), "created_at": result.created_at, } f.write(json.dumps(trace_line, ensure_ascii=False) + "\n") # 4. README.md (dataset card) readme_content = DATASET_CARD_TEMPLATE.format( town_name=town.name, event_title=result.event.title, agent_count=len(result.transitions), created_at=result.created_at, ) with open(export_dir / "README.md", "w") as f: f.write(readme_content) return str(export_dir) def export_zip( self, result: SimulationResult, town: Town, output_dir: str = "outputs", ) -> str: """Export simulation traces as a ZIP file. Returns the path to the ZIP file. """ # First export to temp directory with tempfile.TemporaryDirectory() as tmp_dir: export_dir = self.export_local(result, town, tmp_dir) # Create ZIP zip_dir = Path(output_dir) zip_dir.mkdir(parents=True, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") zip_name = f"analog_town_{result.town_id}_{timestamp}" zip_path = shutil.make_archive( str(zip_dir / zip_name), "zip", export_dir ) return zip_path def export_to_hub( self, result: SimulationResult, town: Town, repo_id: str, token: str | None = None, ) -> str: """Upload simulation traces to Hugging Face Hub as a dataset. Args: result: The simulation result town: The town configuration repo_id: HF repo ID (e.g., "username/analog-town-graybridge") token: HF token (uses env HF_TOKEN if not provided) Returns: The URL of the created dataset """ from huggingface_hub import HfApi token = token or os.getenv("HF_TOKEN") if not token: raise ValueError("HF_TOKEN required for Hub upload") api = HfApi(token=token) # Export to temp directory first with tempfile.TemporaryDirectory() as tmp_dir: export_dir = self.export_local(result, town, tmp_dir) # Create or get repo try: api.create_repo( repo_id=repo_id, repo_type="dataset", exist_ok=True, ) except Exception as e: print(f"Note: repo creation issue (may already exist): {e}") # Upload all files api.upload_folder( folder_path=export_dir, repo_id=repo_id, repo_type="dataset", ) return f"https://huggingface.co/datasets/{repo_id}" def generate_dataset_card(self, town: Town, event: BroadcastEvent) -> str: """Generate a dataset card markdown string.""" return DATASET_CARD_TEMPLATE.format( town_name=town.name, event_title=event.title, agent_count=len(town.agents), created_at=datetime.now().isoformat(), )