|
3 | 3 | from pytest_mock import MockFixture |
4 | 4 |
|
5 | 5 | from assemblyai.streaming.v3 import ( |
| 6 | + SpeechModel, |
6 | 7 | StreamingClient, |
7 | 8 | StreamingClientOptions, |
8 | 9 | StreamingParameters, |
@@ -214,3 +215,50 @@ def mocked_websocket_connect( |
214 | 215 | assert actual_url == f"wss://api.example.com/v3/ws?{urlencode(expected_params)}" |
215 | 216 | assert actual_additional_headers["Authorization"] == "test" |
216 | 217 | assert actual_open_timeout == 15 |
| 218 | + |
| 219 | + |
| 220 | +def test_client_connect_with_u3_pro_and_prompt(mocker: MockFixture): |
| 221 | + actual_url = None |
| 222 | + actual_additional_headers = None |
| 223 | + actual_open_timeout = None |
| 224 | + |
| 225 | + def mocked_websocket_connect( |
| 226 | + url: str, additional_headers: dict, open_timeout: float |
| 227 | + ): |
| 228 | + nonlocal actual_url, actual_additional_headers, actual_open_timeout |
| 229 | + actual_url = url |
| 230 | + actual_additional_headers = additional_headers |
| 231 | + actual_open_timeout = open_timeout |
| 232 | + |
| 233 | + mocker.patch( |
| 234 | + "assemblyai.streaming.v3.client.websocket_connect", |
| 235 | + new=mocked_websocket_connect, |
| 236 | + ) |
| 237 | + |
| 238 | + _disable_rw_threads(mocker) |
| 239 | + |
| 240 | + options = StreamingClientOptions(api_key="test", api_host="api.example.com") |
| 241 | + client = StreamingClient(options) |
| 242 | + |
| 243 | + params = StreamingParameters( |
| 244 | + sample_rate=8000, |
| 245 | + speech_model=SpeechModel.u3_pro, |
| 246 | + min_end_of_turn_silence_when_confident=200, |
| 247 | + prompt="Transcribe this audio with beautiful punctuation and formatting.", |
| 248 | + keyterms_prompt=["yes", "no", "okay"], |
| 249 | + ) |
| 250 | + |
| 251 | + client.connect(params) |
| 252 | + |
| 253 | + # The expected URL should contain all the parameters |
| 254 | + assert "sample_rate=8000" in actual_url |
| 255 | + assert "speech_model=u3-pro" in actual_url |
| 256 | + assert "min_end_of_turn_silence_when_confident=200" in actual_url |
| 257 | + assert "prompt=Transcribe" in actual_url |
| 258 | + assert "keyterms_prompt=" in actual_url # keyterms_prompt is JSON-encoded |
| 259 | + |
| 260 | + assert actual_additional_headers["Authorization"] == "test" |
| 261 | + assert actual_additional_headers["AssemblyAI-Version"] == "2025-05-12" |
| 262 | + assert "AssemblyAI/1.0" in actual_additional_headers["User-Agent"] |
| 263 | + |
| 264 | + assert actual_open_timeout == 15 |
0 commit comments