feat(streaming): complete StreamingPiiPolicy parity with EntityType and server entities

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.6.1"
+__version__ = "1.6.2"

assemblyai/streaming/v3/models.py

@@ -265,29 +265,42 @@ class StreamingPiiPolicy(str, Enum):
     account_number = "account_number"
     banking_information = "banking_information"
     blood_type = "blood_type"
+    corporate_action = "corporate_action"
     credit_card_number = "credit_card_number"
     credit_card_expiration = "credit_card_expiration"
     credit_card_cvv = "credit_card_cvv"
     date = "date"
     date_interval = "date_interval"
     date_of_birth = "date_of_birth"
+    day = "day"
     drivers_license = "drivers_license"
     drug = "drug"
     duration = "duration"
+    effect = "effect"
     email_address = "email_address"
     event = "event"
     filename = "filename"
+    financial_metric = "financial_metric"
     gender_sexuality = "gender_sexuality"
     gender = "gender"
     healthcare_number = "healthcare_number"
     injury = "injury"
     ip_address = "ip_address"
     language = "language"
     location = "location"
+    location_address = "location_address"
+    location_address_street = "location_address_street"
+    location_city = "location_city"
+    location_coordinate = "location_coordinate"
+    location_country = "location_country"
+    location_state = "location_state"
+    location_zip = "location_zip"
     marital_status = "marital_status"
+    medical_code = "medical_code"
     medical_condition = "medical_condition"
     medical_process = "medical_process"
     money_amount = "money_amount"
+    month = "month"
     nationality = "nationality"
     number_sequence = "number_sequence"
     passport_number = "passport_number"
@@ -297,17 +310,22 @@ class StreamingPiiPolicy(str, Enum):
     phone_number = "phone_number"
     physical_attribute = "physical_attribute"
     political_affiliation = "political_affiliation"
+    product = "product"
+    project = "project"
     occupation = "occupation"
     organization = "organization"
+    organization_id = "organization_id"
     organization_medical_facility = "organization_medical_facility"
     religion = "religion"
     sexuality = "sexuality"
     statistics = "statistics"
     time = "time"
+    trend = "trend"
     url = "url"
     us_social_security_number = "us_social_security_number"
     username = "username"
     vehicle_id = "vehicle_id"
+    year = "year"
     zodiac_sign = "zodiac_sign"
 
     def __str__(self):

tests/unit/test_streaming.py

@@ -220,6 +220,112 @@ def mocked_websocket_connect(
     assert "phone_number" in actual_url
 
 
+def test_client_connect_with_granular_location_pii_policies(mocker: MockFixture):
+    # Given: client + granular location PII policies
+    actual_url = None
+
+    def mocked_websocket_connect(
+        url: str, additional_headers: dict, open_timeout: float
+    ):
+        nonlocal actual_url
+        actual_url = url
+
+    mocker.patch(
+        "assemblyai.streaming.v3.client.websocket_connect",
+        new=mocked_websocket_connect,
+    )
+
+    _disable_rw_threads(mocker)
+
+    options = StreamingClientOptions(api_key="test", api_host="api.example.com")
+    client = StreamingClient(options)
+
+    params = StreamingParameters(
+        sample_rate=16000,
+        speech_model=SpeechModel.universal_streaming_english,
+        redact_pii=True,
+        redact_pii_policies=[
+            StreamingPiiPolicy.location_address,
+            StreamingPiiPolicy.location_address_street,
+            StreamingPiiPolicy.location_city,
+            StreamingPiiPolicy.location_coordinate,
+            StreamingPiiPolicy.location_country,
+            StreamingPiiPolicy.location_state,
+            StreamingPiiPolicy.location_zip,
+        ],
+    )
+    client.connect(params)
+
+    assert "redact_pii=True" in actual_url
+    assert "location_address" in actual_url
+    assert "location_address_street" in actual_url
+    assert "location_city" in actual_url
+    assert "location_coordinate" in actual_url
+    assert "location_country" in actual_url
+    assert "location_state" in actual_url
+    assert "location_zip" in actual_url
+
+
+def test_streaming_pii_policy_matches_entity_type():
+    from assemblyai.types import EntityType
+
+    assert set(e.value for e in EntityType) == set(p.value for p in StreamingPiiPolicy)
+
+
+def test_client_connect_with_extended_pii_policies(mocker: MockFixture):
+    # Given: client + extended entity PII policies
+    actual_url = None
+
+    def mocked_websocket_connect(
+        url: str, additional_headers: dict, open_timeout: float
+    ):
+        nonlocal actual_url
+        actual_url = url
+
+    mocker.patch(
+        "assemblyai.streaming.v3.client.websocket_connect",
+        new=mocked_websocket_connect,
+    )
+
+    _disable_rw_threads(mocker)
+
+    options = StreamingClientOptions(api_key="test", api_host="api.example.com")
+    client = StreamingClient(options)
+
+    params = StreamingParameters(
+        sample_rate=16000,
+        speech_model=SpeechModel.universal_streaming_english,
+        redact_pii=True,
+        redact_pii_policies=[
+            StreamingPiiPolicy.corporate_action,
+            StreamingPiiPolicy.day,
+            StreamingPiiPolicy.effect,
+            StreamingPiiPolicy.financial_metric,
+            StreamingPiiPolicy.medical_code,
+            StreamingPiiPolicy.month,
+            StreamingPiiPolicy.organization_id,
+            StreamingPiiPolicy.product,
+            StreamingPiiPolicy.project,
+            StreamingPiiPolicy.trend,
+            StreamingPiiPolicy.year,
+        ],
+    )
+    client.connect(params)
+
+    assert "redact_pii=True" in actual_url
+    assert "corporate_action" in actual_url
+    assert "day" in actual_url
+    assert "effect" in actual_url
+    assert "financial_metric" in actual_url
+    assert "medical_code" in actual_url
+    assert "month" in actual_url
+    assert "organization_id" in actual_url
+    assert "product" in actual_url
+    assert "project" in actual_url
+    assert "trend" in actual_url
+    assert "year" in actual_url
+
+
 def test_client_connect_with_voice_focus(mocker: MockFixture):
     # Given: client + voice_focus parameters
     actual_url = None

tests/unit/test_streaming_async.py

@@ -14,6 +14,7 @@
     StreamingClientOptions,
     StreamingEvents,
     StreamingParameters,
+    StreamingPiiPolicy,
 )
 from assemblyai.streaming.v3.models import TerminateSession
 
