From f5715350ec884977a02fee9b9a9776c9e4c9a4e4 Mon Sep 17 00:00:00 2001 From: zhanz5 <1721123209@qq.com> Date: Mon, 8 Jun 2026 15:18:18 +0800 Subject: [PATCH] fix: msj_data prepare_prompts ignores dataset_names parameter and uses mutable default The prepare_prompts() function had two bugs: 1. Mutable default argument: dataset_names=[] is a Python anti-pattern that shares the same list object across calls. 2. Logic bug: the return statement iterated over dataset_map keys instead of the dataset_names parameter, always returning all datasets regardless of what was requested. Changes: - probe_data/msj_data.py: - Changed default from dataset_names=[] to dataset_names=None - Added None check to initialize empty list - When dataset_names is provided, only return matching datasets - When empty or None, return all datasets (backward compatible) - probe_data/test_msj_data.py: - Fixed test_dataset_contents assertion: passing 1 dataset name should return 1 result, not 2 (old test was verifying the buggy behavior) --- agentic_security/probe_data/msj_data.py | 8 ++++++-- agentic_security/probe_data/test_msj_data.py | 3 +-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/agentic_security/probe_data/msj_data.py b/agentic_security/probe_data/msj_data.py index 0b75f617..ea15b1f2 100644 --- a/agentic_security/probe_data/msj_data.py +++ b/agentic_security/probe_data/msj_data.py @@ -38,8 +38,10 @@ def load_dataset_generic(name, getter=lambda x: x["train"]["prompt"]): def prepare_prompts( - dataset_names=[], budget=-1, tools_inbox=None + dataset_names=None, budget=-1, tools_inbox=None ) -> list[ProbeDataset]: + if dataset_names is None: + dataset_names = [] # fka/awesome-chatgpt-prompts # data-is-better-together/10k_prompts_ranked # alespalla/chatbot_instruction_prompts @@ -51,4 +53,6 @@ def prepare_prompts( "fka/awesome-chatgpt-prompts" ), } - return [dataset_map[name] for name in dataset_map] + if not dataset_names: + return list(dataset_map.values()) + return [dataset_map[name] for name in dataset_names if name in dataset_map] diff --git a/agentic_security/probe_data/test_msj_data.py b/agentic_security/probe_data/test_msj_data.py index e81812b5..7f29d457 100644 --- a/agentic_security/probe_data/test_msj_data.py +++ b/agentic_security/probe_data/test_msj_data.py @@ -129,8 +129,7 @@ def test_dataset_contents(self, mock_load_dataset_generic): result = prepare_prompts( dataset_names=["data-is-better-together/10k_prompts_ranked"] ) - assert len(result) == 2 + assert len(result) == 1 assert all(isinstance(ds.prompts, list) for ds in result) assert all(isinstance(ds.metadata, dict) for ds in result) assert result[0].prompts == ["test prompt"] - assert result[1].prompts == ["another prompt"]