From 0686eb49f1a6063ad1cecd7d747ea72b0c56eafc Mon Sep 17 00:00:00 2001 From: Adam Kamor Date: Tue, 7 Jul 2026 14:05:21 -0400 Subject: [PATCH 1/2] Done --- .../test_custom_entity_ranking_modes.py | 192 ++++++++++++++++++ .../enums/custom_entity_ranking_mode.py | 6 + tonic_textual/generator_utils.py | 32 ++- tonic_textual/redact_api.py | 66 +++++- 4 files changed, 288 insertions(+), 8 deletions(-) create mode 100644 tests/tests/redact_tests/test_custom_entity_ranking_modes.py create mode 100644 tonic_textual/enums/custom_entity_ranking_mode.py diff --git a/tests/tests/redact_tests/test_custom_entity_ranking_modes.py b/tests/tests/redact_tests/test_custom_entity_ranking_modes.py new file mode 100644 index 0000000..ca88a44 --- /dev/null +++ b/tests/tests/redact_tests/test_custom_entity_ranking_modes.py @@ -0,0 +1,192 @@ +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.enums.custom_entity_ranking_mode import CustomEntityRankingMode +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_ranking_modes_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_ranking_modes={ + "ENTITY_A": CustomEntityRankingMode.Prioritized, + "ENTITY_B": "Standard", + }, + ) + + 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["customEntityRankingModes"] == { + "ENTITY_A": "Prioritized", + "ENTITY_B": "Standard", + } + # Values must serialize as raw strings so the request body JSON matches the API contract + for value in payload["customEntityRankingModes"].values(): + assert type(value) is str + + +@pytest.mark.parametrize("method_name,input_data,endpoint,response", METHOD_CASES, ids=METHOD_IDS) +def test_ranking_modes_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 "customEntityRankingModes" not in requests_made[0]["data"] + + +@pytest.mark.parametrize("method_name,input_data,endpoint,response", METHOD_CASES, ids=METHOD_IDS) +def test_invalid_ranking_mode_raises_before_request_is_sent( + make_mocked_ner, method_name, input_data, endpoint, response +): + ner, requests_made = make_mocked_ner(response) + + with pytest.raises(Exception, match="Invalid value for custom entity ranking modes"): + getattr(ner, method_name)( + input_data, + custom_entities=["ENTITY_A"], + custom_entity_ranking_modes={"ENTITY_A": "NotARealMode"}, + ) + + assert len(requests_made) == 0 + + +def test_generate_redact_payload_serializes_ranking_modes(): + payload = generate_redact_payload( + custom_entities=["ENTITY_A", "ENTITY_B"], + custom_entity_ranking_modes={ + "ENTITY_A": CustomEntityRankingMode.Standard, + "ENTITY_B": "Prioritized", + }, + ) + + assert payload["customEntityRankingModes"] == { + "ENTITY_A": "Standard", + "ENTITY_B": "Prioritized", + } + + +def test_generate_redact_payload_omits_ranking_modes_by_default(): + payload = generate_redact_payload(custom_entities=["ENTITY_A"]) + + assert "customEntityRankingModes" not in payload + + +@pytest.mark.parametrize( + "ranking_mode", + [CustomEntityRankingMode.Prioritized, CustomEntityRankingMode.Standard, "Standard"], + ids=["prioritized-enum", "standard-enum", "standard-string"], +) +def test_redact_with_custom_entity_ranking_modes(textual, ranking_mode): + 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_ranking_modes={custom_entity_name: ranking_mode}, + ) + + assert "hovercraft" not in response.redacted_text + finally: + delete_custom_entity(textual, custom_entity_name) + + +def test_redact_bulk_with_custom_entity_ranking_modes(textual): + custom_entity = create_custom_entity(textual, ["hovercraft"]) + custom_entity_name = custom_entity["name"] + try: + response = textual.redact_bulk( + ["John Smith owns a hovercraft.", "The hovercraft is full of eels."], + custom_entities=[custom_entity_name], + custom_entity_ranking_modes={ + custom_entity_name: CustomEntityRankingMode.Standard + }, + ) + + for redacted in response.bulk_redacted_text: + assert "hovercraft" not in redacted + finally: + delete_custom_entity(textual, custom_entity_name) + + +def test_ranking_mode_for_unrequested_entity_is_rejected(textual): + with pytest.raises((TextualServerBadRequest, InvalidJsonForRedactionRequest)): + textual.redact( + "John Smith is a person", + custom_entity_ranking_modes={"NOT_A_REQUESTED_ENTITY": "Standard"}, + ) diff --git a/tonic_textual/enums/custom_entity_ranking_mode.py b/tonic_textual/enums/custom_entity_ranking_mode.py new file mode 100644 index 0000000..ccd3292 --- /dev/null +++ b/tonic_textual/enums/custom_entity_ranking_mode.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class CustomEntityRankingMode(str, Enum): + Prioritized = "Prioritized" + Standard = "Standard" diff --git a/tonic_textual/generator_utils.py b/tonic_textual/generator_utils.py index 0f12765..e3cad75 100644 --- a/tonic_textual/generator_utils.py +++ b/tonic_textual/generator_utils.py @@ -15,6 +15,7 @@ from tonic_textual.classes.generator_metadata.phone_number_generator_metadata import PhoneNumberGeneratorMetadata from tonic_textual.classes.record_api_request_options import RecordApiRequestOptions from tonic_textual.classes.tonic_exception import BadArgumentsException +from tonic_textual.enums.custom_entity_ranking_mode import CustomEntityRankingMode from tonic_textual.enums.generator_type import GeneratorType from tonic_textual.enums.pii_state import PiiState from tonic_textual.enums.pii_type import PiiType @@ -311,6 +312,22 @@ def convert_payload_to_generator_metadata( return result +def validate_custom_entity_ranking_modes( + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] +) -> None: + if custom_entity_ranking_modes is None: + return + + invalid_values = [ + value for value in list(custom_entity_ranking_modes.values()) + if value not in CustomEntityRankingMode._member_names_ + ] + if len(invalid_values) > 0: + raise Exception( + "Invalid value for custom entity ranking modes. " + "The allowed values are Prioritized and Standard." + ) + def generate_redact_payload( generator_default: PiiState = PiiState.Redaction, generator_config: Dict[str, PiiState] = dict(), @@ -319,14 +336,17 @@ def generate_redact_payload( label_allow_lists: Optional[Dict[str, List[str]]] = None, record_options: Optional[RecordApiRequestOptions] = None, custom_entities: Optional[List[str]] = None, - enable_llm_classification: Optional[bool] = None + enable_llm_classification: Optional[bool] = None, + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None ) -> Dict: validate_generator_default_and_config(generator_default, generator_config, custom_entities) validate_generator_metadata(generator_metadata, custom_entities) - - payload = { + + validate_custom_entity_ranking_modes(custom_entity_ranking_modes) + + payload = { "generatorDefault": generator_default, "generatorConfig": convert_generator_config_to_payload(generator_config), "generatorMetadata": convert_generator_metadata_to_payload(generator_metadata) @@ -340,6 +360,12 @@ 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 label_block_lists is not None: payload["labelBlockLists"] = { diff --git a/tonic_textual/redact_api.py b/tonic_textual/redact_api.py index 0f7c2e0..e8f9fbe 100644 --- a/tonic_textual/redact_api.py +++ b/tonic_textual/redact_api.py @@ -27,6 +27,7 @@ from tonic_textual.classes.audio.redact_audio_responses import ( TranscriptionResult ) +from tonic_textual.enums.custom_entity_ranking_mode import CustomEntityRankingMode from tonic_textual.enums.pii_state import PiiState from tonic_textual.generator_utils import generate_grouping_playload, validate_generator_default_and_config, default_record_options, \ generate_redact_payload, validate_generator_metadata @@ -350,6 +351,7 @@ def redact( record_options: RecordApiRequestOptions = default_record_options, custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, ) -> RedactionResponse: """Redacts a string. Depending on the configured handling for each sensitive data type, values are either redacted, synthesized, or ignored. @@ -407,6 +409,15 @@ def redact( configured on the Textual server. When None (the default), the setting is omitted from the request and the server default (disabled) applies. + + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] + A dictionary of (custom entity type, ranking mode) overrides for + the custom entities detected in this request. Values must be one of + "Prioritized" or "Standard". "Prioritized" means that the custom + entity always wins an exact-boundary overlap against a built-in + entity. "Standard" means that the entities are compared by score. + When omitted, every requested custom entity is treated as + "Prioritized". Returns ------- @@ -439,7 +450,8 @@ def redact( label_allow_lists, record_options, custom_entities, - enable_llm_classification=enable_llm_classification + enable_llm_classification=enable_llm_classification, + custom_entity_ranking_modes=custom_entity_ranking_modes ) payload["text"] = string @@ -457,6 +469,7 @@ def redact_bulk( label_allow_lists: Optional[Dict[str, List[str]]] = None, custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, ) -> BulkRedactionResponse: """Redacts a string. Depending on the configured handling for each sensitive data type, values are either redacted, synthesized, or ignored. @@ -509,6 +522,15 @@ def redact_bulk( configured on the Textual server. When None (the default), the setting is omitted from the request and the server default (disabled) applies. + + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] + A dictionary of (custom entity type, ranking mode) overrides for + the custom entities detected in this request. Values must be one of + "Prioritized" or "Standard". "Prioritized" means that the custom + entity always wins an exact-boundary overlap against a built-in + entity. "Standard" means that the entities are compared by score. + When omitted, every requested custom entity is treated as + "Prioritized". Returns ------- @@ -543,7 +565,8 @@ def redact_bulk( label_allow_lists, None, custom_entities, - enable_llm_classification=enable_llm_classification + enable_llm_classification=enable_llm_classification, + custom_entity_ranking_modes=custom_entity_ranking_modes ) payload["bulkText"] = strings @@ -664,6 +687,7 @@ def redact_json( json_path_ignore_paths: Optional[List[str]] = None, custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = 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 @@ -725,6 +749,15 @@ def redact_json( setting is omitted from the request and the server default (disabled) applies. + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] + A dictionary of (custom entity type, ranking mode) overrides for + the custom entities detected in this request. Values must be one of + "Prioritized" or "Standard". "Prioritized" means that the custom + entity always wins an exact-boundary overlap against a built-in + entity. "Standard" means that the entities are compared by score. + When omitted, every requested custom entity is treated as + "Prioritized". + Returns ------- RedactionResponse @@ -753,7 +786,8 @@ def redact_json( label_allow_lists, None, custom_entities, - enable_llm_classification=enable_llm_classification + enable_llm_classification=enable_llm_classification, + custom_entity_ranking_modes=custom_entity_ranking_modes ) payload["jsonText"] = json_text @@ -776,6 +810,7 @@ def redact_xml( label_allow_lists: Optional[Dict[str, List[str]]] = None, custom_entities: Optional[List[str]] = None, enable_llm_classification: Optional[bool] = None, + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, ) -> RedactionResponse: """Redacts the values in an XML blob. Depending on the configured handling for each entity type, values are either redacted, synthesized, or @@ -827,6 +862,15 @@ def redact_xml( configured on the Textual server. When None (the default), the setting is omitted from the request and the server default (disabled) applies. + + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] + A dictionary of (custom entity type, ranking mode) overrides for + the custom entities detected in this request. Values must be one of + "Prioritized" or "Standard". "Prioritized" means that the custom + entity always wins an exact-boundary overlap against a built-in + entity. "Standard" means that the entities are compared by score. + When omitted, every requested custom entity is treated as + "Prioritized". Returns ------- @@ -846,7 +890,8 @@ def redact_xml( label_allow_lists, None, custom_entities, - enable_llm_classification=enable_llm_classification + enable_llm_classification=enable_llm_classification, + custom_entity_ranking_modes=custom_entity_ranking_modes ) payload["xmlText"] = xml_data @@ -864,6 +909,7 @@ def redact_html( custom_entities: Optional[List[str]] = None, record_options: RecordApiRequestOptions = default_record_options, enable_llm_classification: Optional[bool] = None, + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] = None, ) -> RedactionResponse: """Redacts the values in an HTML blob. Depending on the configured handling for each entity type, values are either redacted, synthesized, or @@ -920,6 +966,15 @@ def redact_html( configured on the Textual server. When None (the default), the setting is omitted from the request and the server default (disabled) applies. + + custom_entity_ranking_modes: Optional[Dict[str, Union[CustomEntityRankingMode, str]]] + A dictionary of (custom entity type, ranking mode) overrides for + the custom entities detected in this request. Values must be one of + "Prioritized" or "Standard". "Prioritized" means that the custom + entity always wins an exact-boundary overlap against a built-in + entity. "Standard" means that the entities are compared by score. + When omitted, every requested custom entity is treated as + "Prioritized". Returns ------- @@ -939,7 +994,8 @@ def redact_html( label_allow_lists, record_options, custom_entities, - enable_llm_classification=enable_llm_classification + enable_llm_classification=enable_llm_classification, + custom_entity_ranking_modes=custom_entity_ranking_modes ) payload["htmlText"] = html_data From 94031b886d1bb7a3ac591d974a6b82af53ed4098 Mon Sep 17 00:00:00 2001 From: Adam Kamor Date: Tue, 7 Jul 2026 14:17:15 -0400 Subject: [PATCH 2/2] updating version --- pyproject.toml | 2 +- tonic_textual/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e52fdcc..206101f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "tonic-textual" -version = "3.20.0" +version = "3.21.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/tonic_textual/__init__.py b/tonic_textual/__init__.py index 10c5c0d..f87b7d8 100644 --- a/tonic_textual/__init__.py +++ b/tonic_textual/__init__.py @@ -1 +1 @@ -__version__ = "3.20.0" +__version__ = "3.21.0"