Nvidia NIM (#264)

lib/rules/input.go

@@ -16,27 +16,29 @@ import (
 )
 
 type InputTypes struct {
-	LanguageDetection string
-	PromptInjection   string
-	PIIFilter         string
-	InvisibleChars    string
-	Moderation        string
-	LlamaGuard        string
-	PromptGuard       string
-	LangKit           string
-	VigilLLM          string
+	LanguageDetection  string
+	PromptInjection    string
+	PIIFilter          string
+	InvisibleChars     string
+	Moderation         string
+	LlamaGuard         string
+	PromptGuard        string
+	LangKit            string
+	VigilLLM           string
+	NvidiaNimJailbreak string
 }
 
 var inputTypes = InputTypes{
-	LanguageDetection: "language_detection",
-	PromptInjection:   "prompt_injection",
-	PIIFilter:         "pii_filter",
-	InvisibleChars:    "invisible_chars",
-	Moderation:        "moderation",
-	LlamaGuard:        "llama_guard",
-	PromptGuard:       "prompt_guard",
-	LangKit:           "langkit",
-	VigilLLM:          "vigilllm",
+	LanguageDetection:  "language_detection",
+	PromptInjection:    "prompt_injection",
+	PIIFilter:          "pii_filter",
+	InvisibleChars:     "invisible_chars",
+	Moderation:         "moderation",
+	LlamaGuard:         "llama_guard",
+	PromptGuard:        "prompt_guard",
+	LangKit:            "langkit",
+	VigilLLM:           "vigilllm",
+	NvidiaNimJailbreak: "nvidia_nim_jailbreak",
 }
 
 type Rule struct {
@@ -317,7 +319,10 @@ func handleRuleAction(inputConfig lib.Rule, rule RuleResult, ruleType string, me
 	log.Printf("%s detection result: Match=%v, Score=%f", ruleType, rule.Match, rule.Inspection.Score)
 
 	switch ruleType {
+	case inputTypes.NvidiaNimJailbreak:
+		return genericHandler(inputConfig, rule)
 	case inputTypes.LangKit:
+		return genericHandler(inputConfig, rule)
 	case inputTypes.VigilLLM:
 		return genericHandler(inputConfig, rule)
 	case inputTypes.InvisibleChars:

services/rule/src/plugins/nvidia_nim_jailbreak.py

@@ -0,0 +1,67 @@
+"""
+NVIDIA NIM Jailbreak Detection Plugin
+
+This plugin interfaces with NVIDIA's NIM Jailbreak Detection API
+
+Environment Variables:
+- NVIDIA_API_KEY: Required API key for authentication
+
+"""
+
+import os
+import requests
+from typing import Dict, Any
+from utils.logger_config import setup_logger
+
+logger = setup_logger(__name__)
+
+
+def handler(text: str, threshold: float, config: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Processes text through NVIDIA NIM Jailbreak Detection API
+
+    Args:
+        text: User input text to analyze
+        threshold: e threshold (unused in this plugin)
+        config:  configuration
+
+    Returns:
+        Dict with 'score' (float) and 'check_result' (bool) detection results
+
+    Raises:
+        ValueError: For missing API key or invalid API responses
+        RuntimeError: For API request failures
+    """
+    # Get API key from environment
+    api_key = os.getenv("NVIDIA_API_KEY")
+    if not api_key:
+        logger.error("Missing NVIDIA_API_KEY environment variable")
+        raise ValueError("NVIDIA_API_KEY environment variable required")
+
+
+    url = "https://ai.api.nvidia.com/v1/security/nvidia/nemoguard-jailbreak-detect"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}"
+    }
+    payload = {"input": text}
+
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+    except requests.exceptions.RequestException as e:
+        logger.error(f"NVIDIA API request failed: {str(e)}")
+        raise RuntimeError(f"NVIDIA API request failed: {str(e)}")
+
+    try:
+        response_data = response.json()
+        score = float(response_data["score"])
+        jailbreak = response_data["jailbreak"]
+    except (KeyError, ValueError) as e:
+        logger.error(f"Invalid API response format: {str(e)}")
+        raise ValueError(f"Invalid API response format: {str(e)}")
+
+    return {
+        "check_result": jailbreak,
+        "score": score,
+    }