diff --git a/garak/generators/guardrails.py b/garak/generators/guardrails.py index 96719808a..60eb70366 100644 --- a/garak/generators/guardrails.py +++ b/garak/generators/guardrails.py @@ -10,6 +10,7 @@ from garak import _config from garak.attempt import Message, Conversation from garak.generators.base import Generator +from garak.generators.openai import OpenAICompatible class NeMoGuardrails(Generator): @@ -51,4 +52,30 @@ def _call_model( return [None] +class NeMoGuardrailsServer(OpenAICompatible): + """Generator for NeMo Guardrails Server""" + + supports_multiple_generations = False + generator_family_name = "Guardrails" + + DEFAULT_PARAMS = OpenAICompatible.DEFAULT_PARAMS | { + "uri": "http://localhost:8000/v1/", + "config_ids": set(), + } + + def __init__(self, name="", config_root=...): + super().__init__(name, config_root) + if self.extra_params and not self.extra_params.get("extra_body"): + self.extra_params.append("extra_body") + + guardrails = None + if self.config_ids: + guardrails = {"config_ids": self.config_ids} + if guardrails: + if hasattr(self, "extra_body") and self.extra_body and self.config_ids: + self.extra_body["guardrails"] = guardrails + else: + self.extra_body = {"guardrails": guardrails} + + DEFAULT_CLASS = "NeMoGuardrails" diff --git a/tests/generators/test_generators.py b/tests/generators/test_generators.py index a49c768f5..1379610d2 100644 --- a/tests/generators/test_generators.py +++ b/tests/generators/test_generators.py @@ -131,7 +131,12 @@ def test_instantiate_generators(classname): NON_CONVERSATION_GENERATORS = [ classname for classname in GENERATORS - if not ("openai" in classname or "groq" in classname or "azure" in classname) + if not ( + "openai" in classname + or "groq" in classname + or "azure" in classname + or "NeMoGuardrailsServer" in classname + ) ]