Atomizer/tests/test_llm_complex_request.py

"""
Test LLM-Powered Workflow Analyzer with Complex Invented Request

This test uses a realistic, complex optimization scenario combining:
- Multiple result types (stress, displacement, mass)
- Composite materials (PCOMP)
- Custom constraints
- Multi-objective optimization
- Post-processing calculations

Author: Atomizer Development Team
Version: 0.1.0 (Phase 2.7)
Last Updated: 2025-01-16
"""

import sys
import os
import json
from pathlib import Path

# Set UTF-8 encoding for Windows console
if sys.platform == 'win32':
    import codecs
    if not isinstance(sys.stdout, codecs.StreamWriter):
        if hasattr(sys.stdout, 'buffer'):
            sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, errors='replace')
            sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, errors='replace')

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.llm_workflow_analyzer import LLMWorkflowAnalyzer


def main():
    # Complex invented optimization request
    user_request = """I want to optimize a composite panel structure.

First, I need to extract the maximum von Mises stress from solution 2 subcase 1, and also get the
maximum displacement in Y direction from the same subcase. Then I want to calculate the total mass
using the part expression called 'panel_total_mass' which accounts for all the PCOMP plies.

For my objective function, I want to minimize a weighted combination where stress contributes 70%
and displacement contributes 30%. The combined metric should be normalized by dividing stress by
200 MPa and displacement by 5 mm before applying the weights.

I also need a constraint: keep the displacement under 3.5 mm, and make sure the mass doesn't
increase by more than 10% compared to the baseline which is stored in the expression 'baseline_mass'.

For optimization, I want to vary the ply thicknesses of my PCOMP layup that have the suffix '_design'
in their ply IDs. I want to use Optuna with TPE sampler and run 150 trials.

Can you help me set this up?"""

    print('=' * 80)
    print('PHASE 2.7 TEST: LLM Analysis of Complex Composite Optimization')
    print('=' * 80)
    print()
    print('INVENTED OPTIMIZATION REQUEST:')
    print('-' * 80)
    print(user_request)
    print()
    print('=' * 80)
    print()

    # Check for API key
    api_key = os.environ.get('ANTHROPIC_API_KEY')

    if not api_key:
        print('⚠️  ANTHROPIC_API_KEY not found in environment')
        print()
        print('To run LLM analysis, set your API key:')
        print('  Windows: set ANTHROPIC_API_KEY=your_key_here')
        print('  Linux/Mac: export ANTHROPIC_API_KEY=your_key_here')
        print()
        print('For now, showing EXPECTED intelligent analysis...')
        print()

        # Show what LLM SHOULD detect
        show_expected_analysis()
        return

    # Use LLM to analyze
    print('[1] Calling Claude LLM for Intelligent Analysis...')
    print('-' * 80)
    print()

    analyzer = LLMWorkflowAnalyzer(api_key=api_key)

    try:
        analysis = analyzer.analyze_request(user_request)

        print('✅ LLM Analysis Complete!')
        print()
        print('=' * 80)
        print('INTELLIGENT WORKFLOW BREAKDOWN')
        print('=' * 80)
        print()

        # Display summary
        print(analyzer.get_summary(analysis))

        print()
        print('=' * 80)
        print('DETAILED JSON ANALYSIS')
        print('=' * 80)
        print(json.dumps(analysis, indent=2))
        print()

        # Analyze what LLM detected
        print()
        print('=' * 80)
        print('INTELLIGENCE VALIDATION')
        print('=' * 80)
        print()

        validate_intelligence(analysis)

    except Exception as e:
        print(f'❌ Error calling LLM: {e}')
        import traceback
        traceback.print_exc()


