diff --git a/pyproject.toml b/pyproject.toml index 206101f..ea4a29b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "tonic-textual" -version = "3.21.0" +version = "3.22.0" description = "Wrappers around the Tonic Textual API" authors = ["Adam Kamor ", "Joe Ferrara ", "Ander Steele ", "Ethan Philpott ", "Lyon Van Voorhis ", "Kirill Medvedev ", "Travis Matthews "] license = "MIT" diff --git a/tests/tests/redact_tests/test_custom_entity_confidence_thresholds.py b/tests/tests/redact_tests/test_custom_entity_confidence_thresholds.py new file mode 100644 index 0000000..b20ce27 --- /dev/null +++ b/tests/tests/redact_tests/test_custom_entity_confidence_thresholds.py @@ -0,0 +1,193 @@ +import math + +import pytest + +from tests.utils.redact_utils import create_custom_entity, delete_custom_entity +from tonic_textual.classes.tonic_exception import ( + InvalidJsonForRedactionRequest, + TextualServerBadRequest, +) +from tonic_textual.generator_utils import generate_redact_payload +from tonic_textual.redact_api import TextualNer + +SINGLE_RESPONSE = { + "originalText": "John Smith is a person", + "redactedText": "[NAME_GIVEN_x] [NAME_FAMILY_y] is a person", + "usage": 5, + "deIdentifyResults": [], +} + +BULK_RESPONSE = { + "bulkText": ["John Smith is a person", "I live in Atlanta"], + "bulkRedactedText": [ + "[NAME_GIVEN_x] [NAME_FAMILY_y] is a person", + "I live in [LOCATION_CITY_z]", + ], + "usage": 9, + "deIdentifyResults": [], +} + +METHOD_CASES = [ + ("redact", "John Smith is a person", "/api/redact", SINGLE_RESPONSE), + ( + "redact_bulk", + ["John Smith is a person", "I live in Atlanta"], + "/api/redact/bulk", + BULK_RESPONSE, + ), + ("redact_json", '{"name": "John Smith"}', "/api/redact/json", SINGLE_RESPONSE), + ( + "redact_xml", + "John Smith", + "/api/redact/xml", + SINGLE_RESPONSE, + ), + ("redact_html", "

John Smith

", "/api/redact/html", SINGLE_RESPONSE), +] + +METHOD_IDS = [case[0] for case in METHOD_CASES] + + +@pytest.fixture +def make_mocked_ner(monkeypatch): + """Builds a TextualNer whose http_post is replaced with a stub that records + the request payload and returns a canned response.""" + + def _make(response): + ner = TextualNer(base_url="http://localhost", api_key="fake-key") + requests_made = [] + + def fake_http_post( + url, params={}, data={}, files={}, additional_headers={}, timeout_seconds=None + ): + requests_made.append({"url": url, "data": data}) + return response + + monkeypatch.setattr(ner.client, "http_post", fake_http_post) + return ner, requests_made + + return _make + + +@pytest.mark.parametrize("method_name,input_data,endpoint,response", METHOD_CASES, ids=METHOD_IDS) +def test_confidence_thresholds_are_sent_in_payload( + make_mocked_ner, method_name, input_data, endpoint, response +): + ner, requests_made = make_mocked_ner(response) + + getattr(ner, method_name)( + input_data, + custom_entities=["ENTITY_A", "ENTITY_B"], + custom_entity_confidence_thresholds={ + "ENTITY_A": 0.5, + "ENTITY_B": 1.0, + }, + ) + + assert len(requests_made) == 1 + assert requests_made[0]["url"] == endpoint + + payload = requests_made[0]["data"] + assert payload["customPiiEntityIds"] == ["ENTITY_A", "ENTITY_B"] + assert payload["customEntityConfidenceThresholds"] == { + "ENTITY_A": 0.5, + "ENTITY_B": 1.0, + } + # Preserve numeric type so JSON serialization matches the backend contract. + for value in payload["customEntityConfidenceThresholds"].values(): + assert isinstance(value, (int, float)) + assert not isinstance(value, bool) + + +@pytest.mark.parametrize("method_name,input_data,endpoint,response", METHOD_CASES, ids=METHOD_IDS) +def test_confidence_thresholds_are_omitted_when_not_supplied( + make_mocked_ner, method_name, input_data, endpoint, response +): + ner, requests_made = make_mocked_ner(response) + + getattr(ner, method_name)(input_data, custom_entities=["ENTITY_A"]) + + assert len(requests_made) == 1 + assert "customEntityConfidenceThresholds" not in requests_made[0]["data"] + + +INVALID_THRESHOLD_CASES = [ + ("zero", 0), + ("negative", -0.1), + ("above-one", 1.1), + ("nan", math.nan), + ("positive-infinity", math.inf), + ("negative-infinity", -math.inf), + ("non-numeric-string", "0.5"), + ("non-numeric-none", None), +] + + +@pytest.mark.parametrize("method_name,input_data,endpoint,response", METHOD_CASES, ids=METHOD_IDS) +@pytest.mark.parametrize( + "bad_value", + [case[1] for case in INVALID_THRESHOLD_CASES], + ids=[case[0] for case in INVALID_THRESHOLD_CASES], +) +def test_invalid_confidence_threshold_raises_before_request_is_sent( + make_mocked_ner, method_name, input_data, endpoint, response, bad_value +): + ner, requests_made = make_mocked_ner(response) + + with pytest.raises(Exception, match="Invalid value for custom entity confidence thresholds"): + getattr(ner, method_name)( + input_data, + custom_entities=["ENTITY_A"], + custom_entity_confidence_thresholds={"ENTITY_A": bad_value}, + ) + + assert len(requests_made) == 0 + + +def test_generate_redact_payload_serializes_confidence_thresholds(): + payload = generate_redact_payload( + custom_entities=["ENTITY_A", "ENTITY_B"], + custom_entity_confidence_thresholds={ + "ENTITY_A": 0.25, + "ENTITY_B": 1, + }, + ) + + assert payload["customEntityConfidenceThresholds"] == { + "ENTITY_A": 0.25, + "ENTITY_B": 1, + } + + +def test_generate_redact_payload_omits_confidence_thresholds_by_default(): + payload = generate_redact_payload(custom_entities=["ENTITY_A"]) + + assert "customEntityConfidenceThresholds" not in payload + + +@pytest.mark.parametrize( + "threshold", + [0.1, 0.5, 1.0], + ids=["low", "mid", "boundary-one"], +) +def test_redact_with_custom_entity_confidence_thresholds(textual, threshold): + custom_entity = create_custom_entity(textual, ["hovercraft"]) + custom_entity_name = custom_entity["name"] + try: + response = textual.redact( + "John Smith owns a hovercraft.", + custom_entities=[custom_entity_name], + custom_entity_confidence_thresholds={custom_entity_name: threshold}, + ) + + assert response is not None + finally: + delete_custom_entity(textual, custom_entity_name) + + +def test_confidence_threshold_for_unrequested_entity_is_rejected(textual): + with pytest.raises((TextualServerBadRequest, InvalidJsonForRedactionRequest)): + textual.redact( + "John Smith is a person", + custom_entity_confidence_thresholds={"NOT_A_REQUESTED_ENTITY": 0.5}, + ) diff --git a/tonic_textual/__init__.py b/tonic_textual/__init__.py index f87b7d8..659cac3 100644 --- a/tonic_textual/__init__.py +++ b/tonic_textual/__init__.py @@ -1 +1 @@ -__version__ = "3.21.0" +__version__ = "3.22.0" diff --git a/tonic_textual/generator_utils.py b/tonic_textual/generator_utils.py index e3cad75..004d49f 100644 --- a/tonic_textual/generator_utils.py +++ b/tonic_textual/generator_utils.py @@ -1,3 +1,4 @@ +import math from typing import Dict, List, Optional, Union from tonic_textual.classes.common_api_responses.label_custom_list import LabelCustomList @@ -328,6 +329,24 @@ def validate_custom_entity_ranking_modes( "The allowed values are Prioritized and Standard." ) +def validate_custom_entity_confidence_thresholds( + custom_entity_confidence_thresholds: Optional[Dict[str, float]] +) -> None: + if custom_entity_confidence_thresholds is None: + return + + for value in custom_entity_confidence_thresholds.values(): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise Exception( + "Invalid value for custom entity confidence thresholds. " + "Values must be finite numbers greater than 0 and no greater than 1." + ) + if not math.isfinite(value) or value <= 0 or value > 1: + raise Exception( + "Invalid value for custom entity confidence thresholds. " + "Values must be finite numbers greater than 0 and no greater than 1." + ) + def generate_redact_payload( generator_default: PiiState = PiiState.Redaction, generator_config: Dict[str, PiiState] = dict(), @@ -337,15 +356,18 @@ def generate_redact_payload( record_options: Optional[RecordApiRequestOptions] = None, custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, - custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, + custom_entity_confidence_thresholds: Optional[Dict[str, float]] = None ) -> Dict: - + validate_generator_default_and_config(generator_default, generator_config, custom_entities) validate_generator_metadata(generator_metadata, custom_entities) validate_custom_entity_ranking_modes(custom_entity_ranking_modes) + validate_custom_entity_confidence_thresholds(custom_entity_confidence_thresholds) + payload = { "generatorDefault": generator_default, "generatorConfig": convert_generator_config_to_payload(generator_config), @@ -360,13 +382,16 @@ def generate_redact_payload( payload["llmClassificationPolicy"] = ( "Enabled" if enable_llm_classification else "Disabled" ) - + if custom_entity_ranking_modes is not None: payload["customEntityRankingModes"] = { k: CustomEntityRankingMode(v).value for k, v in custom_entity_ranking_modes.items() } + if custom_entity_confidence_thresholds is not None: + payload["customEntityConfidenceThresholds"] = dict(custom_entity_confidence_thresholds) + if label_block_lists is not None: payload["labelBlockLists"] = { k: LabelCustomList(regexes=v).to_dict() diff --git a/tonic_textual/redact_api.py b/tonic_textual/redact_api.py index e8f9fbe..6a57097 100644 --- a/tonic_textual/redact_api.py +++ b/tonic_textual/redact_api.py @@ -352,6 +352,7 @@ def redact( custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, + custom_entity_confidence_thresholds: Optional[Dict[str, float]] = None, ) -> RedactionResponse: """Redacts a string. Depending on the configured handling for each sensitive data type, values are either redacted, synthesized, or ignored. @@ -419,6 +420,14 @@ def redact( When omitted, every requested custom entity is treated as "Prioritized". + custom_entity_confidence_thresholds: Optional[Dict[str, float]] + A dictionary of (custom entity type, confidence threshold) + overrides. Keys must correspond to a requested custom-PII or + model-based entity name. Values must be finite numbers greater + than 0 and no greater than 1. Detections with a score below the + threshold are dropped. When omitted, the existing no-cutoff + behavior is preserved. + Returns ------- RedactionResponse @@ -451,7 +460,8 @@ def redact( record_options, custom_entities, enable_llm_classification=enable_llm_classification, - custom_entity_ranking_modes=custom_entity_ranking_modes + custom_entity_ranking_modes=custom_entity_ranking_modes, + custom_entity_confidence_thresholds=custom_entity_confidence_thresholds ) payload["text"] = string @@ -470,6 +480,7 @@ def redact_bulk( custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, + custom_entity_confidence_thresholds: Optional[Dict[str, float]] = None, ) -> BulkRedactionResponse: """Redacts a string. Depending on the configured handling for each sensitive data type, values are either redacted, synthesized, or ignored. @@ -532,6 +543,14 @@ def redact_bulk( When omitted, every requested custom entity is treated as "Prioritized". + custom_entity_confidence_thresholds: Optional[Dict[str, float]] + A dictionary of (custom entity type, confidence threshold) + overrides. Keys must correspond to a requested custom-PII or + model-based entity name. Values must be finite numbers greater + than 0 and no greater than 1. Detections with a score below the + threshold are dropped. When omitted, the existing no-cutoff + behavior is preserved. + Returns ------- RedactionResponse @@ -566,7 +585,8 @@ def redact_bulk( None, custom_entities, enable_llm_classification=enable_llm_classification, - custom_entity_ranking_modes=custom_entity_ranking_modes + custom_entity_ranking_modes=custom_entity_ranking_modes, + custom_entity_confidence_thresholds=custom_entity_confidence_thresholds ) payload["bulkText"] = strings @@ -688,6 +708,7 @@ def redact_json( custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, + custom_entity_confidence_thresholds: Optional[Dict[str, float]] = None, ) -> RedactionResponse: """Redacts the values in a JSON blob. Depending on the configured handling for each sensitive data type, values are either redacted, synthesized, or @@ -758,6 +779,14 @@ def redact_json( When omitted, every requested custom entity is treated as "Prioritized". + custom_entity_confidence_thresholds: Optional[Dict[str, float]] + A dictionary of (custom entity type, confidence threshold) + overrides. Keys must correspond to a requested custom-PII or + model-based entity name. Values must be finite numbers greater + than 0 and no greater than 1. Detections with a score below the + threshold are dropped. When omitted, the existing no-cutoff + behavior is preserved. + Returns ------- RedactionResponse @@ -787,7 +816,8 @@ def redact_json( None, custom_entities, enable_llm_classification=enable_llm_classification, - custom_entity_ranking_modes=custom_entity_ranking_modes + custom_entity_ranking_modes=custom_entity_ranking_modes, + custom_entity_confidence_thresholds=custom_entity_confidence_thresholds ) payload["jsonText"] = json_text @@ -811,6 +841,7 @@ def redact_xml( custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, + custom_entity_confidence_thresholds: Optional[Dict[str, float]] = None, ) -> RedactionResponse: """Redacts the values in an XML blob. Depending on the configured handling for each entity type, values are either redacted, synthesized, or @@ -872,6 +903,14 @@ def redact_xml( When omitted, every requested custom entity is treated as "Prioritized". + custom_entity_confidence_thresholds: Optional[Dict[str, float]] + A dictionary of (custom entity type, confidence threshold) + overrides. Keys must correspond to a requested custom-PII or + model-based entity name. Values must be finite numbers greater + than 0 and no greater than 1. Detections with a score below the + threshold are dropped. When omitted, the existing no-cutoff + behavior is preserved. + Returns ------- RedactionResponse @@ -891,7 +930,8 @@ def redact_xml( None, custom_entities, enable_llm_classification=enable_llm_classification, - custom_entity_ranking_modes=custom_entity_ranking_modes + custom_entity_ranking_modes=custom_entity_ranking_modes, + custom_entity_confidence_thresholds=custom_entity_confidence_thresholds ) payload["xmlText"] = xml_data @@ -910,6 +950,7 @@ def redact_html( record_options: RecordApiRequestOptions = default_record_options, enable_llm_classification: Optional[bool] = None, custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, + custom_entity_confidence_thresholds: Optional[Dict[str, float]] = None, ) -> RedactionResponse: """Redacts the values in an HTML blob. Depending on the configured handling for each entity type, values are either redacted, synthesized, or @@ -976,6 +1017,14 @@ def redact_html( When omitted, every requested custom entity is treated as "Prioritized". + custom_entity_confidence_thresholds: Optional[Dict[str, float]] + A dictionary of (custom entity type, confidence threshold) + overrides. Keys must correspond to a requested custom-PII or + model-based entity name. Values must be finite numbers greater + than 0 and no greater than 1. Detections with a score below the + threshold are dropped. When omitted, the existing no-cutoff + behavior is preserved. + Returns ------- RedactionResponse @@ -995,7 +1044,8 @@ def redact_html( record_options, custom_entities, enable_llm_classification=enable_llm_classification, - custom_entity_ranking_modes=custom_entity_ranking_modes + custom_entity_ranking_modes=custom_entity_ranking_modes, + custom_entity_confidence_thresholds=custom_entity_confidence_thresholds ) payload["htmlText"] = html_data