diff --git a/src/fuzzyai/cli.py b/src/fuzzyai/cli.py index 511f860..9681306 100644 --- a/src/fuzzyai/cli.py +++ b/src/fuzzyai/cli.py @@ -169,18 +169,27 @@ async def run_fuzzer(args: argparse.Namespace) -> None: await fuzzer.cleanup() return - await aiofiles.os.makedirs(f'results/{CURRENT_TIMESTAMP}', exist_ok=True) + if isinstance(args.model, (list, set)) and args.model: + model_str = next(iter(args.model)) + else: + model_str = str(args.model) + + model_name = model_str.split("/")[-1] + + folder = f'results/{model_name}/{attack_mode+'__'+CURRENT_TIMESTAMP}' + + await aiofiles.os.makedirs(folder, exist_ok=True) if report.attacking_techniques and any(atp.total_prompts_count > 0 for atp in report.attacking_techniques): if raw_results: - logger.info(f"Dumping raw results to results/{CURRENT_TIMESTAMP}/raw.jsonl") - async with aiofiles.open(f"results/{CURRENT_TIMESTAMP}/raw.jsonl", 'w', encoding="utf-8") as f: + logger.info(f"Dumping raw results to {folder}/raw.jsonl") + async with aiofiles.open(f"{folder}/raw.jsonl", 'w', encoding="utf-8") as f: for raw_result in raw_results: await f.write(raw_result.model_dump_json()) if report: - logger.info(f"Dumping results to results/{CURRENT_TIMESTAMP}/report.json") - async with aiofiles.open(f"results/{CURRENT_TIMESTAMP}/report.json", 'w', encoding="utf-8") as f: + logger.info(f"Dumping results to {folder}/report.json") + async with aiofiles.open(f"{folder}/report.json", 'w', encoding="utf-8") as f: await f.write(report.model_dump_json()) generate_report(report) diff --git a/src/fuzzyai/utils/utils.py b/src/fuzzyai/utils/utils.py index 83f78d2..69130f2 100644 --- a/src/fuzzyai/utils/utils.py +++ b/src/fuzzyai/utils/utils.py @@ -483,8 +483,10 @@ def generate_report(report: FuzzerResult) -> None: # Generate the HTML report using string formatting html_data = REPORT_TEMPLATE.format(report_data=json.dumps(report_data)) + directory_name = model_name.split('/')[1].split(':')[0] + # Save the report - output_path = f'results/{CURRENT_TIMESTAMP}/report.html' + output_path = f'results/{directory_name}/{mode+'__'+CURRENT_TIMESTAMP}/report.html' with open(output_path, 'w') as f: f.write(html_data)