example_custom_rules.py

"""
Example: Custom Rules with LLM Rule Engine

This example demonstrates how to create and use custom rules generated by LLMs,
both from direct code and from Domain Specific Language (DSL) descriptions.
"""

import llm_abm as abm

def main():
    print("LLM Rule Engine Example")
    print("=" * 50)
    
    # Configuration for testing custom rules
    config = {
        "grid": {"width": 15, "height": 15},
        "agents": {
            "herbivore": {"count": 20, "energy": 25},
            "carnivore": {"count": 5, "energy": 40}
        }
    }
    
    # Create base model
    model = abm.create_model(config)
    print(f"Created model with {model['metrics']['total_agents']} agents")
    
    print("\n1. ADDING CUSTOM RULE FROM CODE")
    print("-" * 30)
    
    # Custom rule: Territorial behavior (written as code)
    territorial_code = """
def territorial_behavior(model, params):
    \"\"\"Agents become aggressive when too many others are nearby\"\"\"
    new_model = copy.deepcopy(model)
    
    territory_size = params.get('territory_size', 3)
    energy_boost = params.get('energy_boost', 5)
    
    for agent in new_model['agents']:
        if not agent['alive']:
            continue
            
        # Count nearby agents
        nearby_count = 0
        for other in new_model['agents']:
            if other['id'] != agent['id'] and other['alive']:
                dx = abs(agent['position']['x'] - other['position']['x'])
                dy = abs(agent['position']['y'] - other['position']['y'])
                if dx <= territory_size and dy <= territory_size:
                    nearby_count += 1
        
        # If territory is crowded, get energy boost (aggression)
        if nearby_count > 2:
            agent['energy'] += energy_boost
    
    return new_model
"""
    
    # Add the custom rule
    model = abm.add_custom_rule(model, "territorial_behavior", territorial_code, "code")
    print("✓ Added 'territorial_behavior' rule from code")
    
    print("\n2. ADDING CUSTOM RULE FROM DSL")
    print("-" * 30)
    
    # Custom rule: Foraging behavior (using DSL)
    foraging_dsl = {
        "description": "Hungry agents seek food and gain energy",
        "conditions": [
            "agent['energy'] < params.get('hunger_threshold', 15)",
            "agent['type'] == 'herbivore'"
        ],
        "actions": [
            "agent['energy'] += params.get('food_energy', 8)",
            "# Simulate foraging success"
        ],
        "parameters": {
            "hunger_threshold": 15,
            "food_energy": 8
        }
    }
    
    model = abm.add_custom_rule(model, "foraging", foraging_dsl, "dsl")
    print("✓ Added 'foraging' rule from DSL")
    
    print("\n3. LISTING ALL AVAILABLE RULES")
    print("-" * 30)
    
    # List custom rules
    custom_rules = abm.list_custom_rules()
    print("Custom rules:")
    for rule_name, metadata in custom_rules.items():
        print(f"  - {rule_name}: {metadata['description']}")
        print(f"    Signature: {metadata['signature']}")
    
    print("\n4. BUILDING SIMULATION WITH MIXED RULES")
    print("-" * 30)
    
    # Add built-in rules
    model = abm.add_rule(model, "random_movement", {})
    model = abm.add_rule(model, "predator_prey", {"predator": "carnivore", "prey": "herbivore"})
    
    # Add custom rules with parameters
    model = abm.add_rule(model, "territorial_behavior", {"territory_size": 2, "energy_boost": 3})
    model = abm.add_rule(model, "foraging", {"hunger_threshold": 20, "food_energy": 10})
    
    # Add lifecycle rules
    model = abm.add_rule(model, "energy_decay", {"species": "all", "rate": 2})
    model = abm.add_rule(model, "reproduction", {"species": "herbivore", "energy_threshold": 35})
    model = abm.add_rule(model, "death", {})
    
    print(f"Total rules in simulation: {len(model['rules'])}")
    for i, rule in enumerate(model['rules'], 1):
        rule_type = "Custom" if rule['name'] in custom_rules else "Built-in"
        print(f"  {i}. {rule['name']} ({rule_type})")
    
    print("\n5. RUNNING SIMULATION WITH CUSTOM RULES")
    print("-" * 30)
    
    # Run simulation
    print("Running 30 steps with custom rules...")
    results = abm.run(model, steps=30)
    
    print(f"Simulation completed after {results['final_step']} steps")
    print(f"Initial agents: {results['summary']['initial_agents']}")
    print(f"Final agents: {results['summary']['final_agents']}")
    print(f"Final counts: {results['summary']['final_counts']}")
    
    print("\n6. STREAMING WITH CUSTOM RULES")
    print("-" * 30)
    
    # Reset and stream with custom rules
    model = abm.create_model(config)
    model = abm.add_rule(model, "random_movement", {})
    model = abm.add_rule(model, "territorial_behavior", {"territory_size": 2})
    model = abm.add_rule(model, "foraging", {"hunger_threshold": 15})
    model = abm.add_rule(model, "energy_decay", {"rate": 1})
    model = abm.add_rule(model, "death", {})
    
    def custom_callback(state):
        herbivores = state['agent_counts'].get('herbivore', 0)
        carnivores = state['agent_counts'].get('carnivore', 0)
        print(f"  Step {state['step']:2d}: {herbivores} herbivores, {carnivores} carnivores")
    
    print("Streaming 15 steps with fast updates...")
    stream_count = 0
    for state in abm.run_stream(model, steps=15, delay=0.05, callback=custom_callback):
        stream_count += 1
        if state['total_agents'] == 0:
            break
    
    print(f"Streamed {stream_count} states successfully")
    
    print("\n7. CUSTOM RULE CODE SAFETY")
    print("-" * 30)
    
    # Demonstrate safety validation
    print("Testing rule code safety validation...")
    
    # Safe code should work
    safe_code = """
def safe_rule(model, params):
    import copy
    new_model = copy.deepcopy(model)
    for agent in new_model['agents']:
        if agent['alive']:
            agent['energy'] += 1
    return new_model
"""
    
    try:
        abm.add_custom_rule(model, "safe_rule", safe_code, "code")
        print("✓ Safe code accepted")
    except ValueError as e:
        print(f"✗ Safe code rejected: {e}")
    
    # Unsafe code should be rejected
    unsafe_code = """
def unsafe_rule(model, params):
    import os
    os.system('rm -rf /')  # Dangerous!
    return model
"""
    
    try:
        abm.add_custom_rule(model, "unsafe_rule", unsafe_code, "code")
        print("✗ Unsafe code accepted (THIS SHOULD NOT HAPPEN)")
    except ValueError as e:
        print(f"✓ Unsafe code properly rejected: {e}")
    
    print("\n" + "=" * 50)
    print("CUSTOM RULE ENGINE FEATURES DEMONSTRATED:")
    print("✓ LLM-generated rule code execution")
    print("✓ Domain Specific Language (DSL) rules")
    print("✓ Mixed built-in and custom rules")
    print("✓ Safe code validation and sandboxing")
    print("✓ Custom rule metadata and listing")
    print("✓ Streaming simulation with custom rules")
    print("✓ Parameter passing to custom rules")
    print("=" * 50)

if __name__ == "__main__":
    main()