Spaces:
Running
Running
File size: 4,947 Bytes
9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a f4dbf56 9158b0a 553ced5 9158b0a 553ced5 9158b0a 553ced5 9158b0a f4dbf56 9158b0a 553ced5 9158b0a f4dbf56 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
import hashlib
import re
import pytest
from hypothesis import strategies as st, given, settings
from guardrails_genie.guardrails import SecretsDetectionGuardrail
from guardrails_genie.guardrails.secrets_detection import (
SecretsDetectionSimpleResponse,
SecretsDetectionResponse,
REDACTION,
redact_value,
)
@pytest.fixture
def mock_secrets_guard(monkeypatch):
def _mock_guard(*args, **kwargs):
prompt = kwargs.get("prompt")
return_detected_types = kwargs.get("return_detected_secrets")
if "safe text" in prompt:
if return_detected_types:
return SecretsDetectionResponse(
contains_secrets=False,
explanation="No secrets detected in the text.",
detected_secrets={},
redacted_text=prompt,
risk_score=0.0,
)
else:
return SecretsDetectionSimpleResponse(
contains_secrets=False,
explanation="No secrets detected in the text.",
redacted_text=prompt,
risk_score=0.0,
)
else:
if return_detected_types:
return SecretsDetectionResponse(
contains_secrets=True,
explanation="The output contains secrets.",
detected_secrets={"secrets": ["API_KEY"]},
redacted_text="My secret key is [REDACTED:]************[:REDACTED]",
risk_score=1.0,
)
else:
return SecretsDetectionSimpleResponse(
contains_secrets=True,
explanation="The output contains secrets.",
redacted_text="My secret key is [REDACTED:]************[:REDACTED]",
risk_score=1.0,
)
monkeypatch.setattr(
"guardrails_genie.guardrails.secrets_detection.SecretsDetectionGuardrail.guard",
_mock_guard,
)
def test_redact_partial():
text = "ABCDEFGHIJKL"
redacted_text = redact_value(text, REDACTION.REDACT_PARTIAL)
assert redacted_text == "[REDACTED:]AB..KL[:REDACTED]"
def test_redact_all():
text = "ABCDEFGHIJKL"
redacted_text = redact_value(text, REDACTION.REDACT_ALL)
assert redacted_text == "[REDACTED:]************[:REDACTED]"
def test_redact_hash():
text = "ABCDEFGHIJKL"
hashed_value = hashlib.md5(text.encode()).hexdigest()
redacted_text = redact_value(text, REDACTION.REDACT_HASH)
assert redacted_text == f"[REDACTED:]{hashed_value}[:REDACTED]"
def test_secrets_detection_guardrail_detect_types(mock_secrets_guard):
from guardrails_genie.guardrails.secrets_detection import (
SecretsDetectionGuardrail,
REDACTION,
)
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
prompt = "My secret key is ABCDEFGHIJKL"
result = guardrail.guard(prompt=prompt, return_detected_secrets=True)
assert result.contains_secrets is True
assert result.explanation == "The output contains secrets."
assert result.detected_secrets == {"secrets": ["API_KEY"]}
assert result.redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"
def test_secrets_detection_guardrail_simple_response(mock_secrets_guard):
from guardrails_genie.guardrails.secrets_detection import (
SecretsDetectionGuardrail,
REDACTION,
)
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
prompt = "My secret key is ABCDEFGHIJKL"
result = guardrail.guard(prompt=prompt, return_detected_secrets=False)
assert result.contains_secrets is True
assert result.explanation == "The output contains secrets."
assert result.redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"
def test_secrets_detection_guardrail_no_secrets(mock_secrets_guard):
from guardrails_genie.guardrails.secrets_detection import (
SecretsDetectionGuardrail,
REDACTION,
)
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
prompt = "This is a safe text with no secrets."
result = guardrail.guard(prompt=prompt, return_detected_secrets=True)
assert result.contains_secrets is False
assert result.explanation == "No secrets detected in the text."
assert result.detected_secrets == {}
assert result.redacted_text == prompt
def pattern_strategy(pattern):
return st.from_regex(re.compile(pattern), fullmatch=True)
@settings(deadline=1000)
@given(pattern_strategy(r"AKIA[0-9A-Z]{16}"))
def test_specific_pattern_guardrail(text):
guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
result = guardrail.guard(prompt=text, return_detected_secrets=True)
assert result.contains_secrets is True
assert "AWS Access Key" in result.detected_secrets
|