Skip to content

Commit ec30ffd

Browse files
zkleb-aaiAssemblyAI
andauthored
chore: sync sdk code with DeepLearning repo (#168)
Co-authored-by: AssemblyAI <engineering.sdk@assemblyai.com>
1 parent f1340f2 commit ec30ffd

2 files changed

Lines changed: 50 additions & 0 deletions

File tree

assemblyai/streaming/v3/models.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ class StreamingSessionParameters(BaseModel):
6565
format_turns: Optional[bool] = None
6666
keyterms_prompt: Optional[List[str]] = None
6767
filter_profanity: Optional[bool] = None
68+
prompt: Optional[str] = None
6869

6970

7071
class Encoding(str, Enum):
@@ -78,6 +79,7 @@ def __str__(self):
7879
class SpeechModel(str, Enum):
7980
universal_streaming_multilingual = "universal-streaming-multilingual"
8081
universal_streaming_english = "universal-streaming-english"
82+
u3_pro = "u3-pro"
8183

8284
def __str__(self):
8385
return self.value

tests/unit/test_streaming.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from pytest_mock import MockFixture
44

55
from assemblyai.streaming.v3 import (
6+
SpeechModel,
67
StreamingClient,
78
StreamingClientOptions,
89
StreamingParameters,
@@ -214,3 +215,50 @@ def mocked_websocket_connect(
214215
assert actual_url == f"wss://api.example.com/v3/ws?{urlencode(expected_params)}"
215216
assert actual_additional_headers["Authorization"] == "test"
216217
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

Comments
 (0)