""" Context Builder Builds rich context prompts for Claude sessions based on mode and study. """ import json import sqlite3 from pathlib import Path from typing import Any, Dict, List, Literal, Optional # Atomizer root directory ATOMIZER_ROOT = Path(__file__).parent.parent.parent.parent.parent class ContextBuilder: """Builds context prompts for Claude sessions""" def __init__(self): self.atomizer_root = ATOMIZER_ROOT self.studies_dir = ATOMIZER_ROOT / "studies" def build( self, mode: Literal["user", "power"], study_id: Optional[str] = None, conversation_history: Optional[List[Dict[str, Any]]] = None, ) -> str: """ Build full system prompt with context. Args: mode: "user" for safe operations, "power" for full access study_id: Optional study name to provide context for conversation_history: Optional recent messages for continuity Returns: Complete system prompt string """ parts = [self._base_context(mode)] if study_id: parts.append(self._study_context(study_id)) else: parts.append(self._global_context()) if conversation_history: parts.append(self._conversation_context(conversation_history)) parts.append(self._mode_instructions(mode)) return "\n\n---\n\n".join(parts) def build_study_context(self, study_id: str) -> str: """Build just the study context (for updates)""" return self._study_context(study_id) def _base_context(self, mode: str) -> str: """Base identity and capabilities""" return f"""# Atomizer Assistant You are the Atomizer Assistant - an expert system for structural optimization using FEA. **Current Mode**: {mode.upper()} Your role: - Help engineers with FEA optimization workflows - Create, configure, and run optimization studies - Analyze results and provide insights - Explain FEA concepts and methodology Important guidelines: - Be concise and professional - Use technical language appropriate for engineers - You are "Atomizer Assistant", not a generic AI - Use the available MCP tools to perform actions - When asked about studies, use the appropriate tools to get real data """ def _study_context(self, study_id: str) -> str: """Context for a specific study""" study_dir = self.studies_dir / study_id if not study_dir.exists(): return f"# Current Study: {study_id}\n\n**Status**: Study directory not found." context = f"# Current Study: {study_id}\n\n" # Load configuration config_path = study_dir / "1_setup" / "optimization_config.json" if not config_path.exists(): config_path = study_dir / "optimization_config.json" if config_path.exists(): try: with open(config_path) as f: config = json.load(f) context += "## Configuration\n\n" # Design variables dvs = config.get("design_variables", []) if dvs: context += "**Design Variables:**\n" for dv in dvs[:10]: bounds = f"[{dv.get('lower', '?')}, {dv.get('upper', '?')}]" context += f"- {dv.get('name', 'unnamed')}: {bounds}\n" if len(dvs) > 10: context += f"- ... and {len(dvs) - 10} more\n" # Objectives objs = config.get("objectives", []) if objs: context += "\n**Objectives:**\n" for obj in objs: direction = obj.get("direction", "minimize") context += f"- {obj.get('name', 'unnamed')} ({direction})\n" # Constraints constraints = config.get("constraints", []) if constraints: context += "\n**Constraints:**\n" for c in constraints: context += f"- {c.get('name', 'unnamed')}: {c.get('type', 'unknown')}\n" # Method method = config.get("method", "TPE") max_trials = config.get("max_trials", "not set") context += f"\n**Method**: {method}, max_trials: {max_trials}\n" except (json.JSONDecodeError, IOError) as e: context += f"\n*Config file exists but could not be parsed: {e}*\n" # Check for results db_path = study_dir / "3_results" / "study.db" if db_path.exists(): try: conn = sqlite3.connect(db_path) count = conn.execute( "SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'" ).fetchone()[0] best = conn.execute(""" SELECT MIN(tv.value) FROM trial_values tv JOIN trials t ON tv.trial_id = t.trial_id WHERE t.state = 'COMPLETE' """).fetchone()[0] context += f"\n## Results Status\n\n" context += f"- **Trials completed**: {count}\n" if best is not None: context += f"- **Best objective**: {best:.6f}\n" conn.close() except Exception: pass return context def _global_context(self) -> str: """Context when no study is selected""" context = "# Available Studies\n\n" if self.studies_dir.exists(): studies = [ d.name for d in self.studies_dir.iterdir() if d.is_dir() and not d.name.startswith("_") ] if studies: context += "The following studies are available:\n\n" for name in sorted(studies)[:20]: context += f"- {name}\n" if len(studies) > 20: context += f"\n... and {len(studies) - 20} more\n" else: context += "No studies found. Use `create_study` tool to create one.\n" else: context += "Studies directory not found.\n" context += "\n## Quick Actions\n\n" context += "- **Create study**: Describe what you want to optimize\n" context += "- **List studies**: Use `list_studies` tool for details\n" context += "- **Open study**: Ask about a specific study by name\n" return context def _conversation_context(self, history: List[Dict[str, Any]]) -> str: """Recent conversation for continuity""" if not history: return "" context = "# Recent Conversation\n\n" for msg in history[-10:]: role = "User" if msg.get("role") == "user" else "Assistant" content = msg.get("content", "")[:500] if len(msg.get("content", "")) > 500: content += "..." context += f"**{role}**: {content}\n\n" return context def _mode_instructions(self, mode: str) -> str: """Mode-specific instructions""" if mode == "power": return """# Power Mode Instructions You have **full access** to Atomizer's codebase. You can: - Edit any file using `edit_file` tool - Create new files with `create_file` tool - Create new extractors with `create_extractor` tool - Run shell commands with `run_shell_command` tool - Search codebase with `search_codebase` tool - Commit and push changes **Use these powers responsibly.** Always explain what you're doing and why. For routine operations (list, status, run, analyze), use the standard tools. """ else: return """# User Mode Instructions You can help with optimization workflows: - Create and configure studies - Run optimizations - Analyze results - Generate reports - Explain FEA concepts **For code modifications**, suggest switching to Power Mode. Available tools: - `list_studies`, `get_study_status`, `create_study` - `run_optimization`, `stop_optimization`, `get_optimization_status` - `get_trial_data`, `analyze_convergence`, `compare_trials`, `get_best_design` - `generate_report`, `export_data` - `explain_physics`, `recommend_method`, `query_extractors` """