Atomizer/atomizer-dashboard/backend/api/services/context_builder.py

"""
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`
"""