SK-1758 Add detect support in Python SDK

- Add deidentify and reidentify text support in Python SDK

Author: skyflow-vivek

skyflow/utils/_skyflow_messages.py

@@ -258,22 +258,22 @@ class Info(Enum):
         DEIDENTIFY_TEXT_TRIGGERED = f"{INFO}: [{error_prefix}] Deidentify text method triggered."
         VALIDATING_DEIDENTIFY_TEXT_INPUT = f"{INFO}: [{error_prefix}] Validating deidentify text input."
         DEIDENTIFY_TEXT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Deidentify text request is resolved."
-        DEIDENTIFY_TEXT_SUCCESS = f"{INFO}: [{error_prefix}] data deidentified." 
+        DEIDENTIFY_TEXT_SUCCESS = f"{INFO}: [{error_prefix}] Data deidentified." 
 
         REIDENTIFY_TEXT_TRIGGERED = f"{INFO}: [{error_prefix}] Reidentify text method triggered."
         VALIDATING_REIDENTIFY_TEXT_INPUT = f"{INFO}: [{error_prefix}] Validating reidentify text input."
         REIDENTIFY_TEXT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Reidentify text request is resolved."
-        REIDENTIFY_TEXT_SUCCESS = f"{INFO}: [{error_prefix}] data reidentified." 
+        REIDENTIFY_TEXT_SUCCESS = f"{INFO}: [{error_prefix}] Data reidentified." 
 
         DEIDENTIFY_FILE_TRIGGERED = f"{INFO}: [{error_prefix}] Deidentify file triggered."
         VALIDATING_DETECT_FILE_INPUT = f"{INFO}: [{error_prefix}] Validating deidentify file input."
         DEIDENTIFY_FILE_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Deidentify file request is resolved."
-        DEIDENTIFY_FILE_SUCCESS = f"{INFO}: [{error_prefix}] file deidentified." 
+        DEIDENTIFY_FILE_SUCCESS = f"{INFO}: [{error_prefix}] File deidentified." 
 
         GET_DETECT_RUN_TRIGGERED = f"{INFO}: [{error_prefix}] Get detect run triggered."
         VALIDATING_GET_DETECT_RUN_INPUT = f"{INFO}: [{error_prefix}] Validating get detect run input."
         GET_DETECT_RUN_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Get detect run request is resolved."
-        GET_DETECT_RUN_SUCCESS = f"{INFO}: [{error_prefix}] get detect run success." # will see for a better message
+        GET_DETECT_RUN_SUCCESS = f"{INFO}: [{error_prefix}] Get detect run success." # will see for a better message
 
         DETECT_REQUEST_RESOLVED = f"{INFO}: [{error_prefix}] Detect request is resolved." # Whats this for?
 

skyflow/utils/_utils.py

@@ -18,6 +18,7 @@
 from skyflow.generated.rest.core.http_response import HttpResponse
 from skyflow.utils.logger import log_error_log
 from skyflow.vault.detect import DeidentifyTextResponse, ReidentifyTextResponse
+from skyflow.vault.detect import EntityInfo, TextIndex
 from . import SkyflowMessages, SDK_VERSION
 from .constants import PROTOCOL
 from .enums import Env, ContentType, EnvUrls
@@ -84,6 +85,30 @@ def to_lowercase_keys(dict):
 
     return result
 
+def convert_to_entity_type(detect_entities):
+    entity_types = None
+    if (detect_entities is not None and len(detect_entities) != 0):
+        entity_types = []
+        for entity in detect_entities:
+            entity_types.append(entity.value)
+    return entity_types
+
+def _convert_detected_entity_to_entity_info(detected_entity):
+    return EntityInfo(
+        token=detected_entity.token,
+        value=detected_entity.value,
+        text_index=TextIndex(
+            start=detected_entity.location.start_index,
+            end=detected_entity.location.end_index
+        ),
+        processed_index=TextIndex(
+            start=detected_entity.location.start_index_processed,
+            end=detected_entity.location.end_index_processed
+        ),
+        entity=detected_entity.entity_type,
+        scores=detected_entity.entity_scores
+    )
+
 def construct_invoke_connection_request(request, connection_url, logger) -> PreparedRequest:
     url = parse_path_params(connection_url.rstrip('/'), request.path_params)
 
@@ -367,11 +392,13 @@ def parse_invoke_connection_response(api_response: requests.Response):
         raise SkyflowError(message, status_code)
 
 def parse_deidentify_text_response(api_response: DeidentifyStringResponse):
-    processed_text = api_response.processed_text
-    entities = api_response.entities
-    word_count = api_response.word_count
-    character_count = api_response.character_count
-    return DeidentifyTextResponse(processed_text, entities, word_count, character_count)
+    entities = [_convert_detected_entity_to_entity_info(entity) for entity in api_response.entities]
+    return DeidentifyTextResponse(
+        processed_text=api_response.processed_text,
+        entities=entities,
+        word_count=api_response.word_count,
+        char_count=api_response.character_count
+    )
 
 def parse_reidentify_text_response(api_response: ReidentifyStringResponse):
     return ReidentifyTextResponse(api_response.processed_text)

