diff --git a/aidefense/runtime/__init__.py b/aidefense/runtime/__init__.py index 2c3dd49..2a9f779 100644 --- a/aidefense/runtime/__init__.py +++ b/aidefense/runtime/__init__.py @@ -25,6 +25,7 @@ InspectionConfig, Metadata, InspectResponse, + DetectedPII, ) from .http_models import HttpInspectRequest from .http_models import ( diff --git a/aidefense/runtime/inspection_client.py b/aidefense/runtime/inspection_client.py index 0791f63..14473ba 100644 --- a/aidefense/runtime/inspection_client.py +++ b/aidefense/runtime/inspection_client.py @@ -29,6 +29,7 @@ Severity, Classification, InspectResponse, + DetectedPII, ) from .constants import INTEGRATION_DETAILS from ..request_handler import RequestHandler @@ -148,8 +149,16 @@ def _parse_inspect_response(self, response_data: Dict[str, Any]) -> "InspectResp "attack_technique": "NONE_ATTACK_TECHNIQUE", "explanation": "", "client_transaction_id": "", - "event_id": "b403de99-8d19-408f-8184-ec6d7907f508" - "action": "Allow" + "event_id": "b403de99-8d19-408f-8184-ec6d7907f508", + "action": "Allow", + "detected_pii": [ + { + "message_index": "0", + "type": "PII_ENTITY_TYPE_EMAIL", + "start_index": "22", + "end_index": "48" + } + ] } ``` @@ -170,8 +179,16 @@ def _parse_inspect_response(self, response_data: Dict[str, Any]) -> "InspectResp attack_technique="NONE_ATTACK_TECHNIQUE", explanation="", client_transaction_id="", - event_id="b403de99-8d19-408f-8184-ec6d7907f508" - action="Allow" + event_id="b403de99-8d19-408f-8184-ec6d7907f508", + action="Allow", + detected_pii=[ + DetectedPII( + message_index="0", + type="PII_ENTITY_TYPE_EMAIL", + start_index="22", + end_index="48" + ) + ] ) ``` """ @@ -209,6 +226,19 @@ def _parse_rule_list(rule_list: list) -> List[Rule]: ) return out + def _parse_detected_pii_list(detected_pii_list: list) -> List[DetectedPII]: + out = [] + for detected_pii in detected_pii_list: + out.append( + DetectedPII( + message_index=detected_pii.get("message_index"), + type=detected_pii.get("type"), + start_index=detected_pii.get("start_index"), + end_index=detected_pii.get("end_index"), + ) + ) + return out + # Parse rules if present rules = _parse_rule_list(response_data.get("rules", [])) # Parse processed_rules if present (API may send processed_rules or processedRules) @@ -240,6 +270,7 @@ def _parse_rule_list(rule_list: list) -> List[Rule]: client_transaction_id=response_data.get("client_transaction_id"), event_id=response_data.get("event_id"), action=action, + detected_pii=_parse_detected_pii_list(response_data.get("detected_pii", [])) or None ) def _prepare_inspection_metadata(self, metadata: Metadata) -> Dict: diff --git a/aidefense/runtime/models.py b/aidefense/runtime/models.py index 70b005c..6fd0991 100644 --- a/aidefense/runtime/models.py +++ b/aidefense/runtime/models.py @@ -85,6 +85,10 @@ class RuleName(str, Enum): SEXUAL_CONTENT_EXPLOITATION = "Sexual Content & Exploitation" SOCIAL_DIVISION_POLARIZATION = "Social Division & Polarization" VIOLENCE_PUBLIC_SAFETY_THREATS = "Violence & Public Safety Threats" + TOXICITY = "Toxicity" + GENERAL_HARMS = "General Harms" + TOOL_EXPLOITATION = "Tool Exploitation" + MALICIOUS_URL_DETECTION = "Malicious URL Detection" @dataclass @@ -106,6 +110,24 @@ class Rule: classification: Optional[Classification] = None +@dataclass +class DetectedPII: + """ + Represents a detected PII entity in the inspected content. + + Attributes: + message_index (Optional[str]): Index of the message in the conversation + type (Optional[str]): Type of the PII entity (e.g., Email Address, Phone Number) + start_index (Optional[str]): Start character index of the PII in the message + end_index (Optional[str]): End character index of the PII in the message + """ + + message_index: Optional[str] = None + type: Optional[str] = None + start_index: Optional[str] = None + end_index: Optional[str] = None + + @dataclass class Metadata: """ @@ -165,6 +187,7 @@ class InspectResponse: Attributes: classifications (List[Classification]): List of detected classifications (e.g., PII, PCI, PHI). is_safe (bool): Whether the inspected content is considered safe. + action (Action): Action to take on the detected issue. severity (Optional[Severity]): Severity level of the detected issue (if any). rules (Optional[List[Rule]]): List of rules that matched during inspection. processed_rules (Optional[List[Rule]]): List of rules that were evaluated (same structure as rules). @@ -172,6 +195,7 @@ class InspectResponse: explanation (Optional[str]): Human-readable explanation of the inspection result. client_transaction_id (Optional[str]): Unique client-provided transaction ID for tracing. event_id (Optional[str]): Unique event ID assigned by the backend. + detected_pii (Optional[List[DetectedPII]]): List of detected PII entities. """ classifications: List[Classification] @@ -184,3 +208,4 @@ class InspectResponse: explanation: Optional[str] = None client_transaction_id: Optional[str] = None event_id: Optional[str] = None + detected_pii: Optional[List[DetectedPII]] = None diff --git a/aidefense/tests/test_inspection_client.py b/aidefense/tests/test_inspection_client.py index 8bb98f7..cc072b9 100644 --- a/aidefense/tests/test_inspection_client.py +++ b/aidefense/tests/test_inspection_client.py @@ -17,6 +17,7 @@ from aidefense.runtime.inspection_client import InspectionClient from aidefense.runtime.models import ( Action, + DetectedPII, InspectResponse, Classification, Severity, @@ -338,3 +339,97 @@ def test_parse_inspect_response_complex(): assert result.client_transaction_id == "tx-9876" assert result.event_id == "b403de99-8d19-408f-8184-ec6d7907f508" assert result.action == Action.ALLOW + + +def test_parse_inspect_response_with_detected_pii(): + """Test parsing a response containing detected_pii entries.""" + client = TestInspectionClient(TEST_API_KEY, Config()) + + response_data = { + "is_safe": False, + "classifications": ["SECURITY_VIOLATION"], + "action": "Allow", + "detected_pii": [ + { + "message_index": "0", + "type": "PII_ENTITY_TYPE_EMAIL", + "start_index": "22", + "end_index": "48", + }, + { + "message_index": "1", + "type": "PII_ENTITY_TYPE_PHONE_NUMBER", + "start_index": "5", + "end_index": "17", + }, + ], + } + + result = client._parse_inspect_response(response_data) + + assert result.detected_pii is not None + assert len(result.detected_pii) == 2 + + assert isinstance(result.detected_pii[0], DetectedPII) + assert result.detected_pii[0].message_index == "0" + assert result.detected_pii[0].type == "PII_ENTITY_TYPE_EMAIL" + assert result.detected_pii[0].start_index == "22" + assert result.detected_pii[0].end_index == "48" + + assert result.detected_pii[1].message_index == "1" + assert result.detected_pii[1].type == "PII_ENTITY_TYPE_PHONE_NUMBER" + assert result.detected_pii[1].start_index == "5" + assert result.detected_pii[1].end_index == "17" + + +def test_parse_inspect_response_without_detected_pii(): + """Test that detected_pii is None when absent from the response.""" + client = TestInspectionClient(TEST_API_KEY, Config()) + + response_data = { + "is_safe": True, + "classifications": [], + } + + result = client._parse_inspect_response(response_data) + + assert result.detected_pii is None + + +def test_parse_inspect_response_with_empty_detected_pii(): + """Test that detected_pii is None when the API returns an empty list.""" + client = TestInspectionClient(TEST_API_KEY, Config()) + + response_data = { + "is_safe": True, + "classifications": [], + "detected_pii": [], + } + + result = client._parse_inspect_response(response_data) + + assert result.detected_pii is None + + +def test_parse_inspect_response_detected_pii_partial_fields(): + """Test parsing detected_pii when some optional fields are missing.""" + client = TestInspectionClient(TEST_API_KEY, Config()) + + response_data = { + "is_safe": False, + "classifications": [], + "detected_pii": [ + { + "type": "PII_ENTITY_TYPE_SSN", + }, + ], + } + + result = client._parse_inspect_response(response_data) + + assert result.detected_pii is not None + assert len(result.detected_pii) == 1 + assert result.detected_pii[0].type == "PII_ENTITY_TYPE_SSN" + assert result.detected_pii[0].message_index is None + assert result.detected_pii[0].start_index is None + assert result.detected_pii[0].end_index is None