diff --git a/garak/probes/atkgen.py b/garak/probes/atkgen.py index fc126c869..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( @@ -188,25 +195,34 @@ def probe(self, generator) -> List[garak.attempt.Attempt]: ) # send the challenge in the target language and get the response response = generator.generate(this_attempt.prompt) - if response is None or len(response) == 0: - response_text = "" + if ( + response is None + or len(response) == 0 + or response[0] is None + or response[0].text is None + ): + response_text = None else: - response_text = ( - response[0].text.strip() if response[0] is not None else "" - ) + response_text = response[0].text.strip() this_attempt.outputs = response # log the response - response_to_store = self.reverse_langprovider.get_text([response_text])[ - 0 - ] + response_to_store = ( + self.reverse_langprovider.get_text([response_text])[0] + if response_text is not None + else None + ) if ( self.reverse_langprovider.source_lang != self.reverse_langprovider.target_lang ): - this_attempt.reverse_translation_outputs = [ - garak.attempt.Message(response_to_store) - ] + # when the generator returns [None] and translation is active store store [None] + if response and response[0] is None: + this_attempt.reverse_translation_outputs = response + else: + this_attempt.reverse_translation_outputs = [ + garak.attempt.Message(response_to_store) + ] logging.debug("atkgen: model: %s", response_text) if output_is_conversation: print( @@ -218,12 +234,26 @@ def probe(self, generator) -> List[garak.attempt.Attempt]: calls_made += 1 # last_response needs to be in the attack model's language base update on `response_to_store` # check if the resp is empty or if it matches the previous resp - if not len(response_to_store) and not self.constructive_tension: + if ( + response_to_store + and not len(response_to_store) + and not self.constructive_tension + ): keep_going = False - if response_to_store == last_response: + # if response_to_store is None it will be coerced to "" on iteration hence the extra check + if (response_to_store == last_response) or ( + len(last_response) == 0 and not response_to_store + ): keep_going = False and not self.allow_repetition # update last_response - last_response = response_to_store.replace("\n", " ").strip() + last_response = ( + response_to_store.replace("\n", " ").strip() + if response_to_store + 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"