skyflow/vault/controller/_detect.py

@@ -1,6 +1,9 @@
 import json
+from skyflow.generated.rest.types.token_type import TokenType
+from skyflow.generated.rest.types.transformations import Transformations
+from skyflow.generated.rest.types.transformations_shift_dates import TransformationsShiftDates
 from skyflow.utils._skyflow_messages import SkyflowMessages
-from skyflow.utils._utils import get_metrics, handle_exception, parse_deidentify_text_response, parse_reidentify_text_response
+from skyflow.utils._utils import convert_to_entity_type, get_metrics, handle_exception, parse_deidentify_text_response, parse_reidentify_text_response
 from skyflow.utils.constants import SKY_META_DATA_HEADER
 from skyflow.utils.logger import log_info, log_error_log
 from skyflow.utils.validations._validations import validate_deidentify_text_request, validate_reidentify_text_request
@@ -23,29 +26,42 @@ def __get_headers(self):
 
     def ___build_deidentify_text_body(self, request: DeidentifyTextRequest) -> Dict[str, Any]:
         deidentify_text_body = {}
-        entity_types = []
-        deidentify_text_body['text'] = request.text
-
-        for entity in request.entities:
-            entity_type = entity.value
-            entity_types.append(entity_type)
+        parsed_entity_types = convert_to_entity_type(request.entities)
 
-        deidentify_text_body['entity_types'] = entity_types
-        deidentify_text_body['token_type'] = request.token_format
+        parsed_token_type = TokenType(
+            default = request.token_format.default,
+            vault_token = convert_to_entity_type(request.token_format.vault_token),
+            entity_unq_counter = convert_to_entity_type(request.token_format.entity_unique_counter),
+            entity_only = convert_to_entity_type(request.token_format.entity_only)
+        )
+        parsed_transformations = None
+        if request.transformations is not None:
+            parsed_transformations = Transformations(
+                shift_dates = TransformationsShiftDates(
+                    max_days = request.transformations.shift_days.max,
+                    min_days = request.transformations.shift_days.min,
+                    entity_types = convert_to_entity_type(request.transformations.shift_days.entities)
+                )
+            )
+        
+        deidentify_text_body['text'] = request.text
+        deidentify_text_body['entity_types'] = parsed_entity_types
+        deidentify_text_body['token_type'] = parsed_token_type
         deidentify_text_body['allow_regex'] = request.allow_regex_list
         deidentify_text_body['restrict_regex'] = request.restrict_regex_list 
-        deidentify_text_body['transformations'] = request.transformations 
+        deidentify_text_body['transformations'] = parsed_transformations 
+        
         return deidentify_text_body
 
     def ___build_reidentify_text_body(self, request: ReidentifyTextRequest) -> Dict[str, Any]:
-        format_obj = ReidentifyStringRequestFormat(
-            redacted=request.redacted_entities,
-            masked=request.masked_entities,
-            plaintext=request.plain_text_entities
+        parsed_format = ReidentifyStringRequestFormat(
+            redacted=convert_to_entity_type(request.redacted_entities),
+            masked=convert_to_entity_type(request.masked_entities),
+            plaintext=convert_to_entity_type(request.plain_text_entities)
         )
-        reidentify_text_body = []
+        reidentify_text_body = {}
         reidentify_text_body['text'] = request.text
-        reidentify_text_body['format'] = format_obj
+        reidentify_text_body['format'] = parsed_format
         return reidentify_text_body
 
     def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextResponse:
@@ -59,12 +75,12 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo
         try:
             log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_TRIGGERED.value, self.__vault_client.get_logger())
             api_response = detect_api.deidentify_string(
-              self.__vault_client.get_vault_id(),
-              request.text,
-              entity_types=deidentify_text_body['entities'],
-              allows_regex=deidentify_text_body['allow_regex_list'],
-              restrict_regex=deidentify_text_body['restrict_regex_list'],
-              token_type=deidentify_text_body['token_format'],
+              vault_id=self.__vault_client.get_vault_id(),
+              text=deidentify_text_body['text'],
+              entity_types=deidentify_text_body['entity_types'],
+              allow_regex=deidentify_text_body['allow_regex'],
+              restrict_regex=deidentify_text_body['restrict_regex'],
+              token_type=deidentify_text_body['token_type'],
               transformations=deidentify_text_body['transformations'],
               request_options=self.__get_headers()
             )
@@ -78,7 +94,7 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo
 
     def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextResponse:
         log_info(SkyflowMessages.Info.VALIDATING_REIDENTIFY_TEXT_INPUT.value, self.__vault_client.get_logger())
