SK-1758 Add detect support in Python SDK

- Add unit tests for deidentify and reidentify text

Author: skyflow-vivek

skyflow/utils/__init__.py

@@ -2,4 +2,4 @@
 from ._skyflow_messages import SkyflowMessages
 from ._version import SDK_VERSION
 from ._helpers import get_base_url, format_scope
-from ._utils import get_credentials, get_vault_url, construct_invoke_connection_request, get_metrics, parse_insert_response, handle_exception, parse_update_record_response, parse_delete_response, parse_detokenize_response, parse_tokenize_response, parse_query_response, parse_get_response, parse_invoke_connection_response, validate_api_key, encode_column_values
+from ._utils import get_credentials, get_vault_url, construct_invoke_connection_request, get_metrics, parse_insert_response, handle_exception, parse_update_record_response, parse_delete_response, parse_detokenize_response, parse_tokenize_response, parse_query_response, parse_get_response, parse_invoke_connection_response, validate_api_key, encode_column_values, parse_deidentify_text_response, parse_reidentify_text_response, convert_to_entity_type, convert_detected_entity_to_entity_info 

skyflow/utils/_utils.py

@@ -93,7 +93,7 @@ def convert_to_entity_type(detect_entities):
             entity_types.append(entity.value)
     return entity_types
 
