feat(jailbreak): Validate Jailbreak Detection config at create-time (#1675)

Author: tgasser-nv

nemoguardrails/rails/llm/config.py

@@ -53,6 +53,9 @@
 with open(os.path.join(os.path.dirname(__file__), "default_config_v2.yml")) as _fc:
     _default_config_v2 = yaml.safe_load(_fc)
 
+# Jailbreak-related strings
+JAILBREAK_FLOW_MODEL = "jailbreak detection model"
+JAILBREAK_FLOW_HEURISTICS = "jailbreak detection heuristics"
 
 # Extract the COLANGPATH directories.
 colang_path_dirs = [
@@ -700,9 +703,9 @@ class JailbreakDetectionConfig(BaseModel):
         default=None,
         description="The endpoint for the jailbreak detection heuristics/model container.",
     )
-    length_per_perplexity_threshold: float = Field(default=89.79, description="The length/perplexity threshold.")
+    length_per_perplexity_threshold: float = Field(default=89.79, gt=0, description="The length/perplexity threshold.")
     prefix_suffix_perplexity_threshold: float = Field(
-        default=1845.65, description="The prefix/suffix perplexity threshold."
+        default=1845.65, gt=0, description="The prefix/suffix perplexity threshold."
     )
     nim_base_url: Optional[str] = Field(
         default=None,
@@ -744,6 +747,15 @@ def migrate_deprecated_fields(self) -> "JailbreakDetectionConfig":
             self.nim_base_url = f"http://{self.nim_url}:{port}/v1"
         return self
 
+    @model_validator(mode="after")
+    def validate_urls(self) -> "JailbreakDetectionConfig":
+        """Validate URL formats for endpoints."""
+        if self.nim_base_url and not self.nim_base_url.startswith(("http://", "https://")):
+            raise ValueError(f"nim_base_url must start with 'http://' or 'https://', got '{self.nim_base_url}'")
+        if self.server_endpoint and not self.server_endpoint.startswith(("http://", "https://")):
+            raise ValueError(f"server_endpoint must start with 'http://' or 'https://', got '{self.server_endpoint}'")
+        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):
@@ -1712,6 +1724,60 @@ def check_output_parser_exists(cls, values):
                 )
         return values
 
+    @root_validator(pre=True, allow_reuse=True)
+    def check_jailbreak_detection_config(cls, values):
+        """Validate jailbreak detection configuration against enabled flows."""
+        rails = values.get("rails") or {}
+        config_data = rails.get("config") or {}
+        input_flows = (rails.get("input") or {}).get("flows") or []
+
+        jailbreak_config = config_data.get("jailbreak_detection") or {}
+        has_model_flow = JAILBREAK_FLOW_MODEL in input_flows
+        has_heuristics_flow = JAILBREAK_FLOW_HEURISTICS in input_flows
+        has_any_jailbreak_flow = has_model_flow or has_heuristics_flow
+
+        # Case A: Config present but no flow references it
+        if jailbreak_config and not has_any_jailbreak_flow:
+            log.warning(
+                "Jailbreak detection configuration is present under "
+                "rails.config.jailbreak_detection but no jailbreak detection flow "
+                "is enabled. To use jailbreak detection, add 'jailbreak detection model' "
+                "or 'jailbreak detection heuristics' to rails.input.flows."
+            )
+
+        # Case B: "jailbreak detection model" flow is enabled
+        if has_model_flow:
+            nim_base_url = jailbreak_config.get("nim_base_url")
+            nim_url = jailbreak_config.get("nim_url")  # deprecated, migrated later
+            server_endpoint = jailbreak_config.get("server_endpoint")
+            nim_server_endpoint = jailbreak_config.get("nim_server_endpoint", "classify")
+
+            if nim_base_url or nim_url:
+                if not nim_server_endpoint:
+                    raise InvalidRailsConfigurationError(
+                        "nim_base_url is set for jailbreak detection model but "
+                        "nim_server_endpoint is empty. Both must be configured "
+                        "when using NIM-based jailbreak detection."
+                    )
+            elif not server_endpoint:
+                log.warning(
+                    "No endpoint configured for jailbreak detection model. "
+                    "Will fall back to local in-process detection, which is "
+                    "not recommended for production."
+                )
+
+        # Case C: "jailbreak detection heuristics" flow is enabled
+        if has_heuristics_flow:
+            server_endpoint = jailbreak_config.get("server_endpoint")
+            if not server_endpoint:
+                log.warning(
+                    "No server_endpoint configured for jailbreak detection heuristics. "
+                    "Will fall back to local in-process detection, which is "
+                    "not recommended for production."
+                )
+
+        return values
+
     @root_validator(pre=True, allow_reuse=True)
     def fill_in_default_values_for_v2_x(cls, values):
         instructions = values.get("instructions", {})

