feat: Implement ACE Context Engineering framework (SYS_17)
Complete implementation of Agentic Context Engineering (ACE) framework: Core modules (optimization_engine/context/): - playbook.py: AtomizerPlaybook with helpful/harmful scoring - reflector.py: AtomizerReflector for insight extraction - session_state.py: Context isolation (exposed/isolated state) - feedback_loop.py: Automated learning from trial results - compaction.py: Long-session context management - cache_monitor.py: KV-cache optimization tracking - runner_integration.py: OptimizationRunner integration Dashboard integration: - context.py: 12 REST API endpoints for playbook management Tests: - test_context_engineering.py: 44 unit tests - test_context_integration.py: 16 integration tests Documentation: - CONTEXT_ENGINEERING_REPORT.md: Comprehensive implementation report - CONTEXT_ENGINEERING_API.md: Complete API reference - SYS_17_CONTEXT_ENGINEERING.md: System protocol - Updated cheatsheet with SYS_17 quick reference - Enhanced bootstrap (00_BOOTSTRAP_V2.md) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
463
optimization_engine/context/session_state.py
Normal file
463
optimization_engine/context/session_state.py
Normal file
@@ -0,0 +1,463 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user