-        validate_deidentify_text_request(self.__vault_client.get_logger(), request)
+        validate_reidentify_text_request(self.__vault_client.get_logger(), request)
         log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
         self.__initialize()
         detect_api = self.__vault_client.get_detect_text_api()
@@ -87,8 +103,8 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo
         try:
             log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_TRIGGERED.value, self.__vault_client.get_logger())
             api_response = detect_api.reidentify_string(
-                text=reidentify_text_body['text'],
                 vault_id=self.__vault_client.get_vault_id(),
+                text=reidentify_text_body['text'],
                 format=reidentify_text_body['format'],
                 request_options=self.__get_headers()
             )
@@ -97,5 +113,5 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo
             return reidentify_text_response
 
         except Exception as e:
-            log_error_log(SkyflowMessages.ErrorLogs.DEIDENTIFY_TEXT_REQUEST_REJECTED.value, self.__vault_client.get_logger())
+            log_error_log(SkyflowMessages.ErrorLogs.REIDENTIFY_TEXT_REQUEST_REJECTED.value, self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())

skyflow/vault/detect/_deidentify_text_response.py

@@ -1,50 +1,19 @@
-from dataclasses import dataclass
-from typing import List, Dict
-from ._text_index import TextIndex
+from typing import List
 from ._entity_info import EntityInfo
 
-@dataclass
 class DeidentifyTextResponse:
-    processed_text: str
-    entities: List[EntityInfo]
-    word_count: int
-    char_count: int
-
-    @property
-    def processed_text(self) -> str:
-        return self._processed_text
-
-    @processed_text.setter
-    def processed_text(self, value: str):
-        self._processed_text = value
-
-    @property
-    def entities(self) -> List[EntityInfo]:
-        return self._entities
-
-    @entities.setter
-    def entities(self, value: List[EntityInfo]):
-        self._entities = value
-
-    @property
-    def word_count(self) -> int:
-        return self._word_count
-
-    @word_count.setter
-    def word_count(self, value: int):
-        self._word_count = value
-
-    @property
-    def char_count(self) -> int:
-        return self._char_count
-
-    @char_count.setter
-    def char_count(self, value: int):
-        self._char_count = value
-
-    def __init__(self, processed_text: str, entities: List[EntityInfo], 
-                 word_count: int, char_count: int):
-        self._processed_text = processed_text
-        self._entities = entities
-        self._word_count = word_count
-        self._char_count = char_count
+    def __init__(self, 
+                processed_text: str,
+                entities: List[EntityInfo], 
+                word_count: int,
+                char_count: int):
+        self.processed_text = processed_text
+        self.entities = entities
+        self.word_count = word_count
+        self.char_count = char_count
+
+    def __repr__(self):
+        return f"DeidentifyTextResponse(processed_text='{self.processed_text}', entities={self.entities}, word_count={self.word_count}, char_count={self.char_count})"
+
+    def __str__(self):
+        return self.__repr__()

skyflow/vault/detect/_entity_info.py

@@ -10,3 +10,11 @@ def __init__(self, token: str, value: str, text_index: TextIndex,
         self.processed_index = processed_index
         self.entity = entity
         self.scores = scores
+
+    def __repr__(self) -> str:
+        return (f"EntityInfo(token='{self.token}', value='{self.value}', "
+                f"text_index={self.text_index}, processed_index={self.processed_index}, "
+                f"entity='{self.entity}', scores={self.scores})")
+
+    def __str__(self) -> str:
+        return self.__repr__()

skyflow/vault/detect/_reidentify_text_request.py

@@ -3,9 +3,9 @@
 
 class ReidentifyTextRequest:
     def __init__(self, text: str,
-                 redacted_entities: Optional[List[DetectEntities]] = None,
-                 masked_entities: Optional[List[DetectEntities]] = None,
-                 plain_text_entities: Optional[List[DetectEntities]] = None):
+                redacted_entities: Optional[List[DetectEntities]] = None,
+                masked_entities: Optional[List[DetectEntities]] = None,
+                plain_text_entities: Optional[List[DetectEntities]] = None):
         self.text = text
         self.redacted_entities = redacted_entities
         self.masked_entities = masked_entities

skyflow/vault/detect/_reidentify_text_response.py

@@ -1,3 +1,9 @@
 class ReidentifyTextResponse:
     def __init__(self, processed_text: str):
         self.processed_text = processed_text
+
+    def __repr__(self) -> str:
+        return f"ReidentifyTextResponse(processed_text='{self.processed_text}')"
+
+    def __str__(self) -> str:
+        return self.__repr__()

skyflow/vault/detect/_text_index.py

@@ -2,3 +2,9 @@ class TextIndex:
     def __init__(self, start: int, end: int):
         self.start = start
         self.end = end
+
+    def __repr__(self):
+        return f"TextIndex(start={self.start}, end={self.end})"
+        
+    def __str__(self):
+        return self.__repr__()