@@ -165,6 +166,74 @@ async def test_client_connect_with_speaker_labels_revision_interval(
     await client.disconnect()
 
 
+async def test_client_connect_with_granular_location_pii_policies(
+    mocker: MockFixture,
+):
+    # Given: client + granular location PII policies
+    fake_ws = _FakeAsyncWebSocket()
+    fake_connect = _patch_connect(mocker, fake_ws)
+
+    client = AsyncStreamingClient(
+        StreamingClientOptions(api_key="test", api_host="api.example.com")
+    )
+
+    # When: connecting
+    await client.connect(
+        StreamingParameters(
+            sample_rate=16000,
+            speech_model=SpeechModel.universal_streaming_english,
+            redact_pii=True,
+            redact_pii_policies=[
+                StreamingPiiPolicy.location_city,
+                StreamingPiiPolicy.location_state,
+            ],
+        )
+    )
+
+    # Then: query string carries both granular location policies
+    assert "redact_pii=True" in fake_connect.uri
+    assert "location_city" in fake_connect.uri
+    assert "location_state" in fake_connect.uri
+
+    await client.disconnect()
+
+
+async def test_client_connect_with_extended_pii_policies(
+    mocker: MockFixture,
+):
+    # Given: client + extended entity PII policies
+    fake_ws = _FakeAsyncWebSocket()
+    fake_connect = _patch_connect(mocker, fake_ws)
+
+    client = AsyncStreamingClient(
+        StreamingClientOptions(api_key="test", api_host="api.example.com")
+    )
+
+    # When: connecting
+    await client.connect(
+        StreamingParameters(
+            sample_rate=16000,
+            speech_model=SpeechModel.universal_streaming_english,
+            redact_pii=True,
+            redact_pii_policies=[
+                StreamingPiiPolicy.corporate_action,
+                StreamingPiiPolicy.financial_metric,
+                StreamingPiiPolicy.medical_code,
+                StreamingPiiPolicy.product,
+            ],
+        )
+    )
+
+    # Then: query string carries extended policies
+    assert "redact_pii=True" in fake_connect.uri
+    assert "corporate_action" in fake_connect.uri
+    assert "financial_metric" in fake_connect.uri
+    assert "medical_code" in fake_connect.uri
+    assert "product" in fake_connect.uri
+
+    await client.disconnect()
+
+
 async def test_stream_bytes_writes_to_socket(mocker: MockFixture):
     fake_ws = _FakeAsyncWebSocket()
     _patch_connect(mocker, fake_ws)