""" Atomizer Session State - Context Isolation Management Part of the ACE (Agentic Context Engineering) implementation for Atomizer. Implements the "Write-Select-Compress-Isolate" pattern: - Exposed fields are sent to LLM at every turn - Isolated fields are accessed selectively when needed - Automatic compression of old data This ensures efficient context usage while maintaining access to full historical data when needed. """ from typing import Dict, List, Optional, Any from datetime import datetime from enum import Enum from dataclasses import dataclass, field import json from pathlib import Path class TaskType(Enum): """Types of tasks Claude can perform in Atomizer.""" CREATE_STUDY = "create_study" RUN_OPTIMIZATION = "run_optimization" MONITOR_PROGRESS = "monitor_progress" ANALYZE_RESULTS = "analyze_results" DEBUG_ERROR = "debug_error" CONFIGURE_SETTINGS = "configure_settings" EXPORT_DATA = "export_data" NEURAL_ACCELERATION = "neural_acceleration" @dataclass class ExposedState: """ State exposed to LLM at every turn. Keep this minimal - only what's needed for immediate context. Everything here counts against token budget every turn. """ # Current task context task_type: Optional[TaskType] = None current_objective: str = "" # Recent history (compressed) recent_actions: List[str] = field(default_factory=list) recent_errors: List[str] = field(default_factory=list) # Active study summary study_name: Optional[str] = None study_status: str = "unknown" trials_completed: int = 0 trials_total: int = 0 best_value: Optional[float] = None best_trial: Optional[int] = None # Playbook excerpt (most relevant items) active_playbook_items: List[str] = field(default_factory=list) # Constraints for context size MAX_ACTIONS: int = 10 MAX_ERRORS: int = 5 MAX_PLAYBOOK_ITEMS: int = 15 @dataclass class IsolatedState: """ State isolated from LLM - accessed selectively. This data is NOT included in every context window. Load specific fields when explicitly needed. """ # Full optimization history (can be large) full_trial_history: List[Dict[str, Any]] = field(default_factory=list) # NX session state (heavy, complex) nx_model_path: Optional[str] = None nx_expressions: Dict[str, Any] = field(default_factory=dict) nx_sim_path: Optional[str] = None # Neural network cache neural_predictions: Dict[str, float] = field(default_factory=dict) surrogate_model_path: Optional[str] = None # Full playbook (loaded on demand) full_playbook_path: Optional[str] = None # Debug information last_solver_output: str = "" last_f06_content: str = "" last_solver_returncode: Optional[int] = None # Configuration snapshots optimization_config: Dict[str, Any] = field(default_factory=dict) study_config: Dict[str, Any] = field(default_factory=dict) @dataclass class AtomizerSessionState: """ Complete session state with exposure control. The exposed state is automatically injected into every LLM context. The isolated state is accessed only when explicitly needed. Usage: session = AtomizerSessionState(session_id="session_001") session.exposed.task_type = TaskType.CREATE_STUDY session.add_action("Created study directory") # Get context for LLM context = session.get_llm_context() # Access isolated data when needed f06 = session.load_isolated_data("last_f06_content") """ session_id: str created_at: str = field(default_factory=lambda: datetime.now().isoformat()) last_updated: str = field(default_factory=lambda: datetime.now().isoformat()) exposed: ExposedState = field(default_factory=ExposedState) isolated: IsolatedState = field(default_factory=IsolatedState) def get_llm_context(self) -> str: """ Generate context string for LLM consumption. Only includes exposed state - isolated state requires explicit access via load_isolated_data(). Returns: Formatted markdown context string """ lines = [ "## Current Session State", "", f"**Task**: {self.exposed.task_type.value if self.exposed.task_type else 'Not set'}", f"**Objective**: {self.exposed.current_objective or 'None specified'}", "", ] # Study context if self.exposed.study_name: progress = "" if self.exposed.trials_total > 0: pct = (self.exposed.trials_completed / self.exposed.trials_total) * 100 progress = f" ({pct:.0f}%)" lines.extend([ f"### Active Study: {self.exposed.study_name}", f"- Status: {self.exposed.study_status}", f"- Trials: {self.exposed.trials_completed}/{self.exposed.trials_total}{progress}", ]) if self.exposed.best_value is not None: lines.append(f"- Best: {self.exposed.best_value:.6g} (trial #{self.exposed.best_trial})") lines.append("") # Recent actions if self.exposed.recent_actions: lines.append("### Recent Actions") for action in self.exposed.recent_actions[-5:]: lines.append(f"- {action}") lines.append("") # Recent errors (highlight these) if self.exposed.recent_errors: lines.append("### Recent Errors (address these)") for error in self.exposed.recent_errors: lines.append(f"- {error}") lines.append("") # Relevant playbook items if self.exposed.active_playbook_items: lines.append("### Relevant Knowledge") for item in self.exposed.active_playbook_items: lines.append(f"- {item}") lines.append("") return "\n".join(lines) def add_action(self, action: str) -> None: """ Record an action (auto-compresses old actions). Args: action: Description of the action taken """ timestamp = datetime.now().strftime("%H:%M:%S") self.exposed.recent_actions.append(f"[{timestamp}] {action}") # Compress if over limit if len(self.exposed.recent_actions) > self.exposed.MAX_ACTIONS: # Keep first, summarize middle, keep last 5 first = self.exposed.recent_actions[0] last_five = self.exposed.recent_actions[-5:] middle_count = len(self.exposed.recent_actions) - 6 self.exposed.recent_actions = ( [first] + [f"... ({middle_count} earlier actions)"] + last_five ) self.last_updated = datetime.now().isoformat() def add_error(self, error: str, error_type: str = "") -> None: """ Record an error for LLM attention. Errors are preserved more aggressively than actions because they need to be addressed. Args: error: Error message error_type: Optional error classification """ prefix = f"[{error_type}] " if error_type else "" self.exposed.recent_errors.append(f"{prefix}{error}") # Keep most recent errors self.exposed.recent_errors = self.exposed.recent_errors[-self.exposed.MAX_ERRORS:] self.last_updated = datetime.now().isoformat() def clear_errors(self) -> None: """Clear all recorded errors (after they're addressed).""" self.exposed.recent_errors = [] self.last_updated = datetime.now().isoformat() def update_study_status( self, name: str, status: str, trials_completed: int, trials_total: int, best_value: Optional[float] = None, best_trial: Optional[int] = None ) -> None: """ Update the study status in exposed state. Args: name: Study name status: Current status (running, completed, failed, etc.) trials_completed: Number of completed trials trials_total: Total planned trials best_value: Best objective value found best_trial: Trial number with best value """ self.exposed.study_name = name self.exposed.study_status = status self.exposed.trials_completed = trials_completed self.exposed.trials_total = trials_total self.exposed.best_value = best_value self.exposed.best_trial = best_trial self.last_updated = datetime.now().isoformat() def set_playbook_items(self, items: List[str]) -> None: """ Set the active playbook items for context. Args: items: List of playbook item context strings """ self.exposed.active_playbook_items = items[:self.exposed.MAX_PLAYBOOK_ITEMS] self.last_updated = datetime.now().isoformat() def load_isolated_data(self, key: str) -> Any: """ Explicitly load isolated data when needed. Use this when you need access to heavy data that shouldn't be in every context window. Args: key: Attribute name in IsolatedState Returns: The isolated data value, or None if not found """ return getattr(self.isolated, key, None) def set_isolated_data(self, key: str, value: Any) -> None: """ Set isolated data. Args: key: Attribute name in IsolatedState value: Value to set """ if hasattr(self.isolated, key): setattr(self.isolated, key, value) self.last_updated = datetime.now().isoformat() def add_trial_to_history(self, trial_data: Dict[str, Any]) -> None: """ Add a trial to the full history (isolated state). Args: trial_data: Dictionary with trial information """ trial_data["recorded_at"] = datetime.now().isoformat() self.isolated.full_trial_history.append(trial_data) self.last_updated = datetime.now().isoformat() def get_trial_history_summary(self, last_n: int = 10) -> List[Dict[str, Any]]: """ Get summary of recent trials from isolated history. Args: last_n: Number of recent trials to return Returns: List of trial summary dictionaries """ return self.isolated.full_trial_history[-last_n:] def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for serialization.""" return { "session_id": self.session_id, "created_at": self.created_at, "last_updated": self.last_updated, "exposed": { "task_type": self.exposed.task_type.value if self.exposed.task_type else None, "current_objective": self.exposed.current_objective, "recent_actions": self.exposed.recent_actions, "recent_errors": self.exposed.recent_errors, "study_name": self.exposed.study_name, "study_status": self.exposed.study_status, "trials_completed": self.exposed.trials_completed, "trials_total": self.exposed.trials_total, "best_value": self.exposed.best_value, "best_trial": self.exposed.best_trial, "active_playbook_items": self.exposed.active_playbook_items }, "isolated": { "nx_model_path": self.isolated.nx_model_path, "nx_sim_path": self.isolated.nx_sim_path, "surrogate_model_path": self.isolated.surrogate_model_path, "full_playbook_path": self.isolated.full_playbook_path, "trial_history_count": len(self.isolated.full_trial_history) } } def save(self, path: Path) -> None: """ Save session state to JSON. Note: Full trial history is saved to a separate file to keep the main state file manageable. Args: path: Path to save state file """ path.parent.mkdir(parents=True, exist_ok=True) # Save main state with open(path, 'w', encoding='utf-8') as f: json.dump(self.to_dict(), f, indent=2) # Save trial history separately if large if len(self.isolated.full_trial_history) > 0: history_path = path.with_suffix('.history.json') with open(history_path, 'w', encoding='utf-8') as f: json.dump(self.isolated.full_trial_history, f, indent=2) @classmethod def load(cls, path: Path) -> "AtomizerSessionState": """ Load session state from JSON. Args: path: Path to state file Returns: Loaded session state (or new state if file doesn't exist) """ if not path.exists(): return cls(session_id=f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}") with open(path, encoding='utf-8') as f: data = json.load(f) state = cls( session_id=data.get("session_id", "unknown"), created_at=data.get("created_at", datetime.now().isoformat()), last_updated=data.get("last_updated", datetime.now().isoformat()) ) # Load exposed state exposed = data.get("exposed", {}) if exposed.get("task_type"): state.exposed.task_type = TaskType(exposed["task_type"]) state.exposed.current_objective = exposed.get("current_objective", "") state.exposed.recent_actions = exposed.get("recent_actions", []) state.exposed.recent_errors = exposed.get("recent_errors", []) state.exposed.study_name = exposed.get("study_name") state.exposed.study_status = exposed.get("study_status", "unknown") state.exposed.trials_completed = exposed.get("trials_completed", 0) state.exposed.trials_total = exposed.get("trials_total", 0) state.exposed.best_value = exposed.get("best_value") state.exposed.best_trial = exposed.get("best_trial") state.exposed.active_playbook_items = exposed.get("active_playbook_items", []) # Load isolated state metadata isolated = data.get("isolated", {}) state.isolated.nx_model_path = isolated.get("nx_model_path") state.isolated.nx_sim_path = isolated.get("nx_sim_path") state.isolated.surrogate_model_path = isolated.get("surrogate_model_path") state.isolated.full_playbook_path = isolated.get("full_playbook_path") # Load trial history from separate file if exists history_path = path.with_suffix('.history.json') if history_path.exists(): with open(history_path, encoding='utf-8') as f: state.isolated.full_trial_history = json.load(f) return state # Convenience functions for session management _active_session: Optional[AtomizerSessionState] = None def get_session() -> AtomizerSessionState: """ Get the active session state. Creates a new session if none exists. Returns: The active AtomizerSessionState """ global _active_session if _active_session is None: _active_session = AtomizerSessionState( session_id=f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}" ) return _active_session def set_session(session: AtomizerSessionState) -> None: """ Set the active session. Args: session: Session state to make active """ global _active_session _active_session = session def clear_session() -> None: """Clear the active session.""" global _active_session _active_session = None