From ba7868819604be598432fcfc6f2bd3858eb0e8ef Mon Sep 17 00:00:00 2001
From: AssemblyAI <engineering.sdk@assemblyai.com>
Date: Wed, 11 Feb 2026 15:25:19 -0500
Subject: [PATCH] Project import generated by Copybara.

GitOrigin-RevId: 480c9df6e8678f8391a236ba419239a534fc9ea3
---
 assemblyai/streaming/v3/models.py |  2 ++
 tests/unit/test_streaming.py      | 48 +++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/assemblyai/streaming/v3/models.py b/assemblyai/streaming/v3/models.py
index bce64f2..c2a9aa4 100644
--- a/assemblyai/streaming/v3/models.py
+++ b/assemblyai/streaming/v3/models.py
@@ -65,6 +65,7 @@ class StreamingSessionParameters(BaseModel):
     format_turns: Optional[bool] = None
     keyterms_prompt: Optional[List[str]] = None
     filter_profanity: Optional[bool] = None
+    prompt: Optional[str] = None
 
 
 class Encoding(str, Enum):
@@ -78,6 +79,7 @@ def __str__(self):
 class SpeechModel(str, Enum):
     universal_streaming_multilingual = "universal-streaming-multilingual"
     universal_streaming_english = "universal-streaming-english"
+    u3_pro = "u3-pro"
 
     def __str__(self):
         return self.value
diff --git a/tests/unit/test_streaming.py b/tests/unit/test_streaming.py
index 41681a4..84027ce 100644
--- a/tests/unit/test_streaming.py
+++ b/tests/unit/test_streaming.py
@@ -3,6 +3,7 @@
 from pytest_mock import MockFixture
 
 from assemblyai.streaming.v3 import (
+    SpeechModel,
     StreamingClient,
     StreamingClientOptions,
     StreamingParameters,
@@ -214,3 +215,50 @@ def mocked_websocket_connect(
     assert actual_url == f"wss://api.example.com/v3/ws?{urlencode(expected_params)}"
     assert actual_additional_headers["Authorization"] == "test"
     assert actual_open_timeout == 15
+
+
+def test_client_connect_with_u3_pro_and_prompt(mocker: MockFixture):
+    actual_url = None
+    actual_additional_headers = None
+    actual_open_timeout = None
+
+    def mocked_websocket_connect(
+        url: str, additional_headers: dict, open_timeout: float
+    ):
+        nonlocal actual_url, actual_additional_headers, actual_open_timeout
+        actual_url = url
+        actual_additional_headers = additional_headers
+        actual_open_timeout = open_timeout
+
+    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=8000,
+        speech_model=SpeechModel.u3_pro,
+        min_end_of_turn_silence_when_confident=200,
+        prompt="Transcribe this audio with beautiful punctuation and formatting.",
+        keyterms_prompt=["yes", "no", "okay"],
+    )
+
+    client.connect(params)
+
+    # The expected URL should contain all the parameters
+    assert "sample_rate=8000" in actual_url
+    assert "speech_model=u3-pro" in actual_url
+    assert "min_end_of_turn_silence_when_confident=200" in actual_url
+    assert "prompt=Transcribe" in actual_url
+    assert "keyterms_prompt=" in actual_url  # keyterms_prompt is JSON-encoded
+
+    assert actual_additional_headers["Authorization"] == "test"
+    assert actual_additional_headers["AssemblyAI-Version"] == "2025-05-12"
+    assert "AssemblyAI/1.0" in actual_additional_headers["User-Agent"]
+
+    assert actual_open_timeout == 15