jlov7/Dynamic-Function-Calling-Agent

1"""2Robustness Testing for Dynamic Function-Calling Agent3 4Tests model stability with:51. Shuffled JSON key order62. Distractor text before schema73. Noisy prompts8 9Quick test that doesn't require retraining.10"""11 12import json13import random14from test_constrained_model import load_trained_model, constrained_json_generate, create_json_schema15 16def shuffle_json_keys(obj):17    """Recursively shuffle the order of keys in JSON objects"""18    if isinstance(obj, dict):19        items = list(obj.items())20        random.shuffle(items)21        return {k: shuffle_json_keys(v) for k, v in items}22    elif isinstance(obj, list):23        return [shuffle_json_keys(item) for item in obj]24    return obj25 26def add_distractor_text(schema_str):27    """Add distracting text before the schema"""28    distractors = [29        "Note: This is a complex API with many parameters.",30        "Important: Please review all requirements carefully.",31        "Warning: Some fields may be optional depending on context.",32        "Info: This function supports multiple data formats.",33        "Reminder: Check authentication before making calls."34    ]35    distractor = random.choice(distractors)36    return f"{distractor}\n\n{schema_str}"37 38def test_robustness():39    """Run robustness tests on the function calling agent"""40    print("🧪 Starting Robustness Tests...")41    42    # Load model43    model, tokenizer = load_trained_model()44    45    # Test schema46    base_schema = {47        "name": "get_weather_forecast",48        "description": "Get weather forecast for a location",49        "parameters": {50            "type": "object",51            "properties": {52                "location": {"type": "string", "description": "City name"},53                "days": {"type": "integer", "description": "Number of days", "minimum": 1},54                "units": {"type": "string", "enum": ["metric", "imperial"]},55                "include_hourly": {"type": "boolean", "default": False}56            },57            "required": ["location", "days"]58        }59    }60    61    test_queries = [62        "Get 3-day weather for Paris",63        "Weather forecast for Tokyo, 5 days, metric units",64        "I need the weather for London for the next week"65    ]66    67    results = {68        "baseline": [],69        "shuffled_keys": [],70        "with_distractors": [],71        "both_shuffled_and_distractors": []72    }73    74    print("\n🔍 Running test scenarios...")75    76    for query in test_queries:77        print(f"\n📝 Query: '{query}'")78        79        # 1. Baseline test80        schema = create_json_schema(base_schema)81        prompt = f"""<|im_start|>system82You are a helpful assistant that calls functions by responding with valid JSON when given a schema. Always respond with JSON function calls only, never prose.<|im_end|>83 84<schema>85{json.dumps(base_schema, indent=2)}86</schema>87 88<|im_start|>user89{query}<|im_end|>90<|im_start|>assistant91"""92        93        response, success, error = constrained_json_generate(model, tokenizer, prompt, schema)94        results["baseline"].append(success)95        print(f"  ✅ Baseline: {'✓' if success else '✗'}")96        97        # 2. Shuffled keys test98        shuffled_schema = shuffle_json_keys(base_schema)99        schema = create_json_schema(shuffled_schema)100        prompt = f"""<|im_start|>system101You are a helpful assistant that calls functions by responding with valid JSON when given a schema. Always respond with JSON function calls only, never prose.<|im_end|>102 103<schema>104{json.dumps(shuffled_schema, indent=2)}105</schema>106 107<|im_start|>user108{query}<|im_end|>109<|im_start|>assistant110"""111        112        response, success, error = constrained_json_generate(model, tokenizer, prompt, schema)113        results["shuffled_keys"].append(success)114        print(f"  🔀 Shuffled: {'✓' if success else '✗'}")115        116        # 3. Distractor text test117        schema = create_json_schema(base_schema)118        schema_with_distractor = add_distractor_text(json.dumps(base_schema, indent=2))119        prompt = f"""<|im_start|>system120You are a helpful assistant that calls functions by responding with valid JSON when given a schema. Always respond with JSON function calls only, never prose.<|im_end|>121 122<schema>123{schema_with_distractor}124</schema>125 126<|im_start|>user127{query}<|im_end|>128<|im_start|>assistant129"""130        131        response, success, error = constrained_json_generate(model, tokenizer, prompt, schema)132        results["with_distractors"].append(success)133        print(f"  🎭 Distractor: {'✓' if success else '✗'}")134        135        # 4. Both shuffled and distractors136        shuffled_schema = shuffle_json_keys(base_schema)137        schema = create_json_schema(shuffled_schema)138        schema_with_distractor = add_distractor_text(json.dumps(shuffled_schema, indent=2))139        prompt = f"""<|im_start|>system140You are a helpful assistant that calls functions by responding with valid JSON when given a schema. Always respond with JSON function calls only, never prose.<|im_end|>141 142<schema>143{schema_with_distractor}144</schema>145 146<|im_start|>user147{query}<|im_end|>148<|im_start|>assistant149"""150        151        response, success, error = constrained_json_generate(model, tokenizer, prompt, schema)152        results["both_shuffled_and_distractors"].append(success)153        print(f"  🔀🎭 Both: {'✓' if success else '✗'}")154    155    # Calculate success rates156    print("\n📊 Robustness Test Results:")157    print("=" * 50)158    159    for test_name, test_results in results.items():160        success_rate = (sum(test_results) / len(test_results)) * 100161        print(f"{test_name.replace('_', ' ').title()}: {success_rate:.1f}% ({sum(test_results)}/{len(test_results)})")162    163    print("\n🎯 Analysis:")164    baseline_rate = (sum(results["baseline"]) / len(results["baseline"])) * 100165    166    for test_name, test_results in results.items():167        if test_name != "baseline":168            test_rate = (sum(test_results) / len(test_results)) * 100169            diff = test_rate - baseline_rate170            status = "🟢" if diff >= -10 else "🟡" if diff >= -20 else "🔴"171            print(f"{status} {test_name.replace('_', ' ').title()}: {diff:+.1f}% vs baseline")172    173    return results174 175if __name__ == "__main__":176    test_robustness()