| 49 | super().__init__(types=enum_types) |
| 50 | |
| 51 | def assess( |
| 52 | self, |
| 53 | model_callback: CallbackType, |
| 54 | purpose: Optional[str] = None, |
| 55 | ) -> Dict[BiasType, List[RTTestCase]]: |
| 56 | # Validate model callback signature once |
| 57 | validate_model_callback_signature( |
| 58 | model_callback=model_callback, |
| 59 | async_mode=self.async_mode, |
| 60 | ) |
| 61 | |
| 62 | if self.async_mode: |
| 63 | loop = get_or_create_event_loop() |
| 64 | return loop.run_until_complete( |
| 65 | self.a_assess(model_callback=model_callback, purpose=purpose) |
| 66 | ) |
| 67 | |
| 68 | # Run simulated attacks |
| 69 | simulated_test_cases = self.simulate_attacks(purpose) |
| 70 | |
| 71 | results: Dict[BiasType, List[RTTestCase]] = {} |
| 72 | res: Dict[BiasType, BiasMetric] = {} |
| 73 | simulated_attacks: Dict[str, str] = {} |
| 74 | |
| 75 | for test_case in simulated_test_cases: |
| 76 | vuln_type = test_case.vulnerability_type |
| 77 | input_text = test_case.input |
| 78 | |
| 79 | output = model_callback(input_text) |
| 80 | |
| 81 | rt_test_case = RTTestCase( |
| 82 | vulnerability=test_case.vulnerability, |
| 83 | vulnerability_type=vuln_type, |
| 84 | attackMethod=test_case.attack_method, |
| 85 | riskCategory=getRiskCategory(vuln_type), |
| 86 | input=input_text, |
| 87 | actual_output=output, |
| 88 | ) |
| 89 | |
| 90 | metric = self._get_metric(vuln_type) |
| 91 | metric.measure(rt_test_case) |
| 92 | |
| 93 | rt_test_case.score = metric.score |
| 94 | rt_test_case.reason = metric.reason |
| 95 | |
| 96 | res[vuln_type] = metric |
| 97 | simulated_attacks[vuln_type.value] = input_text |
| 98 | |
| 99 | results.setdefault(vuln_type, []).append(rt_test_case) |
| 100 | |
| 101 | # Store results in instance for later access |
| 102 | self.res = res |
| 103 | self.simulated_attacks = simulated_attacks |
| 104 | |
| 105 | return results |
| 106 | |
| 107 | async def a_assess( |
| 108 | self, |