handle target model returning None in atkgen (#678)

* handle target model returning None in atkgen * validate response structure without try/except
NVIDIA · May 14, 2024 · dc4a1c1 · dc4a1c1
1 parent 9a59e21
commit dc4a1c1
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/garak/probes/atkgen.py b/garak/probes/atkgen.py
@@ -135,13 +135,12 @@ def probe(self, generator) -> List[garak.attempt.Attempt]:
                         f" turn {t.n:02}: waiting for [{generator.name[:10]:<10}]"
                     )
                 # send the challenge and get the response
-                try:
-                    response = generator.generate(challenge)[0].strip()
-                except AttributeError as ae:
-                    if generator.generate(challenge)[0] is None:
-                        response = ""
-                    else:
-                        raise AttributeError from ae
+                response = generator.generate(challenge)
+                if response is None or len(response) == 0:
+                    response = ""
+                else:
+                    response = response[0].strip() if response[0] is not None else ""
+
                 # log the response
                 turn = ("model", response)
                 turns.append(turn)