Custom entity ranking modes

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 <adam@tonic.ai>", "Joe Ferrara <joeferrara@tonic.ai>", "Ander Steele <ander@tonic.ai>", "Ethan Philpott <ephilpott@tonic.ai>", "Lyon Van Voorhis <lyon@tonic.ai>", "Kirill Medvedev <kirill@tonic.ai>", "Travis Matthews <travis@tonic.ai>"]
 license = "MIT"

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",
+        "<root><name>John Smith</name></root>",
+        "/api/redact/xml",
+        SINGLE_RESPONSE,
+    ),
+    ("redact_html", "<p>John Smith</p>", "/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"},
+        )

tonic_textual/__init__.py

@@ -1 +1 @@
-__version__ = "3.20.0"
+__version__ = "3.21.0"

tonic_textual/enums/custom_entity_ranking_mode.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class CustomEntityRankingMode(str, Enum):
+    Prioritized = "Prioritized"
+    Standard = "Standard"

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"] = {