-def _convert_detected_entity_to_entity_info(detected_entity):
+def convert_detected_entity_to_entity_info(detected_entity):
     return EntityInfo(
         token=detected_entity.token,
         value=detected_entity.value,
@@ -392,7 +392,7 @@ def parse_invoke_connection_response(api_response: requests.Response):
         raise SkyflowError(message, status_code)
 
 def parse_deidentify_text_response(api_response: DeidentifyStringResponse):
-    entities = [_convert_detected_entity_to_entity_info(entity) for entity in api_response.entities]
+    entities = [convert_detected_entity_to_entity_info(entity) for entity in api_response.entities]
     return DeidentifyTextResponse(
         processed_text=api_response.processed_text,
         entities=entities,

skyflow/vault/controller/_detect.py

@@ -28,12 +28,14 @@ def ___build_deidentify_text_body(self, request: DeidentifyTextRequest) -> Dict[
         deidentify_text_body = {}
         parsed_entity_types = convert_to_entity_type(request.entities)
 
-        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_token_type = None
+        if request.token_format is not None:
+            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(

tests/utils/test__utils.py

@@ -10,7 +10,8 @@
 from skyflow.utils import get_credentials, SkyflowMessages, get_vault_url, construct_invoke_connection_request, \
     parse_insert_response, parse_update_record_response, parse_delete_response, parse_get_response, \
     parse_detokenize_response, parse_tokenize_response, parse_query_response, parse_invoke_connection_response, \
-    handle_exception, validate_api_key, encode_column_values
+    handle_exception, validate_api_key, encode_column_values, parse_deidentify_text_response, \
+    parse_reidentify_text_response, convert_to_entity_type, convert_detected_entity_to_entity_info
 from skyflow.utils._utils import parse_path_params, to_lowercase_keys, get_metrics
 from skyflow.utils.enums import EnvUrls, Env, ContentType
 from skyflow.vault.connection import InvokeConnectionResponse
@@ -418,3 +419,133 @@ def test_encode_column_values(self):
 
         result = encode_column_values(get_request)
         self.assertEqual(result, expected_encoded_values)
+
+    def test_parse_deidentify_text_response(self):
+        """Test parsing deidentify text response with multiple entities."""
+        mock_entity = Mock()
+        mock_entity.token = "token123"
+        mock_entity.value = "sensitive_value"
+        mock_entity.entity_type = "EMAIL"
+        mock_entity.entity_scores = {"EMAIL": 0.95}
+        mock_entity.location = Mock(
+            start_index=10,
+            end_index=20,
+            start_index_processed=15,
+            end_index_processed=25
+        )
+
+        mock_api_response = Mock()
+        mock_api_response.processed_text = "Sample processed text"
+        mock_api_response.entities = [mock_entity]
+        mock_api_response.word_count = 3
+        mock_api_response.character_count = 20
+
+        result = parse_deidentify_text_response(mock_api_response)
+
+        self.assertEqual(result.processed_text, "Sample processed text")
+        self.assertEqual(result.word_count, 3)
+        self.assertEqual(result.char_count, 20)
+        self.assertEqual(len(result.entities), 1)
+
+        entity = result.entities[0]
+        self.assertEqual(entity.token, "token123")
+        self.assertEqual(entity.value, "sensitive_value")
+        self.assertEqual(entity.entity, "EMAIL")
+        self.assertEqual(entity.scores, {"EMAIL": 0.95})
+        self.assertEqual(entity.text_index.start, 10)
+        self.assertEqual(entity.text_index.end, 20)
+        self.assertEqual(entity.processed_index.start, 15)
+        self.assertEqual(entity.processed_index.end, 25)
+
+    def test_parse_deidentify_text_response_no_entities(self):
+        """Test parsing deidentify text response with no entities."""
+        mock_api_response = Mock()
+        mock_api_response.processed_text = "Sample processed text"
+        mock_api_response.entities = []
+        mock_api_response.word_count = 3
+        mock_api_response.character_count = 20
+
+        result = parse_deidentify_text_response(mock_api_response)
+
+        self.assertEqual(result.processed_text, "Sample processed text")
+        self.assertEqual(result.word_count, 3)
+        self.assertEqual(result.char_count, 20)
+        self.assertEqual(len(result.entities), 0)
+
+    def test_parse_reidentify_text_response(self):
+        """Test parsing reidentify text response."""
+        mock_api_response = Mock()
+        mock_api_response.processed_text = "Reidentified text with actual values"
+
+        result = parse_reidentify_text_response(mock_api_response)
+
+        self.assertEqual(result.processed_text, "Reidentified text with actual values")
+
+    def test_convert_to_entity_type_with_valid_entities(self):
+        """Test converting entity types with valid input."""
+        from skyflow.utils.enums import DetectEntities
+
+        detect_entities = [DetectEntities.EMAIL_ADDRESS, DetectEntities.PHONE_NUMBER]
+        result = convert_to_entity_type(detect_entities)
+
+        self.assertEqual(result, ["email_address", "phone_number"])
+
+    def test_convert_to_entity_type_with_empty_list(self):
+        """Test converting entity types with empty list."""
+        result = convert_to_entity_type([])
+        self.assertIsNone(result)
+
+    def test_convert_to_entity_type_with_none(self):
+        """Test converting entity types with None input."""
+        result = convert_to_entity_type(None)
+        self.assertIsNone(result)
+
+    def test__convert_detected_entity_to_entity_info(self):
+        """Test converting detected entity to EntityInfo object."""
+        mock_detected_entity = Mock()
+        mock_detected_entity.token = "token123"
+        mock_detected_entity.value = "sensitive_value"
+        mock_detected_entity.entity_type = "EMAIL"
+        mock_detected_entity.entity_scores = {"EMAIL": 0.95}
+        mock_detected_entity.location = Mock(
+            start_index=10,
+            end_index=20,
+            start_index_processed=15,
+            end_index_processed=25
+        )
+
+        result = convert_detected_entity_to_entity_info(mock_detected_entity)
+
+        self.assertEqual(result.token, "token123")
+        self.assertEqual(result.value, "sensitive_value")
+        self.assertEqual(result.entity, "EMAIL")
+        self.assertEqual(result.scores, {"EMAIL": 0.95})
+        self.assertEqual(result.text_index.start, 10)
+        self.assertEqual(result.text_index.end, 20)
+        self.assertEqual(result.processed_index.start, 15)
+        self.assertEqual(result.processed_index.end, 25)
+
+    def test__convert_detected_entity_to_entity_info_with_minimal_data(self):
+        """Test converting detected entity with minimal required data."""
+        mock_detected_entity = Mock()
+        mock_detected_entity.token = "token123"
+        mock_detected_entity.value = None
+        mock_detected_entity.entity_type = "UNKNOWN"
+        mock_detected_entity.entity_scores = {}
+        mock_detected_entity.location = Mock(
+            start_index=0,
+            end_index=0,
+            start_index_processed=0,
+            end_index_processed=0
+        )
+
+        result = convert_detected_entity_to_entity_info(mock_detected_entity)
+
+        self.assertEqual(result.token, "token123")
+        self.assertIsNone(result.value)
+        self.assertEqual(result.entity, "UNKNOWN")
+        self.assertEqual(result.scores, {})
+        self.assertEqual(result.text_index.start, 0)
+        self.assertEqual(result.text_index.end, 0)
+        self.assertEqual(result.processed_index.start, 0)
+        self.assertEqual(result.processed_index.end, 0)

tests/vault/controller/test__detect.py

@@ -0,0 +1,118 @@
+import unittest
+from unittest.mock import Mock, patch
+from skyflow.error import SkyflowError
+from skyflow.vault.controller import Detect
+from skyflow.vault.detect import DeidentifyTextRequest, ReidentifyTextRequest, \
+    TokenFormat, DateTransformation, Transformations
+from skyflow.utils.enums import DetectEntities, TokenType
+
+VAULT_ID = "test_vault_id"
+
+class TestDetect(unittest.TestCase):
+    def setUp(self):
+        # Mock vault client
+        self.vault_client = Mock()
+        self.vault_client.get_vault_id.return_value = VAULT_ID
+        self.vault_client.get_logger.return_value = Mock()
+
+        # Create a Detect instance with the mock client
+        self.detect = Detect(self.vault_client)
+
+    @patch("skyflow.vault.controller._detect.validate_deidentify_text_request")
+    @patch("skyflow.vault.controller._detect.parse_deidentify_text_response")
+    def test_deidentify_text_success(self, mock_parse_response, mock_validate):
+        # Mock API response
+        mock_api_response = Mock()
+        mock_api_response.data = {
+            'text': '[TOKEN_1] lives in [TOKEN_2]',
+            'entities': [
+                {
+                    'token': 'Token1',
+                    'value': 'John',
+                    'text_index': {'start': 0, 'end': 4},
+                    'processed_index': {'start': 0, 'end': 8},
+                    'entity': 'NAME',
+                    'scores': {'confidence': 0.9}
+                }
+            ],
+            'word_count': 4,
+            'char_count': 20
+        }
+
+        # Create request
+        request = DeidentifyTextRequest(
+            text="John lives in NYC",
+            entities=[DetectEntities.NAME],
+            token_format=TokenFormat(default=TokenType.ENTITY_ONLY)
+        )
+
+        # Mock detect API
+        detect_api = self.vault_client.get_detect_text_api.return_value
+        detect_api.deidentify_string.return_value = mock_api_response
+
+        # Call deidentify_text
+        response = self.detect.deidentify_text(request)
+
+        # Assertions
+        mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
+        mock_parse_response.assert_called_once_with(mock_api_response)
+        detect_api.deidentify_string.assert_called_once()
+
+    @patch("skyflow.vault.controller._detect.validate_reidentify_text_request")
+    @patch("skyflow.vault.controller._detect.parse_reidentify_text_response")
+    def test_reidentify_text_success(self, mock_parse_response, mock_validate):
+        # Mock API response
+        mock_api_response = Mock()
+        mock_api_response.data = {
+            'text': 'John lives in NYC'
+        }
+
+        # Create request
+        request = ReidentifyTextRequest(
+            text="Token1 lives in Token2",
+            redacted_entities=[DetectEntities.NAME]
+        )
+
+        # Mock detect API
+        detect_api = self.vault_client.get_detect_text_api.return_value
+        detect_api.reidentify_string.return_value = mock_api_response
+
+        # Call reidentify_text
+        response = self.detect.reidentify_text(request)
+
+        # Assertions
+        mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
+        mock_parse_response.assert_called_once_with(mock_api_response)
+        detect_api.reidentify_string.assert_called_once()
+
+    @patch("skyflow.vault.controller._detect.validate_deidentify_text_request")
+    def test_deidentify_text_handles_generic_error(self, mock_validate):
+        request = DeidentifyTextRequest(
+            text="John lives in NYC",
+            entities=[DetectEntities.NAME]
+        )
+        detect_api = self.vault_client.get_detect_text_api.return_value
+        detect_api.deidentify_string.side_effect = Exception("Generic Error")
+
+        with self.assertRaises(Exception):
+            self.detect.deidentify_text(request)
+
+        detect_api.deidentify_string.assert_called_once()
+
+    @patch("skyflow.vault.controller._detect.validate_reidentify_text_request")
+    def test_reidentify_text_handles_generic_error(self, mock_validate):
+        request = ReidentifyTextRequest(
+            text="Token1 lives in Token2",
+            redacted_entities=[DetectEntities.NAME]
+        )
+        detect_api = self.vault_client.get_detect_text_api.return_value
+        detect_api.reidentify_string.side_effect = Exception("Generic Error")
+
+        with self.assertRaises(Exception):
+            self.detect.reidentify_text(request)
+
+        detect_api.reidentify_string.assert_called_once()
+
+
+# if __name__ == '__main__':
+#     unittest.main()