Anto01
a24e3f750c
Core plugin architecture for LLM-driven optimization: New Features: - Hook system with 6 lifecycle points (pre_mesh, post_mesh, pre_solve, post_solve, post_extraction, custom_objectives) - HookManager for centralized registration and execution - Code validation with AST-based safety checks - Feature registry (JSON) for LLM capability discovery - Example plugin: log_trial_start - 23 comprehensive tests (all passing) Integration: - OptimizationRunner now loads plugins automatically - Hooks execute at 5 points in optimization loop - Custom objectives can override total_objective via hooks Safety: - Module whitelist (numpy, scipy, pandas, optuna, pyNastran) - Dangerous operation blocking (eval, exec, os.system, subprocess) - Optional file operation permission flag Files Added: - optimization_engine/plugins/__init__.py - optimization_engine/plugins/hooks.py - optimization_engine/plugins/hook_manager.py - optimization_engine/plugins/validators.py - optimization_engine/feature_registry.json - optimization_engine/plugins/pre_solve/log_trial_start.py - tests/test_plugin_system.py (23 tests) Files Modified: - optimization_engine/runner.py (added hook integration) Ready for Phase 2: LLM interface layer 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
"""
|
|
Core Hook System for Atomizer
|
|
|
|
Defines hook points in the optimization lifecycle and hook registration mechanism.
|
|
"""
|
|
|
|
from enum import Enum
|
|
from typing import Callable, Dict, Any, Optional
|
|
from dataclasses import dataclass
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class HookPoint(Enum):
|
|
"""Enumeration of available hook points in the optimization lifecycle."""
|
|
|
|
PRE_MESH = "pre_mesh" # Before meshing
|
|
POST_MESH = "post_mesh" # After meshing, before solve
|
|
PRE_SOLVE = "pre_solve" # Before solver execution
|
|
POST_SOLVE = "post_solve" # After solve, before extraction
|
|
POST_EXTRACTION = "post_extraction" # After result extraction
|
|
CUSTOM_OBJECTIVE = "custom_objective" # Custom objective functions
|
|
|
|
|
|
@dataclass
|
|
class Hook:
|
|
"""
|
|
Represents a single hook function to be executed at a specific point.
|
|
|
|
Attributes:
|
|
name: Unique identifier for this hook
|
|
hook_point: When this hook should execute (HookPoint enum)
|
|
function: The callable to execute
|
|
description: Human-readable description of what this hook does
|
|
priority: Execution order (lower = earlier, default=100)
|
|
enabled: Whether this hook is currently active
|
|
"""
|
|
|
|
name: str
|
|
hook_point: HookPoint
|
|
function: Callable[[Dict[str, Any]], Optional[Dict[str, Any]]]
|
|
description: str
|
|
priority: int = 100
|
|
enabled: bool = True
|
|
|
|
def execute(self, context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Execute this hook with the given context.
|
|
|
|
Args:
|
|
context: Dictionary containing relevant data for this hook point
|
|
Common keys:
|
|
- trial_number: Current trial number
|
|
- design_variables: Current design variable values
|
|
- sim_file: Path to simulation file
|
|
- working_dir: Current working directory
|
|
|
|
Returns:
|
|
Optional dictionary with results or modifications to context
|
|
|
|
Raises:
|
|
Exception: Any exception from the hook function is logged and re-raised
|
|
"""
|
|
if not self.enabled:
|
|
logger.debug(f"Hook '{self.name}' is disabled, skipping")
|
|
return None
|
|
|
|
try:
|
|
logger.info(f"Executing hook '{self.name}' at {self.hook_point.value}")
|
|
result = self.function(context)
|
|
logger.debug(f"Hook '{self.name}' completed successfully")
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"Hook '{self.name}' failed: {e}", exc_info=True)
|
|
raise
|
|
|
|
def __repr__(self) -> str:
|
|
status = "enabled" if self.enabled else "disabled"
|
|
return f"Hook(name='{self.name}', point={self.hook_point.value}, priority={self.priority}, {status})"
|
|
|
|
|
|
class HookContext:
|
|
"""
|
|
Context object passed to hooks containing all relevant data.
|
|
|
|
This is a convenience wrapper around a dictionary that provides
|
|
both dict-like access and attribute access.
|
|
"""
|
|
|
|
def __init__(self, **kwargs):
|
|
self._data = kwargs
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
return self._data[key]
|
|
|
|
def __setitem__(self, key: str, value: Any):
|
|
self._data[key] = value
|
|
|
|
def __contains__(self, key: str) -> bool:
|
|
return key in self._data
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
return self._data.get(key, default)
|
|
|
|
def update(self, other: Dict[str, Any]):
|
|
self._data.update(other)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return self._data.copy()
|
|
|
|
def __repr__(self) -> str:
|
|
keys = list(self._data.keys())
|
|
return f"HookContext({', '.join(keys)})"
|