Atomizer/tests/test_phase_3_2_llm_mode.py

"""
Test Phase 3.2: LLM Mode Integration

Tests the new generic run_optimization.py with --llm flag support.

This test verifies:
1. Natural language request parsing with LLM
2. Workflow generation (engineering features, calculations, hooks)
3. Integration with LLMOptimizationRunner
4. Argument parsing and validation

Author: Antoine Letarte
Date: 2025-11-17
"""

import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from optimization_engine.llm_workflow_analyzer import LLMWorkflowAnalyzer


def test_llm_workflow_analysis():
    """Test that LLM can analyze a natural language optimization request."""
    print("=" * 80)
    print("Test: LLM Workflow Analysis")
    print("=" * 80)
    print()

    # Natural language request (same as bracket study)
    request = """
    Maximize displacement while ensuring safety factor is greater than 4.

    Material: Aluminum 6061-T6 with yield strength of 276 MPa
    Design variables:
    - tip_thickness: 15 to 25 mm
    - support_angle: 20 to 40 degrees

    Run 20 trials using TPE algorithm.
    """

    print("Natural Language Request:")
    print(request)
    print()

    # Initialize analyzer (using Claude Code integration)
    print("Initializing LLM Workflow Analyzer (Claude Code mode)...")
    analyzer = LLMWorkflowAnalyzer(use_claude_code=True)
    print()

    # Analyze request
    print("Analyzing request with LLM...")
    print("(This will call Claude Code to parse the natural language)")
    print()

    try:
        workflow = analyzer.analyze_request(request)

        print("=" * 80)
        print("LLM Analysis Results")
        print("=" * 80)
        print()

        # Engineering features
        print(f"Engineering Features ({len(workflow.get('engineering_features', []))}):")
        for i, feature in enumerate(workflow.get('engineering_features', []), 1):
            print(f"  {i}. {feature.get('action')}: {feature.get('description')}")
            print(f"     Domain: {feature.get('domain')}")
            print(f"     Params: {feature.get('params')}")
            print()

        # Inline calculations
        print(f"Inline Calculations ({len(workflow.get('inline_calculations', []))}):")
        for i, calc in enumerate(workflow.get('inline_calculations', []), 1):
            print(f"  {i}. {calc.get('action')}")
            print(f"     Params: {calc.get('params')}")
            print(f"     Code hint: {calc.get('code_hint')}")
            print()

        # Post-processing hooks
        print(f"Post-Processing Hooks ({len(workflow.get('post_processing_hooks', []))}):")
        for i, hook in enumerate(workflow.get('post_processing_hooks', []), 1):
            print(f"  {i}. {hook.get('action')}")
            print(f"     Params: {hook.get('params')}")
            print()

        # Optimization config
        opt_config = workflow.get('optimization', {})
        print("Optimization Configuration:")
        print(f"  Algorithm: {opt_config.get('algorithm')}")
        print(f"  Direction: {opt_config.get('direction')}")
        print(f"  Design Variables ({len(opt_config.get('design_variables', []))}):")
        for var in opt_config.get('design_variables', []):
            print(f"    - {var.get('parameter')}: {var.get('min')} to {var.get('max')} {var.get('units', '')}")
        print()

        print("=" * 80)
        print("TEST PASSED: LLM successfully analyzed the request!")
        print("=" * 80)
        print()

        return True

    except Exception as e:
        print()
        print("=" * 80)
        print(f"TEST FAILED: {e}")
        print("=" * 80)
        print()
        import traceback
        traceback.print_exc()
        return False


def test_argument_parsing():
    """Test that run_optimization.py argument parsing works."""
    print("=" * 80)
    print("Test: Argument Parsing")
    print("=" * 80)
    print()

    import subprocess

    # Test help message
    # Need to go up one directory since we're in tests/
    result = subprocess.run(
        ["python", "../optimization_engine/run_optimization.py", "--help"],
        capture_output=True,
        text=True,
        cwd=Path(__file__).parent
    )

    if result.returncode == 0 and "--llm" in result.stdout:
        print("[OK] Help message displays correctly")
        print("[OK] --llm flag is present")
        print()
        print("TEST PASSED: Argument parsing works!")
        return True
    else:
        print("[FAIL] Help message failed or --llm flag missing")
        print(result.stdout)
        print(result.stderr)
        return False


def main():
    """Run all tests."""
    print()
    print("=" * 80)
    print("PHASE 3.2 INTEGRATION TESTS")
    print("=" * 80)
    print()

    tests = [
        ("Argument Parsing", test_argument_parsing),
        ("LLM Workflow Analysis", test_llm_workflow_analysis),
    ]

    results = []
    for test_name, test_func in tests:
        print()
        passed = test_func()
        results.append((test_name, passed))

    # Summary
    print()
    print("=" * 80)
    print("TEST SUMMARY")
    print("=" * 80)
    for test_name, passed in results:
        status = "[PASSED]" if passed else "[FAILED]"
        print(f"{status}: {test_name}")
    print()

    all_passed = all(passed for _, passed in results)
    if all_passed:
        print("All tests passed!")
    else:
        print("Some tests failed")

    return all_passed


if __name__ == '__main__':
    success = main()
    sys.exit(0 if success else 1)