def show_expected_analysis():
    """Show what the LLM SHOULD intelligently detect."""
    print('=' * 80)
    print('EXPECTED LLM ANALYSIS (What Intelligence Should Detect)')
    print('=' * 80)
    print()

    expected = {
        "engineering_features": [
            {
                "action": "extract_von_mises_stress",
                "domain": "result_extraction",
                "description": "Extract maximum von Mises stress from OP2 file",
                "params": {
                    "result_type": "von_mises_stress",
                    "metric": "maximum",
                    "solution": 2,
                    "subcase": 1
                },
                "why_engineering": "Requires pyNastran to read OP2 binary format"
            },
            {
                "action": "extract_displacement_y",
                "domain": "result_extraction",
                "description": "Extract maximum Y displacement from OP2 file",
                "params": {
                    "result_type": "displacement",
                    "direction": "Y",
                    "metric": "maximum",
                    "solution": 2,
                    "subcase": 1
                },
                "why_engineering": "Requires pyNastran OP2 extraction"
            },
            {
                "action": "read_panel_mass_expression",
                "domain": "geometry",
                "description": "Read panel_total_mass expression from .prt file",
                "params": {
                    "expression_name": "panel_total_mass",
                    "source": "part_file"
                },
                "why_engineering": "Requires NX API to read part expressions"
            },
            {
                "action": "read_baseline_mass_expression",
                "domain": "geometry",
                "description": "Read baseline_mass expression for constraint",
                "params": {
                    "expression_name": "baseline_mass",
                    "source": "part_file"
                },
                "why_engineering": "Requires NX API to read part expressions"
            },
            {
                "action": "update_pcomp_ply_thicknesses",
                "domain": "fea_properties",
                "description": "Modify PCOMP ply thicknesses with _design suffix",
                "params": {
                    "property_type": "PCOMP",
                    "parameter_filter": "_design",
                    "property": "ply_thickness"
                },
                "why_engineering": "Requires understanding of PCOMP card format and NX API"
            }
        ],
        "inline_calculations": [
            {
                "action": "normalize_stress",
                "description": "Normalize stress by 200 MPa",
                "params": {
                    "input": "max_stress",
                    "divisor": 200.0,
                    "units": "MPa"
                },
                "code_hint": "norm_stress = max_stress / 200.0"
            },
            {
                "action": "normalize_displacement",
                "description": "Normalize displacement by 5 mm",
                "params": {
                    "input": "max_disp_y",
                    "divisor": 5.0,
                    "units": "mm"
                },
                "code_hint": "norm_disp = max_disp_y / 5.0"
            },
            {
                "action": "calculate_mass_increase",
                "description": "Calculate mass increase percentage vs baseline",
                "params": {
                    "current": "panel_total_mass",
                    "baseline": "baseline_mass"
                },
                "code_hint": "mass_increase_pct = ((panel_total_mass - baseline_mass) / baseline_mass) * 100"
            }
        ],
        "post_processing_hooks": [
            {
                "action": "weighted_objective_function",
                "description": "Combine normalized stress (70%) and displacement (30%)",
                "params": {
                    "inputs": ["norm_stress", "norm_disp"],
                    "weights": [0.7, 0.3],
                    "formula": "0.7 * norm_stress + 0.3 * norm_disp",
                    "objective": "minimize"
                },
                "why_hook": "Custom weighted combination of multiple normalized metrics"
            }
        ],
        "constraints": [
            {
                "type": "displacement_limit",
                "parameter": "max_disp_y",
                "condition": "<=",
                "value": 3.5,
                "units": "mm"
            },
            {
                "type": "mass_increase_limit",
                "parameter": "mass_increase_pct",
                "condition": "<=",
                "value": 10.0,
                "units": "percent"
            }
        ],
        "optimization": {
            "algorithm": "optuna",
            "sampler": "TPE",
            "trials": 150,
            "design_variables": [
                {
                    "parameter_type": "pcomp_ply_thickness",
                    "filter": "_design",
                    "property_card": "PCOMP"
                }
            ],
            "objectives": [
                {
                    "type": "minimize",
                    "target": "weighted_objective_function"
                }
            ]
        },
        "summary": {
            "total_steps": 11,
            "engineering_features": 5,
            "inline_calculations": 3,
            "post_processing_hooks": 1,
            "constraints": 2,
            "complexity": "high",
            "multi_objective": "weighted_combination"
        }
    }

    # Print formatted analysis
    print('Engineering Features (Need Research): 5')
    print('  1. extract_von_mises_stress - OP2 extraction')
    print('  2. extract_displacement_y - OP2 extraction')
    print('  3. read_panel_mass_expression - NX part expression')
    print('  4. read_baseline_mass_expression - NX part expression')
    print('  5. update_pcomp_ply_thicknesses - PCOMP property modification')
    print()

    print('Inline Calculations (Auto-Generate): 3')
    print('  1. normalize_stress → norm_stress = max_stress / 200.0')
    print('  2. normalize_displacement → norm_disp = max_disp_y / 5.0')
    print('  3. calculate_mass_increase → mass_increase_pct = ...')
    print()

    print('Post-Processing Hooks (Generate Middleware): 1')
    print('  1. weighted_objective_function')
    print('     Formula: 0.7 * norm_stress + 0.3 * norm_disp')
    print('     Objective: minimize')
    print()

    print('Constraints: 2')
    print('  1. max_disp_y <= 3.5 mm')
    print('  2. mass_increase <= 10%')
    print()

    print('Optimization:')
    print('  Algorithm: Optuna TPE')
    print('  Trials: 150')
    print('  Design Variables: PCOMP ply thicknesses with _design suffix')
    print()

    print('=' * 80)
    print('INTELLIGENCE ASSESSMENT')
    print('=' * 80)
    print()
    print('What makes this INTELLIGENT (not dumb regex):')
    print()
    print('  ✓ Detected solution 2 subcase 1 (specific subcase targeting)')
    print('  ✓ Distinguished OP2 extraction vs part expression reading')
    print('  ✓ Identified PCOMP as composite material requiring special handling')
    print('  ✓ Recognized weighted combination as post-processing hook')
    print('  ✓ Understood normalization as simple inline calculation')
    print('  ✓ Detected constraint logic (displacement limit, mass increase %)')
    print('  ✓ Identified TPE sampler specifically (not just "Optuna")')
    print('  ✓ Understood _design suffix as parameter filter')
    print('  ✓ Separated engineering features from trivial math')
    print()
    print('This level of understanding requires LLM intelligence!')
    print()


