type check generators for Turn patterns

Author: leondz

garak/generators/huggingface.py

@@ -242,7 +242,9 @@ def __init__(self, name="", config_root=_config):
         ),
         max_value=125,
     )
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         import json
         import requests
 
@@ -350,7 +352,9 @@ def __init__(self, name="", config_root=_config):
         ),
         max_value=125,
     )
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         import requests
 
         payload = {
@@ -440,7 +444,9 @@ def _clear_client(self):
         self.tokenizer = None
         self.generation_config = None
 
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         self._load_client()
         self.generation_config.max_new_tokens = self.max_tokens
         self.generation_config.do_sample = self.hf_args["do_sample"]

garak/generators/replicate.py

@@ -79,12 +79,12 @@ class InferenceEndpoint(ReplicateGenerator):
         backoff.fibo, replicate.exceptions.ReplicateError, max_value=70
     )
     def _call_model(
-        self, prompt, generations_this_call: int = 1
-    ) -> List[Union[str, None]]:
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Union[Turn, None]]:
         deployment = self.replicate.deployments.get(self.name)
         prediction = deployment.predictions.create(
             input={
-                "prompt": prompt,
+                "prompt": prompt.text,
                 "max_length": self.max_tokens,
                 "temperature": self.temperature,
                 "top_p": self.top_p,
@@ -98,7 +98,7 @@ def _call_model(
             raise IOError(
                 "Replicate endpoint didn't generate a response. Make sure the endpoint is active."
             ) from exc
-        return [response]
+        return [Turn(r) for r in response]
 
 
 DEFAULT_CLASS = "ReplicateGenerator"

garak/generators/test.py

@@ -18,7 +18,9 @@ class Blank(Generator):
     generator_family_name = "Test"
     name = "Blank"
 
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         return [Turn("")] * generations_this_call
 
 
@@ -29,7 +31,9 @@ class Repeat(Generator):
     generator_family_name = "Test"
     name = "Repeat"
 
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         return [prompt] * generations_this_call
 
 
@@ -41,7 +45,9 @@ class Single(Generator):
     name = "Single"
     test_generation_string = "ELIM"
 
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         if generations_this_call == 1:
             return [Turn(self.test_generation_string)]
         else:
@@ -71,7 +77,9 @@ class BlankVision(Generator):
     name = "BlankVision"
     modality = {"in": {"text", "image"}, "out": {"text"}}
 
-    def _call_model(self, prompt: Turn, generations_this_call: int = 1) -> List[Turn]:
+    def _call_model(
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Turn | None]:
         return [Turn("")] * generations_this_call
 
 

garak/generators/watsonx.py

@@ -126,14 +126,14 @@ def _validate_env_var(self):
         return super()._validate_env_var()
 
     def _call_model(
-        self, prompt: str, generations_this_call: int = 1
-    ) -> List[Union[str, None]]:
+        self, prompt: Turn, generations_this_call: int = 1
+    ) -> List[Union[Turn, None]]:
         if not self.bearer_token:
             self._set_bearer_token()
 
         # Check if message is empty. If it is, append null byte.
-        if not prompt:
-            prompt = "\x00"
+        if not prompt or not prompt.text:
+            prompt = Turn("\x00")
             print(
                 "WARNING: Empty prompt was found. Null byte character appended to prevent API failure."
             )

tests/generators/test_generators.py

@@ -6,6 +6,8 @@
 import pytest
 import random
 
+from typing import List, Union
+
 from garak import _plugins
 from garak import _config
 from garak.attempt import Turn
@@ -222,3 +224,31 @@ def test_instantiate_generators(classname):
     m = importlib.import_module("garak." + ".".join(classname.split(".")[:-1]))
     g = getattr(m, classname.split(".")[-1])(config_root=config_root)
     assert isinstance(g, Generator)
+
+
+NON_CONVERSATION_GENERATORS = [
+    classname
+    for classname in GENERATORS
+    if not ("openai" in classname or "groq" in classname or "azure" in classname)
+]
+
+
+@pytest.mark.parametrize("classname", NON_CONVERSATION_GENERATORS)
+def test_generator_signature(classname):
+    _, namespace, klass = classname.split(".")
+    m = importlib.import_module(f"garak.generators.{namespace}")
+    g = getattr(m, klass)
+    generate_signature = inspect.signature(g.generate)
+    assert (
+        generate_signature.parameters.get("prompt").annotation == Turn
+    ), "generate should take a Turn and return list of Turns or Nones"
+    assert (
+        generate_signature.return_annotation == List[Union[None, Turn]]
+    ), "generate should take a Turn and return list of Turns or Nones"
+    _call_model_signature = inspect.signature(g._call_model)
+    assert (
+        _call_model_signature.parameters.get("prompt").annotation == Turn
+    ), "_call_model should take a Turn and return list of Turns or Nones"
+    assert (
+        _call_model_signature.return_annotation == List[Union[None, Turn]]
+    ), "_call_model should take a Turn and return list of Turns or Nones"