Add detected_pii field in inspection response

aidefense/runtime/__init__.py

@@ -25,6 +25,7 @@
     InspectionConfig,
     Metadata,
     InspectResponse,
+    DetectedPII,
 )
 from .http_models import HttpInspectRequest
 from .http_models import (

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:

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,13 +187,15 @@ 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).
         attack_technique (Optional[str]): Attack technique detected, if applicable.
         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

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