def validate_intelligence(analysis):
    """Validate that LLM detected key intelligent aspects."""
    print('Checking LLM Intelligence...')
    print()

    checks = []

    # Check 1: Multiple result extractions
    eng_features = analysis.get('engineering_features', [])
    result_extractions = [f for f in eng_features if 'extract' in f.get('action', '').lower()]
    checks.append(('Multiple result extractions detected', len(result_extractions) >= 2))

    # Check 2: Normalization calculations
    inline_calcs = analysis.get('inline_calculations', [])
    normalizations = [c for c in inline_calcs if 'normal' in c.get('action', '').lower()]
    checks.append(('Normalization calculations detected', len(normalizations) >= 2))

    # Check 3: Weighted combination hook
    hooks = analysis.get('post_processing_hooks', [])
    weighted = [h for h in hooks if 'weight' in h.get('description', '').lower()]
    checks.append(('Weighted combination hook detected', len(weighted) >= 1))

    # Check 4: PCOMP understanding
    pcomp_features = [f for f in eng_features if 'pcomp' in str(f).lower()]
    checks.append(('PCOMP composite understanding', len(pcomp_features) >= 1))

    # Check 5: Constraints
    constraints = analysis.get('constraints', []) or []
    checks.append(('Constraints detected', len(constraints) >= 2))

    # Check 6: Optuna configuration
    opt = analysis.get('optimization', {})
    has_optuna = 'optuna' in str(opt).lower()
    checks.append(('Optuna optimization detected', has_optuna))

    # Print results
    for check_name, passed in checks:
        status = '✅' if passed else '❌'
        print(f'  {status} {check_name}')

    print()
    passed_count = sum(1 for _, p in checks if p)
    total_count = len(checks)

    if passed_count == total_count:
        print(f'🎉 Perfect! LLM detected {passed_count}/{total_count} intelligent aspects!')
    elif passed_count >= total_count * 0.7:
        print(f'✅ Good! LLM detected {passed_count}/{total_count} intelligent aspects')
    else:
        print(f'⚠️  Needs improvement: {passed_count}/{total_count} aspects detected')
    print()


if __name__ == '__main__':
    main()