From 0de53aa920a470065e8a79aca6a1eb6f212e11dc Mon Sep 17 00:00:00 2001 From: pylrn Date: Fri, 25 Sep 2026 05:33:18 +0530 Subject: [PATCH 1/2] feat(streaming): add granular location PII policies to StreamingPiiPolicy --- assemblyai/__version__.py | 2 +- assemblyai/streaming/v3/models.py | 7 +++++ tests/unit/test_streaming.py | 46 ++++++++++++++++++++++++++++++ tests/unit/test_streaming_async.py | 33 +++++++++++++++++++++ 4 files changed, 87 insertions(+), 1 deletion(-) diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index f49459c..51bbb3f 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.6.1" +__version__ = "1.6.2" diff --git a/assemblyai/streaming/v3/models.py b/assemblyai/streaming/v3/models.py index b278dee..0d4fcb5 100644 --- a/assemblyai/streaming/v3/models.py +++ b/assemblyai/streaming/v3/models.py @@ -284,6 +284,13 @@ class StreamingPiiPolicy(str, Enum): 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_condition = "medical_condition" medical_process = "medical_process" diff --git a/tests/unit/test_streaming.py b/tests/unit/test_streaming.py index 70049a4..8086f08 100644 --- a/tests/unit/test_streaming.py +++ b/tests/unit/test_streaming.py @@ -220,6 +220,52 @@ 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_client_connect_with_voice_focus(mocker: MockFixture): # Given: client + voice_focus parameters actual_url = None diff --git a/tests/unit/test_streaming_async.py b/tests/unit/test_streaming_async.py index 77a455c..c6d999e 100644 --- a/tests/unit/test_streaming_async.py +++ b/tests/unit/test_streaming_async.py @@ -14,6 +14,7 @@ StreamingClientOptions, StreamingEvents, StreamingParameters, + StreamingPiiPolicy, ) from assemblyai.streaming.v3.models import TerminateSession @@ -165,6 +166,38 @@ 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_stream_bytes_writes_to_socket(mocker: MockFixture): fake_ws = _FakeAsyncWebSocket() _patch_connect(mocker, fake_ws) From e08ce62195e089b49af95a5cf9d6c35cb7756a7a Mon Sep 17 00:00:00 2001 From: pylrn Date: Fri, 25 Sep 2026 10:29:34 +0530 Subject: [PATCH 2/2] feat(streaming): complete StreamingPiiPolicy parity with all EntityType members --- assemblyai/streaming/v3/models.py | 11 ++++++ tests/unit/test_streaming.py | 60 ++++++++++++++++++++++++++++++ tests/unit/test_streaming_async.py | 36 ++++++++++++++++++ 3 files changed, 107 insertions(+) diff --git a/assemblyai/streaming/v3/models.py b/assemblyai/streaming/v3/models.py index 0d4fcb5..56b4e07 100644 --- a/assemblyai/streaming/v3/models.py +++ b/assemblyai/streaming/v3/models.py @@ -265,18 +265,22 @@ 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" @@ -292,9 +296,11 @@ class StreamingPiiPolicy(str, Enum): 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" @@ -304,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): diff --git a/tests/unit/test_streaming.py b/tests/unit/test_streaming.py index 8086f08..c81f817 100644 --- a/tests/unit/test_streaming.py +++ b/tests/unit/test_streaming.py @@ -266,6 +266,66 @@ def mocked_websocket_connect( 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 diff --git a/tests/unit/test_streaming_async.py b/tests/unit/test_streaming_async.py index c6d999e..aeb39f7 100644 --- a/tests/unit/test_streaming_async.py +++ b/tests/unit/test_streaming_async.py @@ -198,6 +198,42 @@ async def test_client_connect_with_granular_location_pii_policies( 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)