tests/test_jailbreak_config.py

@@ -12,10 +12,13 @@
 # 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 logging
 import os
 from unittest.mock import patch
 
-from nemoguardrails.rails.llm.config import JailbreakDetectionConfig
+import pytest
+
+from nemoguardrails.rails.llm.config import JailbreakDetectionConfig, RailsConfig
 
 
 class TestJailbreakDetectionConfig:
@@ -184,3 +187,247 @@ def test_get_api_key_api_key_env_var_not_set(self):
 
             auth_token = config.get_api_key()
             assert auth_token is None
+
+    def test_negative_length_per_perplexity_threshold_raises(self):
+        """Threshold <= 0 should raise ValueError."""
+        with pytest.raises(ValueError, match="greater than 0"):
+            JailbreakDetectionConfig(length_per_perplexity_threshold=-1.0)
+
+    def test_negative_prefix_suffix_perplexity_threshold_raises(self):
+        """Threshold <= 0 should raise ValueError."""
+        with pytest.raises(ValueError, match="greater than 0"):
+            JailbreakDetectionConfig(prefix_suffix_perplexity_threshold=0)
+
+    def test_invalid_nim_base_url_raises(self):
+        """nim_base_url without http(s) scheme should raise ValueError."""
+        with pytest.raises(ValueError, match="nim_base_url must start with"):
+            JailbreakDetectionConfig(nim_base_url="ftp://localhost:8000/v1")
+
+    def test_invalid_server_endpoint_raises(self):
+        """server_endpoint without http(s) scheme should raise ValueError."""
+        with pytest.raises(ValueError, match="server_endpoint must start with"):
+            JailbreakDetectionConfig(server_endpoint="localhost:1337/model")
+
+    def test_valid_urls_accepted(self):
+        """Valid http and https URLs should be accepted."""
+        config = JailbreakDetectionConfig(
+            nim_base_url="https://nim.example.com/v1",
+            server_endpoint="http://localhost:1337/model",
+        )
+        assert config.nim_base_url == "https://nim.example.com/v1"
+        assert config.server_endpoint == "http://localhost:1337/model"
+
+
+def _make_rails_config(**kwargs):
+    """Helper to build a RailsConfig with minimal required fields."""
+    defaults = {
+        "models": [{"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}],
+    }
+    defaults.update(kwargs)
+    return RailsConfig(**defaults)
+
+
+class TestJailbreakDetectionCrossValidation:
+    def test_model_flow_with_nim_url_but_no_endpoint_raises(self):
+        """nim_base_url set but nim_server_endpoint empty should raise."""
+        with pytest.raises(Exception, match="nim_server_endpoint is empty"):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_base_url": "http://localhost:8000/v1",
+                            "nim_server_endpoint": "",
+                        }
+                    },
+                },
+            )
+
+    def test_model_flow_with_no_endpoints_warns(self, caplog):
+        """No nim_base_url or server_endpoint should warn about local fallback."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {"jailbreak_detection": {}},
+                },
+            )
+        assert "No endpoint configured for jailbreak detection model" in caplog.text
+
+    def test_heuristics_flow_with_no_server_endpoint_warns(self, caplog):
+        """No server_endpoint for heuristics flow should warn."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection heuristics"]},
+                    "config": {"jailbreak_detection": {}},
+                },
+            )
+        assert "No server_endpoint configured for jailbreak detection heuristics" in caplog.text
+
+    def test_jailbreak_config_present_but_no_flow_warns(self, caplog):
+        """Orphaned jailbreak_detection config should warn."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_base_url": "http://localhost:8000/v1",
+                        }
+                    },
+                },
+            )
+        assert "no jailbreak detection flow is enabled" in caplog.text
+
+    def test_model_flow_with_nim_fully_configured_passes(self, caplog):
+        """Fully configured NIM-based model flow should produce no warnings."""
+        with caplog.at_level(logging.WARNING):
+            config = _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_base_url": "http://localhost:8000/v1",
+                            "nim_server_endpoint": "classify",
+                        }
+                    },
+                },
+            )
+        assert "jailbreak" not in caplog.text.lower()
+        assert config.rails.config.jailbreak_detection.nim_base_url == "http://localhost:8000/v1"
+
+    def test_model_flow_with_deprecated_nim_url_no_spurious_warning(self, caplog):
+        """Deprecated nim_url/nim_port should not trigger 'no endpoint' warning."""
+        with caplog.at_level(logging.WARNING):
+            config = _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_url": "localhost",
+                            "nim_port": 8000,
+                        }
+                    },
+                },
+            )
+        assert "No endpoint configured" not in caplog.text
+        # Verify migration happened
+        assert config.rails.config.jailbreak_detection.nim_base_url == "http://localhost:8000/v1"
+
+    def test_model_flow_with_server_endpoint_passes(self, caplog):
+        """Model flow with server_endpoint (no NIM) should pass without warning."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "server_endpoint": "http://localhost:1337/model",
+                        }
+                    },
+                },
+            )
+        assert "jailbreak" not in caplog.text.lower()
+
+    def test_heuristics_flow_with_server_endpoint_passes(self, caplog):
+        """Heuristics flow with server_endpoint should pass without warning."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection heuristics"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "server_endpoint": "http://localhost:1337/heuristics",
+                        }
+                    },
+                },
+            )
+        assert "jailbreak" not in caplog.text.lower()
+
+    def test_model_flow_deprecated_nim_url_empty_server_endpoint_raises(self):
+        """Deprecated nim_url with empty nim_server_endpoint should raise."""
+        with pytest.raises(Exception, match="nim_server_endpoint is empty"):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_url": "localhost",
+                            "nim_server_endpoint": "",
+                        }
+                    },
+                },
+            )
+
+    def test_model_flow_nim_port_only_warns(self, caplog):
+        """nim_port alone (no nim_url or nim_base_url) should warn about local fallback."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_port": 9000,
+                        }
+                    },
+                },
+            )
+        assert "No endpoint configured for jailbreak detection model" in caplog.text
+
+    def test_both_flows_nim_only_warns_heuristics(self, caplog):
+        """Both flows with only NIM configured should warn for heuristics only."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {
+                        "flows": [
+                            "jailbreak detection model",
+                            "jailbreak detection heuristics",
+                        ]
+                    },
+                    "config": {
+                        "jailbreak_detection": {
+                            "nim_base_url": "http://localhost:8000/v1",
+                        }
+                    },
+                },
+            )
+        assert "No endpoint configured for jailbreak detection model" not in caplog.text
+        assert "No server_endpoint configured for jailbreak detection heuristics" in caplog.text
+
+    def test_both_flows_server_endpoint_only_passes(self, caplog):
+        """Both flows with server_endpoint should pass without warnings."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {
+                        "flows": [
+                            "jailbreak detection model",
+                            "jailbreak detection heuristics",
+                        ]
+                    },
+                    "config": {
+                        "jailbreak_detection": {
+                            "server_endpoint": "http://localhost:1337/model",
+                        }
+                    },
+                },
+            )
+        assert "jailbreak" not in caplog.text.lower()
+
+    def test_explicit_null_jailbreak_detection_config(self, caplog):
+        """Explicit None for jailbreak_detection should not raise AttributeError."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config(
+                rails={
+                    "input": {"flows": ["jailbreak detection model"]},
+                    "config": {"jailbreak_detection": None},
+                },
+            )
+        assert "No endpoint configured for jailbreak detection model" in caplog.text
+
+    def test_no_jailbreak_config_no_flow_no_warnings(self, caplog):
+        """Default config with no jailbreak config or flows should produce no warnings."""
+        with caplog.at_level(logging.WARNING):
+            _make_rails_config()
+        assert "jailbreak" not in caplog.text.lower()