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
},
"toxicity": {
"enabled": False
},
"topic_detector": {
"topic": ["science"],
"enabled": False
},
"injection_attack": {
"enabled": True
},
"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""
},
"bias": {
"enabled": False
}
}
# 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())