Guardrails Policy

import os
from enkryptai_sdk import *
from dotenv import load_dotenv

load_dotenv()

ENKRYPT_API_KEY = os.getenv("ENKRYPTAI_API_KEY")
ENKRYPT_BASE_URL = os.getenv("ENKRYPTAI_BASE_URL") or "https://api.enkryptai.com"

guardrails_client = GuardrailsClient(api_key=ENKRYPT_API_KEY, base_url=ENKRYPT_BASE_URL)

test_guardrails_policy_name = "Test Guardrails Policy"

sample_detectors = {
        "pii": {
            "enabled": False,
            "entities": [
                "pii",
                "secrets",
                "ip_address",
                "url"
            ]
        },
        "nsfw": {
            "enabled": True,
            "block_message": "Your custom message"
        },
        "toxicity": {
            "enabled": False,
            "block_message": "Your custom message"
        },
        "topic_detector": {
            "topic": ["science"],
            "enabled": False,
            "block_message": "Your custom message"
        },
        "injection_attack": {
            "enabled": True,
            "block_message": "Your custom message"
        },
        "keyword_detector": {
            "enabled": False,
            "banned_keywords": []
        },
        "policy_violation": {
            "enabled": True,
            "need_explanation": True,
            "policy_text": "Do not allow any illegal or immoral activities.",
            # Or we can also give coc_policy_name of a saved Code of Conduct Policy
            # Instead of policy_text
            # "coc_policy_name": "Test CoC Policy""
            "block_message": "Your custom message"
        },
        "bias": {
            "enabled": False
        },
        "sponge_attack": {
            "enabled": False,
            "block_message": "Your custom message"
        }
    }

# Create a policy with a dictionary
add_policy_response = guardrails_client.add_policy(
    policy_name=test_guardrails_policy_name,
    config=copy.deepcopy(sample_detectors),
    description="Sample custom security policy"
)

# Or create a policy with GuardrailsConfig object
injection_config = GuardrailsConfig.injection_attack()
add_policy_response = guardrails_client.add_policy(
    policy_name=test_guardrails_policy_name,
    config=injection_config,
    description="Detects prompt injection attacks"
)

print(add_policy_response)

assert add_policy_response.message == "Policy details added successfully"

# Print as a dictionary
print(add_policy_response.to_dict())

# Retrieve policy configuration
policy = guardrails_client.get_policy(policy_name=test_guardrails_policy_name)

print(policy)

# Get other fields
print(policy.name)
print(policy.detectors)

# Print as a dictionary
print(policy.to_dict())
print(policy.detectors.to_dict())

# Update policy with new configuration
# Similar to add, we can use a dictionary or GuardrailsConfig object
new_detectors_dict = copy.deepcopy(sample_detectors)
# Modify the detectors as needed
# Example: Enable bias detection
new_detectors_dict["bias"]["enabled"] = True

new_config = new_detectors_dict or GuardrailsConfig.bias()  # Switch to bias detection

modify_policy_response = guardrails_client.modify_policy(
    policy_name=test_guardrails_policy_name,
    guardrails_config=new_config,
    description="Updated to detect bias"
)

print(modify_policy_response)

assert modify_policy_response.message == "Policy details updated successfully"

# Print as a dictionary
print(modify_policy_response.to_dict())

# Remove a policy
delete_policy_response = guardrails_client.delete_policy(policy_name=test_guardrails_policy_name)

print(delete_policy_response)

assert delete_policy_response.message == "Policy details deleted successfully"

# Print as a dictionary
print(delete_policy_response.to_dict())

# List all policies
policies = guardrails_client.get_policy_list()

print(policies)

# Get the first policy
print(policies.policies[0])
print(policies.policies[0].name)

# Print as a dictionary
print(policies.to_dict())

# Use policy to detect
policy_detect_response = guardrails_client.policy_detect(
    policy_name=test_guardrails_policy_name,
    text="Check this text for policy violations"
)

print(policy_detect_response)

# Print as a dictionary
print(policy_detect_response.to_dict())