@@ -114,7 +114,9 @@ async def process_prompt(
114114
115115 if response .status_code >= 400 :
116116 logger .error (f"HTTP { response .status_code } { response .content = } " )
117- fuzzer_state .add_error (module_name , prompt , response .status_code , response .text )
117+ fuzzer_state .add_error (
118+ module_name , prompt , response .status_code , response .text
119+ )
118120 return tokens , True
119121
120122 # Process successful response
@@ -124,7 +126,9 @@ async def process_prompt(
124126 # Check if the response indicates a refusal
125127 refused = refusal_heuristic (response .json ())
126128 if refused :
127- fuzzer_state .add_refusal (module_name , prompt , response .status_code , response_text )
129+ fuzzer_state .add_refusal (
130+ module_name , prompt , response .status_code , response_text
131+ )
128132
129133 fuzzer_state .add_output (module_name , prompt , response_text , refused )
130134 return tokens , refused
@@ -168,7 +172,10 @@ async def process_prompt_batch(
168172 - Total number of tokens processed.
169173 - Number of failed prompts.
170174 """
171- tasks = [process_prompt (request_factory , p , tokens , module_name , fuzzer_state ) for p in prompts ]
175+ tasks = [
176+ process_prompt (request_factory , p , tokens , module_name , fuzzer_state )
177+ for p in prompts
178+ ]
172179 results = await asyncio .gather (* tasks )
173180 total_tokens = sum (r [0 ] for r in results )
174181 failures = sum (1 for r in results if r [1 ])
@@ -212,7 +219,11 @@ async def scan_module(
212219
213220 # Initialize optimizer if optimization is enabled
214221 optimizer = (
215- Optimizer ([Real (0 , 1 )], base_estimator = "GP" , n_initial_points = INITIAL_OPTIMIZER_POINTS ) if optimize else None
222+ Optimizer (
223+ [Real (0 , 1 )], base_estimator = "GP" , n_initial_points = INITIAL_OPTIMIZER_POINTS
224+ )
225+ if optimize
226+ else None
216227 )
217228
218229 module_size = 0 if module .lazy else len (module .prompts )
@@ -544,7 +555,9 @@ async def perform_many_shot_scan(
544555 ).model_dump_json ()
545556
546557 if optimize and len (failure_rates ) >= MIN_FAILURE_SAMPLES :
547- yield ScanResult .status_msg (f"High failure rate detected ({ failure_rate :.2%} ). Stopping this module..." )
558+ yield ScanResult .status_msg (
559+ f"High failure rate detected ({ failure_rate :.2%} ). Stopping this module..."
560+ )
548561 break
549562
550563 yield ScanResult .status_msg ("Scan completed." )
0 commit comments