generator: NeMoGuardrails server support

garak/configurable.py

@@ -109,7 +109,7 @@ def _load_config(self, config_root=_config):
                 self._apply_config(plugins_config[namespaced_klass])
         self._apply_run_defaults()
         self._apply_missing_instance_defaults()
-        if hasattr(self, "ENV_VAR"):
+        if hasattr(self, "ENV_VAR") and self.ENV_VAR:
             if not hasattr(self, "key_env_var"):
                 self.key_env_var = self.ENV_VAR
         self._validate_env_var()
@@ -167,7 +167,7 @@ def _apply_missing_instance_defaults(self):
                     setattr(self, k, v)
 
     def _validate_env_var(self):
-        if hasattr(self, "key_env_var"):
+        if hasattr(self, "key_env_var") and self.key_env_var:
             if not hasattr(self, "api_key") or self.api_key is None:
                 self.api_key = os.getenv(self.key_env_var, default=None)
                 if self.api_key is None:

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,37 @@ def _call_model(
             return [None]
 
 
+class NeMoGuardrailsServer(OpenAICompatible):
+    """Generator for NeMo Guardrails Server
+
+    To select specific rails in a multi rail deployment set `config_ids` to match the rail configuration names
+    as documented by the `NeMo guardrails SDK <https://docs.nvidia.com/nemo/guardrails/0.21.0/run-rails/using-fastapi-server/chat-with-guardrailed-model.html#using-the-openai-python-sdk>`_.
+    """
+
+    ENV_VAR = None
+
+    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=_config):
+        self.api_key = "not-used"  # suppress any api_key from being sent as the server does not utilize one
+        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"

tests/generators/test_generators.py

@@ -134,7 +134,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
+    )
 ]
 
 

tests/generators/test_guardrails.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import httpx
+import pytest
+
+from garak.attempt import Message, Turn, Conversation
+from garak.generators.guardrails import NeMoGuardrailsServer
+
+
+def guardrails_config(selected_rails):
+    """helper method to provide generator configuration"""
+    return {
+        "generators": {
+            "guardrails": {
+                "NeMoGuardrailsServer": {
+                    "name": "UnknownModel",
+                    "config_ids": selected_rails,
+                }
+            }
+        }
+    }
+
+
+@pytest.mark.parametrize(
+    "selected_rails",
+    [
+        [],
+        ["rail1"],
+        ["rail1", "rail2"],
+    ],
+)
+@pytest.mark.respx(base_url=NeMoGuardrailsServer.DEFAULT_PARAMS["uri"])
+def test_guardrail_selection(selected_rails, respx_mock, openai_compat_mocks):
+    """validate selected rails are passed as headers on the request"""
+    mock_response = openai_compat_mocks["chat"]
+    mock_request = respx_mock.post("chat/completions")
+    mock_request.mock(
+        return_value=httpx.Response(
+            mock_response["code"],
+            json=mock_response["json"],
+        )
+    )
+    config_root = guardrails_config(selected_rails)
+    g = NeMoGuardrailsServer(config_root=config_root)
+    conv = Conversation(turns=[Turn(role="user", content=Message("Testing text"))])
+    g.generate(conv)
+    assert mock_request.called
+    for rail in selected_rails:
+        content = str(mock_request.calls.last.request.content)
+        assert "guardrails" in content
+        assert "config_ids" in content
+        assert rail in content