Custom Guardrails Example

colang_version: "2.x"

models:
  - type: main
    engine: openai
    model: gpt-3.5-turbo-instruct

import core
import guardrails
import llm

flow main
  activate llm continuation
  activate greeting

flow greeting
  user expressed greeting
  bot express greeting

flow user expressed greeting
  user said "hi" or user said "hello"

flow bot express greeting
  bot say "Hello world!"

flow input rails $input_text
  $input_safe = await check user utterance $input_text

  if not $input_safe
    bot say "I'm sorry, I can't respond to that."
    abort

flow check user utterance $input_text -> $input_safe
  $is_safe = ..."Consider the following user utterance: '{$input_text}'. Assign 'True' if appropriate, 'False' if inappropriate."
  print $is_safe
  return $is_safe

colang_version: 2.x

models:
  - type: main
    engine: nim
    model: meta/llama-3.3-70b-instruct
  
  - type: content_safety
    engine: nim
    model: nvidia/llama-3.1-nemoguard-8b-content-safety
  
  - type: topic_control
    engine: nim
    model: nvidia/llama-3.1-nemoguard-8b-topic-control

rails:
  config:
    jailbreak_detection:
      nim_base_url: "https://ai.api.nvidia.com"
      nim_server_endpoint: "/v1/security/nvidia/nemoguard-jailbreak-detect"
      api_key_env_var: NVIDIA_API_KEY

import guardrails
import nemoguardrails.library.content_safety
import nemoguardrails.library.topic_safety
import nemoguardrails.library.jailbreak_detection

flow input rails $input_text
    content safety check input $model="content_safety"
    topic safety check input $model="topic_control"
    jailbreak detection model

flow output rails $output_text
    content safety check output $model="content_safety"

prompts:
  - task: content_safety_check
    content: |-
      Check if the following content violates safety policies.
      
      Categories to check:
      - Violence and hate speech
      - Sexual content
      - Self-harm
      - Criminal activity
      - Privacy violations
      
      Content: {{ content }}
      
      Safe: yes/no

  - task: topic_control_check  
    content: |-
      Determine if the query is within allowed topics.
      
      Allowed topics:
      - Company policies
      - Employee benefits
      - HR procedures
      - General workplace questions
      
      Query: {{ query }}
      
      On topic: yes/no

from nemoguardrails import LLMRails
from nemoguardrails.actions.actions import ActionResult
import re

async def check_pii(context: dict) -> ActionResult:
    """Check if user input contains PII (emails, phone numbers, SSN)."""
    user_message = context.get("last_user_message")
    
    # Check for email addresses
    email_pattern = r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'
    if re.search(email_pattern, user_message):
        return ActionResult(
            return_value=False,
            context_updates={"pii_detected": "email"}
        )
    
    # Check for phone numbers
    phone_pattern = r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b'
    if re.search(phone_pattern, user_message):
        return ActionResult(
            return_value=False,
            context_updates={"pii_detected": "phone"}
        )
    
    # Check for SSN
    ssn_pattern = r'\b\d{3}-\d{2}-\d{4}\b'
    if re.search(ssn_pattern, user_message):
        return ActionResult(
            return_value=False,
            context_updates={"pii_detected": "ssn"}
        )
    
    return ActionResult(return_value=True)

async def mask_sensitive_data(context: dict) -> ActionResult:
    """Mask sensitive data in bot responses."""
    bot_message = context.get("bot_message")
    
    # Mask credit card numbers
    cc_pattern = r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b'
    masked_message = re.sub(cc_pattern, '****-****-****-****', bot_message)
    
    # Mask API keys
    api_pattern = r'\b[A-Za-z0-9]{32,}\b'
    masked_message = re.sub(api_pattern, '[REDACTED]', masked_message)
    
    return ActionResult(return_value=masked_message)

def init(app: LLMRails):
    app.register_action(check_pii, "check_pii")
    app.register_action(mask_sensitive_data, "mask_sensitive_data")

define flow pii protection
  user ...
  $is_safe = execute check_pii()
  
  if not $is_safe
    bot inform pii detected
    stop

define bot inform pii detected
  "I noticed you may have shared personal information. For your security, please don't share emails, phone numbers, or other sensitive data."

define flow output masking
  bot ...
  $masked_response = execute mask_sensitive_data()
  bot $masked_response