simple_test.py

#!/usr/bin/env python3
"""
Simple test script to verify the Adaptive CoT framework works.
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.models.model_factory import ModelFactory
from src.adaptive.adaptive_cot import AdaptiveCoT


def test_basic_functionality():
    """Test basic functionality with a simple problem."""
    print("🧪 Testing Basic Adaptive CoT Functionality")
    print("=" * 50)
    
    # Simple config
    config = {
        "adaptive_branching": {
            "enabled": True,
            "min_branches": 1,
            "max_branches": 5,
            "default_branches": 3
        },
        "research_logging": False
    }
    
    try:
        # Create model
        print("📦 Creating model...")
        model = ModelFactory.create_model("generic", "/raid/LLM/llama3.1-8b-instruct", {})
        model.load_model()
        print("✅ Model loaded successfully")
        
        # Create AdaptiveCoT
        print("🔧 Creating AdaptiveCoT...")
        cot = AdaptiveCoT(model, config)
        print("✅ AdaptiveCoT created successfully")
        
        # Test problem
        problem = "What is 2+2?"
        print(f"\n🔍 Testing problem: {problem}")
        
        # Solve problem
        result = cot.solve_problem(problem)
        
        # Display results
        print(f"\n📊 Results:")
        print(f"   Final Answer: {result['final_answer']}")
        print(f"   Branches Used: {result['num_branches']}")
        print(f"   Strategy: {result['allocation_info']['strategy']}")
        print(f"   Execution Time: {result['execution_time']:.2f}s")
        
        if 'prefill_signals' in result:
            prefill = result['prefill_signals']
            print(f"   Prefill Signals:")
            print(f"     Entropy: {prefill['entropy']:.3f}")
            print(f"     KL Divergence: {prefill['kl_divergence']:.3f}")
            print(f"     Confidence: {prefill['confidence']:.3f}")
        
        print(f"\n✅ Test completed successfully!")
        return True
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = test_basic_functionality()
    sys.exit(0 if success else 1)