| 105 | return results |
| 106 | |
| 107 | async def a_assess( |
| 108 | self, |
| 109 | model_callback: CallbackType, |
| 110 | purpose: Optional[str] = None, |
| 111 | ) -> Dict[BiasType, List[RTTestCase]]: |
| 112 | # Validate the async model callback |
| 113 | validate_model_callback_signature( |
| 114 | model_callback=model_callback, |
| 115 | async_mode=self.async_mode, |
| 116 | ) |
| 117 | |
| 118 | # Run simulated attack generation |
| 119 | simulated_test_cases = await self.a_simulate_attacks(purpose) |
| 120 | |
| 121 | results: Dict[BiasType, List[RTTestCase]] = {} |
| 122 | res: Dict[BiasType, BiasMetric] = {} |
| 123 | simulated_attacks: Dict[str, str] = {} |
| 124 | |
| 125 | async def process_attack(test_case: RTTestCase): |
| 126 | vuln_type = test_case.vulnerability_type |
| 127 | input_text = test_case.input |
| 128 | |
| 129 | output = await model_callback(input_text) |
| 130 | |
| 131 | rt_test_case = RTTestCase( |
| 132 | vulnerability=test_case.vulnerability, |
| 133 | vulnerability_type=vuln_type, |
| 134 | attackMethod=test_case.attack_method, |
| 135 | riskCategory=getRiskCategory(vuln_type), |
| 136 | input=input_text, |
| 137 | actual_output=output, |
| 138 | ) |
| 139 | |
| 140 | metric = self._get_metric(vuln_type) |
| 141 | await metric.a_measure(rt_test_case) |
| 142 | |
| 143 | rt_test_case.score = metric.score |
| 144 | rt_test_case.reason = metric.reason |
| 145 | |
| 146 | res[vuln_type] = metric |
| 147 | simulated_attacks[vuln_type.value] = input_text |
| 148 | |
| 149 | return vuln_type, rt_test_case |
| 150 | |
| 151 | # Run all processing concurrently for supported types |
| 152 | tasks = [ |
| 153 | process_attack(test_case) |
| 154 | for test_case in simulated_test_cases |
| 155 | if test_case.vulnerability_type in self.types |
| 156 | ] |
| 157 | |
| 158 | for coro in asyncio.as_completed(tasks): |
| 159 | vuln_type, test_case = await coro |
| 160 | results.setdefault(vuln_type, []).append(test_case) |
| 161 | |
| 162 | # Persist results |
| 163 | self.res = res |
| 164 | self.simulated_attacks = simulated_attacks |