feat(jailbreak): Add direct API key configuration support (#1260)

* Support direct jailbreak api key, not via environment variable

* Add unit-test to cover api_key_env_var being set, but no environment variable exists with the value

* Removed unused imports, fixed test docstring copy-and-paste

* Rename get_auth_token() to get_api_key()

Author: tgasser-nv

examples/configs/jailbreak_detection_nim/README.md

@@ -0,0 +1,7 @@
+# Jailbreak Detection using NIMs
+
+This examples showcases the jailbreak detection capabilities of NeMo Guardrails using a NIM hosted on NVCF.
+
+The structure of the config folder is the following:
+
+- `config.yml` - The config file holding all the configuration options.

examples/configs/jailbreak_detection_nim/config.yml

@@ -0,0 +1,41 @@
+models:
+  - type: main
+    engine: nvidia_ai_endpoints
+    model: mistralai/mixtral-8x7b-instruct-v0.1
+    parameters:
+      temperature: 0.7
+      max_tokens: 1000
+      timeout: 120
+      api_key: "<insert NVCF API key here>"
+
+rails:
+  config:
+    jailbreak_detection:
+      nim_base_url: "https://ai.api.nvidia.com"
+      nim_server_endpoint: "/v1/security/nvidia/nemoguard-jailbreak-detect"
+      api_key:  "<insert NVCF API key here>"
+  input:
+    flows:
+      - jailbreak detection model
+  output:
+    flows: []
+  retrieval:
+    flows: []
+
+instructions:
+  - type: general
+    content: |
+      Below is a conversation between a helpful AI assistant and a user.
+      The assistant is direct, honest, and concise.
+      If the assistant does not know something, it says so.
+      The assistant does not engage in harmful, unethical, or illegal behavior.
+
+sample_conversation: |
+  user "Hello there!"
+    express greeting
+  bot express greeting
+    "Hello! How can I assist you today?"
+  user "What can you do for me?"
+    ask about capabilities
+  bot respond about capabilities
+    "As an AI assistant, I can help you with a wide range of tasks. This includes question answering on various topics, generating text for various purposes and providing suggestions based on your preferences."

nemoguardrails/library/jailbreak_detection/actions.py

@@ -97,15 +97,7 @@ async def jailbreak_detection_model(
     jailbreak_api_url = jailbreak_config.server_endpoint
     nim_base_url = jailbreak_config.nim_base_url
     nim_classification_path = jailbreak_config.nim_server_endpoint
-    if jailbreak_config.api_key_env_var is not None:
-        nim_auth_token = os.getenv(jailbreak_config.api_key_env_var)
-        if nim_auth_token is None:
-            log.warning(
-                "Specified a value for jailbreak config api_key_env var at %s but the environment variable was not set!"
-                % jailbreak_config.api_key_env_var
-            )
-    else:
-        nim_auth_token = None
+    nim_auth_token = jailbreak_config.get_api_key()
 
     if context is not None:
         prompt = context.get("user_message", "")

nemoguardrails/rails/llm/config.py

@@ -26,6 +26,7 @@
     BaseModel,
     ConfigDict,
     Field,
+    SecretStr,
     model_validator,
     root_validator,
     validator,
@@ -572,6 +573,10 @@ class JailbreakDetectionConfig(BaseModel):
         default="classify",
         description="Classification path uri. Defaults to 'classify' for NemoGuard JailbreakDetect.",
     )
+    api_key: Optional[SecretStr] = Field(
+        default=None,
+        description="Secret String with API key for use in Jailbreak requests. Takes precedence over api_key_env_var",
+    )
     api_key_env_var: Optional[str] = Field(
         default=None,
         description="Environment variable containing API key for jailbreak detection model",
@@ -600,6 +605,31 @@ def migrate_deprecated_fields(self) -> "JailbreakDetectionConfig":
             self.nim_base_url = f"http://{self.nim_url}:{port}/v1"
         return self
 
+    def get_api_key(self) -> Optional[str]:
+        """Helper to return an API key (if it exists) from a Jailbreak configuration.
+          This can come from (in descending order of priority):
+
+        1. The `api_key` field, a Pydantic SecretStr from which we extract the full string.
+        2. The `api_key_env_var` field, a string stored in this environment variable.
+
+        If neither is found, None is returned.
+        """
+
+        if self.api_key:
+            return self.api_key.get_secret_value()
+
+        if self.api_key_env_var:
+            nim_auth_token = os.getenv(self.api_key_env_var)
+            if nim_auth_token:
+                return nim_auth_token
+
+            log.warning(
+                "Specified a value for jailbreak config api_key_env var at %s but the environment variable was not set!"
+                % self.api_key_env_var
+            )
+
+        return None
+
 
 class AutoAlignOptions(BaseModel):
     """List of guardrails that are activated"""

tests/test_jailbreak_config.py

@@ -12,7 +12,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
+from unittest.mock import patch
 
+from pydantic import SecretStr
 
 from nemoguardrails.rails.llm.config import JailbreakDetectionConfig
 
@@ -128,3 +131,60 @@ def test_empty_configuration(self):
         assert config.nim_url is None
         assert config.nim_port is None
         assert config.embedding is None
+
+    def test_get_api_key_no_key(self):
+        """Check when neither `api_key` nor `api_key_env_var` are provided, auth token is None"""
+
+        config = JailbreakDetectionConfig(
+            nim_base_url="http://localhost:8000/v1",
+            nim_server_endpoint="classify",
+        )
+
+        auth_token = config.get_api_key()
+        assert auth_token is None
+
+    def test_get_api_key_api_key(self):
+        """Check when both `api_key` and `api_key_env_var` are provided, `api_key` takes precedence"""
+        api_key_value = "nvapi-abcdef12345"
+        api_key_env_var_name = "CUSTOM_API_KEY"
+        api_key_env_var_value = "env-var-nvapi-abcdef12345"
+
+        with patch.dict(os.environ, {api_key_env_var_name: api_key_env_var_value}):
+            config = JailbreakDetectionConfig(
+                nim_base_url="http://localhost:8000/v1",
+                nim_server_endpoint="classify",
+                api_key=api_key_value,
+                api_key_env_var=api_key_env_var_name,
+            )
+
+            auth_token = config.get_api_key()
+            assert auth_token == api_key_value
+
+    def test_get_api_key_api_key_env_var(self):
+        """Check when only `api_key_env_var` is provided, the env-var value is correctly returned"""
+        api_key_env_var_name = "CUSTOM_API_KEY"
+        api_key_env_var_value = "env-var-nvapi-abcdef12345"
+
+        with patch.dict(os.environ, {api_key_env_var_name: api_key_env_var_value}):
+            config = JailbreakDetectionConfig(
+                nim_base_url="http://localhost:8000/v1",
+                nim_server_endpoint="classify",
+                api_key_env_var=api_key_env_var_name,
+            )
+
+            auth_token = config.get_api_key()
+            assert auth_token == api_key_env_var_value
+
+    def test_get_api_key_api_key_env_var_not_set(self):
+        """Check configuring an `api_key_env_var` that isn't set in the shell returns None"""
+        api_key_env_var_name = "CUSTOM_API_KEY"
+
+        with patch.dict(os.environ, {}):
+            config = JailbreakDetectionConfig(
+                nim_base_url="http://localhost:8000/v1",
+                nim_server_endpoint="classify",
+                api_key_env_var=api_key_env_var_name,
+            )
+
+            auth_token = config.get_api_key()
+            assert auth_token is None