diff --git a/garak/probes/atkgen.py b/garak/probes/atkgen.py index 41684c367..659c67aa3 100644 --- a/garak/probes/atkgen.py +++ b/garak/probes/atkgen.py @@ -102,6 +102,7 @@ def probe(self, generator) -> List[garak.attempt.Attempt]: keep_going = True last_response = None + new_conversation = True last_attempt = None if not output_is_conversation: @@ -116,16 +117,22 @@ def probe(self, generator) -> List[garak.attempt.Attempt]: f" turn {t.n:02}: red teaming [{red_team_model_short[:10]:<10}]" ) # generate a challenge - if last_response == "" and self.constructive_tension: + if ( + last_response == "" + and self.constructive_tension + and not new_conversation + ): challenge = "" + else: - if last_response: - if self.use_only_first_sent: - last_response = last_response.split(". ")[ - 0 - ] # nltk sentence splitter install too unfriendly - else: - last_response = "" # maybe it's none or 0.0 or [] or goodness knows. let's get an empty string. + if new_conversation: + last_response = "" + new_conversation = False + + if self.use_only_first_sent: + last_response = last_response.split(". ")[ + 0 + ] # nltk sentence splitter install too unfriendly if len(self.red_team_prompt_template): query = self.red_team_prompt_template.replace( @@ -245,6 +252,9 @@ def probe(self, generator) -> List[garak.attempt.Attempt]: else None ) + if last_response is None: + keep_going = False + _config.transient.reportfile.write( json.dumps(this_attempt.as_dict()) + "\n" ) diff --git a/tests/probes/test_probes_atkgen.py b/tests/probes/test_probes_atkgen.py index 83b9ad5d4..c30223259 100644 --- a/tests/probes/test_probes_atkgen.py +++ b/tests/probes/test_probes_atkgen.py @@ -169,3 +169,25 @@ def test_atkgen_verbose_output(capsys): # Verify that attempts were created assert isinstance(result, list), "probe results should be a list" assert len(result) > 0, "probe should return at least one attempt" + + +def test_atkgen_nones(): + _config.load_base_config() + _config.plugins.probes["atkgen"]["generations"] = 1 # we only need one conversation + p = _plugins.load_plugin("probes.atkgen.Tox", config_root=garak._config) + p.max_calls_per_conv = 1 # we don't need a full conversation + g = _plugins.load_plugin("generators.test.Nones", config_root=garak._config) + + with tempfile.NamedTemporaryFile(mode="w+", encoding="utf-8") as temp_report_file: + _config.transient.reportfile = temp_report_file + _config.transient.report_filename = temp_report_file.name + result = p.probe(g) + + assert result is not None, "Malformed None result - should be full result object" + assert ( + len(result) == p.convs_per_generation + ), "generators returning Nones should still give correct cardinality of results" + assert result[0].outputs == [None], "generator Nones should be propagated back" + assert ( + result[0].prompt.turns[0].content.text is not None + ), "Attack text should be stored"