Added changes for prompt group

pebblo/entity_classifier/entity_classifier.py

@@ -6,7 +6,7 @@
     ConfidenceScore,
     Entities,
     SecretEntities,
-    entity_conf_mapping,
+    entity_group_conf_mapping,
 )
 from pebblo.entity_classifier.utils.utils import (
     add_custom_regex_analyzer_registry,
@@ -55,32 +55,27 @@ def analyze_response(self, input_text, anonymize_all_entities=True):
         """
         # Analyze the text to detect entities using the Presidio analyzer
         analyzer_results = self.analyzer.analyze(text=input_text, language="en")
-
         # Initialize the list to hold the final classified entities
         final_results = []
-
         # Iterate through the detected entities
         for entity in analyzer_results:
             try:
                 mapped_entity = None
-
                 # Map entity type to predefined entities if it exists in the Entities enumeration
                 if entity.entity_type in Entities.__members__:
                     mapped_entity = Entities[entity.entity_type].value
-
                 # Check if the entity type exists in SecretEntities enumeration
                 elif entity.entity_type in SecretEntities.__members__:
                     mapped_entity = SecretEntities[entity.entity_type].value
-
                 # Append entity to final results if it meets the confidence threshold and is in the desired entities list
                 if (
                     mapped_entity
-                    and entity.score >= float(entity_conf_mapping[mapped_entity])
+                    and entity.score
+                    >= float(entity_group_conf_mapping[mapped_entity][0])
                     and entity.entity_type in self.entities
                 ):
                     final_results.append(entity)
 
-            # Handle any exceptions that occur during entity classification
             except Exception as ex:
                 logger.warning(
                     f"Error in analyze_response in entity classification. {str(ex)}"
@@ -101,8 +96,13 @@ def anonymize_response(self, analyzer_results, input_text):
     def get_analyzed_entities_response(data, anonymized_response=None):
         # Returns entities with its location i.e. start to end and confidence score
         response = []
-
         for index, value in enumerate(data):
+            mapped_entity = None
+            if value.entity_type in Entities.__members__:
+                mapped_entity = Entities[value.entity_type].value
+            elif value.entity_type in SecretEntities.__members__:
+                mapped_entity = SecretEntities[value.entity_type].value
+
             location = f"{value.start}_{value.end}"
             if anonymized_response:
                 anonymized_data = anonymized_response[len(data) - index - 1]
@@ -112,6 +112,7 @@ def get_analyzed_entities_response(data, anonymized_response=None):
                     "entity_type": value.entity_type,
                     "location": location,
                     "confidence_score": value.score,
+                    "entity_group": entity_group_conf_mapping[mapped_entity][1],
                 }
             )
         return response

pebblo/entity_classifier/utils/config.py

@@ -21,6 +21,9 @@ class Entities(Enum):
     US_PASSPORT = "us-passport-number"
     US_DRIVER_LICENSE = "us-drivers-license"
 
+    # network
+    IP_ADDRESS = "ip-address"
+
     # Financial
     CREDIT_CARD = "credit-card-number"
     US_BANK_NUMBER = "us-bank-account-number"
@@ -38,24 +41,33 @@ class SecretEntities(Enum):
     GOOGLE_API_KEY = "google-api-key"
 
 
-entity_conf_mapping = {
+class PIIGroups(Enum):
+    Identification = "pii-identification"
+    Financial = "pii-financial"
+    Secrets = "secrets_and_tokens"
+    Network = "pii-network"
+
+
+entity_group_conf_mapping = {
     # Identification
-    Entities.US_SSN.value: 0.8,
-    Entities.US_PASSPORT.value: 0.4,
-    Entities.US_DRIVER_LICENSE.value: 0.4,
+    Entities.US_SSN.value: (0.8, PIIGroups.Identification.value),
+    Entities.US_PASSPORT.value: (0.4, PIIGroups.Identification.value),
+    Entities.US_DRIVER_LICENSE.value: (0.4, PIIGroups.Identification.value),
     # Financial
-    Entities.US_ITIN.value: 0.8,
-    Entities.CREDIT_CARD.value: 0.8,
-    Entities.US_BANK_NUMBER.value: 0.4,
-    Entities.IBAN_CODE.value: 0.8,
+    Entities.US_ITIN.value: (0.8, PIIGroups.Financial.value),
+    Entities.CREDIT_CARD.value: (0.8, PIIGroups.Financial.value),
+    Entities.US_BANK_NUMBER.value: (0.4, PIIGroups.Financial.value),
+    Entities.IBAN_CODE.value: (0.8, PIIGroups.Financial.value),
     # Secret
-    SecretEntities.GITHUB_TOKEN.value: 0.8,
-    SecretEntities.SLACK_TOKEN.value: 0.8,
-    SecretEntities.AWS_ACCESS_KEY.value: 0.45,
-    SecretEntities.AWS_SECRET_KEY.value: 0.8,
-    SecretEntities.AZURE_KEY_ID.value: 0.8,
-    SecretEntities.AZURE_CLIENT_SECRET.value: 0.8,
-    SecretEntities.GOOGLE_API_KEY.value: 0.8,
+    SecretEntities.GITHUB_TOKEN.value: (0.8, PIIGroups.Secrets.value),
+    SecretEntities.SLACK_TOKEN.value: (0.8, PIIGroups.Secrets.value),
+    SecretEntities.AWS_ACCESS_KEY.value: (0.45, PIIGroups.Secrets.value),
+    SecretEntities.AWS_SECRET_KEY.value: (0.8, PIIGroups.Secrets.value),
+    SecretEntities.AZURE_KEY_ID.value: (0.8, PIIGroups.Secrets.value),
+    SecretEntities.AZURE_CLIENT_SECRET.value: (0.8, PIIGroups.Secrets.value),
+    SecretEntities.GOOGLE_API_KEY.value: (0.8, PIIGroups.Secrets.value),
+    # Network
+    Entities.IP_ADDRESS.value: (0.4, PIIGroups.Network.value),
 }
 
 

pebblo/entity_classifier/utils/utils.py

@@ -8,6 +8,7 @@
     ConfidenceScore,
     Entities,
     SecretEntities,
+    entity_group_conf_mapping,
     secret_entities_context_mapping,
 )
 from pebblo.entity_classifier.utils.regex_pattern import regex_secrets_patterns
@@ -32,6 +33,7 @@ def get_entities(entities_list, response):
                 "confidence_score": get_confidence_score_label(
                     entity["confidence_score"]
                 ),
+                "entity_group": entity_group_conf_mapping[mapped_entity][1],
             }
             if mapped_entity in entity_details.keys():
                 entity_details[mapped_entity].append(entity_data)

tests/app/service/test_prompt_gov.py

@@ -25,6 +25,7 @@ def test_process_request_success(mock_entity_classifier):
                 {
                     "location": "16_27",
                     "confidence_score": "HIGH",
+                    "entity_group": "pii-identification",
                 }
             ]
         },

tests/entity_classifier/mock_response.py

@@ -24,4 +24,5 @@
 Azure Client Secret - c4cb6f91-15a7-4e6d-a824-abcdef012345 
 <PERSON> - <SLACK_TOKEN> 
 Google API key- KLzaSyB_tWrbmfWx8g2bzL7Vhq7znuTUn0JPKmY"
+IP Address - <IP_ADDRESS>
 """

tests/entity_classifier/test_data.py

@@ -25,6 +25,7 @@
 Azure Client Secret - c4cb6f91-15a7-4e6d-a824-abcdef012345 
 Slack Token - xoxb-3204014939555-4519358291237-TTIf0243T8YFSAGEVr1wBrWE 
 Google API key- KLzaSyB_tWrbmfWx8g2bzL7Vhq7znuTUn0JPKmY"
+My IP Address - 10.55.60.61 
 """
 
 negative_data = """