PII Detection Guardrail for Teams

from agno.team import Team
from agno.guardrails import PIIDetectionGuardrail
from agno.models.openai import OpenAIResponses
from agno.run import RunStatus


def check_input(team: Team, label: str, text: str, expect_blocked: bool) -> None:
    response = team.run(input=text)
    blocked = response.status == RunStatus.error
    print(f"{label}: {'blocked' if blocked else 'allowed'}")
    if response.content:
        print(response.content)
    assert blocked == expect_blocked


def main():
    team = Team(
        name="Privacy-Protected Team",
        members=[],
        model=OpenAIResponses(id="gpt-5.2"),
        pre_hooks=[PIIDetectionGuardrail()],
        description="A team that helps with customer service while protecting privacy.",
        instructions="You are a helpful customer service assistant. Always protect user privacy and handle sensitive information appropriately.",
    )

    cases = [
        ("Normal request", "Can you help me understand your return policy?", False),
        ("SSN", "My Social Security Number is 123-45-6789.", True),
        ("Credit card", "My card number is 4532 1234 5678 9012.", True),
        ("Email", "Send the receipt to john.doe@example.com.", True),
        ("Phone", "My phone number is 555-123-4567.", True),
        (
            "Multiple PII types",
            "My email is john@company.com and phone is 555.987.6543.",
            True,
        ),
        ("Unseparated credit card", "My card is 4532123456789012.", True),
    ]
    for label, text, expect_blocked in cases:
        check_input(team, label, text, expect_blocked)

    masked_team = Team(
        name="Privacy-Protected Team (Masked)",
        members=[],
        model=OpenAIResponses(id="gpt-5.2"),
        pre_hooks=[PIIDetectionGuardrail(mask_pii=True)],
        description="A team that helps with customer service while protecting privacy.",
        instructions="You are a helpful customer service assistant. Always protect user privacy and handle sensitive information appropriately.",
    )
    masked_team.print_response(
        input="Hi, my Social Security Number is 123-45-6789. Can you help me with my account?",
    )


if __name__ == "__main__":
    main()