Merge branch 'main' into bugfix/encoding_payload_explosion

leondz · Jul 5, 2024 · 57d7cb0 · 57d7cb0
2 parents 2b4cd1d + e5d2458
commit 57d7cb0
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 16 deletions.
diff --git a/garak/exception.py b/garak/exception.py
@@ -30,3 +30,9 @@ class BadGeneratorException(GarakException):
     """Generator config/description is not usable"""
 
     pass
+
+
+class RateLimitHit(Exception):
+    """Raised when a rate limiting response is returned"""
+
+    pass
diff --git a/garak/generators/base.py b/garak/generators/base.py
@@ -11,6 +11,7 @@
 
 from garak import _config
 from garak.configurable import Configurable
+import garak.resources.theme
 
 
 class Generator(Configurable):
@@ -107,11 +108,12 @@ def generate(
         if generations_this_call == 1:
             outputs = self._call_model(prompt, 1)
 
-        if self.supports_multiple_generations:
+        elif self.supports_multiple_generations:
             outputs = self._call_model(prompt, generations_this_call)
 
         else:
             outputs = []
+
             if (
                 hasattr(_config.system, "parallel_requests")
                 and _config.system.parallel_requests
@@ -120,19 +122,25 @@ def generate(
             ):
                 from multiprocessing import Pool
 
-                bar = tqdm.tqdm(total=generations_this_call, leave=False)
-                bar.set_description(self.fullname[:55])
+                multi_generator_bar = tqdm.tqdm(
+                    total=generations_this_call,
+                    leave=False,
+                    colour=f"#{garak.resources.theme.GENERATOR_RGB}",
+                )
+                multi_generator_bar.set_description(self.fullname[:55])
 
                 with Pool(_config.system.parallel_requests) as pool:
                     for result in pool.imap_unordered(
                         self._call_model, [prompt] * generations_this_call
                     ):
                         outputs.append(result)
-                        bar.update(1)
+                        multi_generator_bar.update(1)
 
             else:
                 generation_iterator = tqdm.tqdm(
-                    list(range(generations_this_call)), leave=False
+                    list(range(generations_this_call)),
+                    leave=False,
+                    colour=f"#{garak.resources.theme.GENERATOR_RGB}",
                 )
                 generation_iterator.set_description(self.fullname[:55])
                 for i in generation_iterator:

diff --git a/garak/generators/rest.py b/garak/generators/rest.py
@@ -16,16 +16,10 @@
 from jsonpath_ng.exceptions import JsonPathParserError
 
 from garak import _config
-from garak.exception import APIKeyMissingError
+from garak.exception import APIKeyMissingError, RateLimitHit
 from garak.generators.base import Generator
 
 
-class RESTRateLimitError(Exception):
-    """Raised when a rate limiting response is returned"""
-
-    pass
-
-
 class RestGenerator(Generator):
     """Generic API interface for REST models
 
@@ -247,7 +241,7 @@ def _populate_template(
         return output.replace("$INPUT", self.escape_function(text))
 
     # we'll overload IOError as the rate limit exception
-    @backoff.on_exception(backoff.fibo, RESTRateLimitError, max_value=70)
+    @backoff.on_exception(backoff.fibo, RateLimitHit, max_value=70)
     def _call_model(
         self, prompt: str, generations_this_call: int = 1
     ) -> List[Union[str, None]]:
@@ -274,9 +268,7 @@ def _call_model(
         }
         resp = self.http_function(self.uri, **req_kArgs)
         if resp.status_code in self.ratelimit_codes:
-            raise RESTRateLimitError(
-                f"Rate limited: {resp.status_code} - {resp.reason}"
-            )
+            raise RateLimitHit(f"Rate limited: {resp.status_code} - {resp.reason}")
 
         elif str(resp.status_code)[0] == "3":
             raise NotImplementedError(