Transcribe streaming audio

1
import pyaudio
2
import websocket
3
import json
4
import threading
5
import time
6
import wave
7
from urllib.parse import urlencode
8
from datetime import datetime
9
10
# --- Configuration ---
11
YOUR_API_KEY = "YOUR-API-KEY"  # Replace with your actual API key
12
13
CONNECTION_PARAMS = {
14
    "speech_model": "u3-rt-pro",
15
    "sample_rate": 16000,
16
}
17
API_ENDPOINT_BASE_URL = "wss://streaming.assemblyai.com/v3/ws"
18
API_ENDPOINT = f"{API_ENDPOINT_BASE_URL}?{urlencode(CONNECTION_PARAMS)}"
19
20
# Audio Configuration
21
FRAMES_PER_BUFFER = 800  # 50ms of audio (0.05s * 16000Hz)
22
SAMPLE_RATE = CONNECTION_PARAMS["sample_rate"]
23
CHANNELS = 1
24
FORMAT = pyaudio.paInt16
25
26
# Global variables for audio stream and websocket
27
audio = None
28
stream = None
29
ws_app = None
30
audio_thread = None
31
stop_event = threading.Event()  # To signal the audio thread to stop
32
33
# WAV recording variables
34
recorded_frames = []  # Store audio frames for WAV file
35
recording_lock = threading.Lock()  # Thread-safe access to recorded_frames
36
37
def save_wav_file():
38
    """Save recorded audio frames to a WAV file."""
39
    if not recorded_frames:
40
        print("No audio data recorded.")
41
        return
42
43
    # Generate filename with timestamp
44
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
45
    filename = f"recorded_audio_{timestamp}.wav"
46
47
    try:
48
        with wave.open(filename, 'wb') as wf:
49
            wf.setnchannels(CHANNELS)
50
            wf.setsampwidth(2)  # 16-bit = 2 bytes
51
            wf.setframerate(SAMPLE_RATE)
52
53
            # Write all recorded frames
54
            with recording_lock:
55
                wf.writeframes(b''.join(recorded_frames))
56
57
        print(f"Audio saved to: {filename}")
58
        print(f"Duration: {len(recorded_frames) * FRAMES_PER_BUFFER / SAMPLE_RATE:.2f} seconds")
59
60
    except Exception as e:
61
        print(f"Error saving WAV file: {e}")
62
63
# --- WebSocket Event Handlers ---
64
65
def on_open(ws):
66
    """Called when the WebSocket connection is established."""
67
    print("WebSocket connection opened.")
68
    print(f"Connected to: {API_ENDPOINT}")
69
70
    # Start sending audio data in a separate thread
71
    def stream_audio():
72
        global stream
73
        print("Starting audio streaming...")
74
        while not stop_event.is_set():
75
            try:
76
                audio_data = stream.read(FRAMES_PER_BUFFER, exception_on_overflow=False)
77
78
                # Store audio data for WAV recording
79
                with recording_lock:
80
                    recorded_frames.append(audio_data)
81
82
                # Send audio data as binary message
83
                ws.send(audio_data, websocket.ABNF.OPCODE_BINARY)
84
            except Exception as e:
85
                print(f"Error streaming audio: {e}")
86
                # If stream read fails, likely means it's closed, stop the loop
87
                break
88
        print("Audio streaming stopped.")
89
90
    global audio_thread
91
    audio_thread = threading.Thread(target=stream_audio)
92
    audio_thread.daemon = (
93
        True  # Allow main thread to exit even if this thread is running
94
    )
95
    audio_thread.start()
96
97
def on_message(ws, message):
98
    try:
99
        data = json.loads(message)
100
        msg_type = data.get('type')
101
102
        if msg_type == "Begin":
103
            session_id = data.get('id')
104
            expires_at = data.get('expires_at')
105
            print(f"\nSession began: ID={session_id}, ExpiresAt={datetime.fromtimestamp(expires_at)}")
106
        elif msg_type == "Turn":
107
            transcript = data.get('transcript', '')
108
            if data.get('end_of_turn'):
109
                print('\r' + ' ' * 80 + '\r', end='')
110
                print(transcript)
111
            else:
112
                print(f"\r{transcript}", end='')
113
        elif msg_type == "Termination":
114
            audio_duration = data.get('audio_duration_seconds', 0)
115
            session_duration = data.get('session_duration_seconds', 0)
116
            print(f"\nSession Terminated: Audio Duration={audio_duration}s, Session Duration={session_duration}s")
117
    except json.JSONDecodeError as e:
118
        print(f"Error decoding message: {e}")
119
    except Exception as e:
120
        print(f"Error handling message: {e}")
121
122
def on_error(ws, error):
123
    """Called when a WebSocket error occurs."""
124
    print(f"\nWebSocket Error: {error}")
125
    # Attempt to signal stop on error
126
    stop_event.set()
127
128
129
def on_close(ws, close_status_code, close_msg):
130
    """Called when the WebSocket connection is closed."""
131
    print(f"\nWebSocket Disconnected: Status={close_status_code}, Msg={close_msg}")
132
133
    # Save recorded audio to WAV file
134
    save_wav_file()
135
136
    # Ensure audio resources are released
137
    global stream, audio
138
    stop_event.set()  # Signal audio thread just in case it's still running
139
140
    if stream:
141
        if stream.is_active():
142
            stream.stop_stream()
143
        stream.close()
144
        stream = None
145
    if audio:
146
        audio.terminate()
147
        audio = None
148
    # Try to join the audio thread to ensure clean exit
149
    if audio_thread and audio_thread.is_alive():
150
        audio_thread.join(timeout=1.0)
151
152
# --- Main Execution ---
153
def run():
154
    global audio, stream, ws_app
155
156
    # Initialize PyAudio
157
    audio = pyaudio.PyAudio()
158
159
    # Open microphone stream
160
    try:
161
        stream = audio.open(
162
            input=True,
163
            frames_per_buffer=FRAMES_PER_BUFFER,
164
            channels=CHANNELS,
165
            format=FORMAT,
166
            rate=SAMPLE_RATE,
167
        )
168
        print("Microphone stream opened successfully.")
169
        print("Speak into your microphone. Press Ctrl+C to stop.")
170
        print("Audio will be saved to a WAV file when the session ends.")
171
    except Exception as e:
172
        print(f"Error opening microphone stream: {e}")
173
        if audio:
174
            audio.terminate()
175
        return  # Exit if microphone cannot be opened
176
177
    # Create WebSocketApp
178
    ws_app = websocket.WebSocketApp(
179
        API_ENDPOINT,
180
        header={"Authorization": YOUR_API_KEY},
181
        on_open=on_open,
182
        on_message=on_message,
183
        on_error=on_error,
184
        on_close=on_close,
185
    )
186
187
    # Run WebSocketApp in a separate thread to allow main thread to catch KeyboardInterrupt
188
    ws_thread = threading.Thread(target=ws_app.run_forever)
189
    ws_thread.daemon = True
190
    ws_thread.start()
191
192
    try:
193
        # Keep main thread alive until interrupted
194
        while ws_thread.is_alive():
195
            time.sleep(0.1)
196
    except KeyboardInterrupt:
197
        print("\nCtrl+C received. Stopping...")
198
        stop_event.set()  # Signal audio thread to stop
199
200
        # Send termination message to the server
201
        if ws_app and ws_app.sock and ws_app.sock.connected:
202
            try:
203
                terminate_message = {"type": "Terminate"}
204
                print(f"Sending termination message: {json.dumps(terminate_message)}")
205
                ws_app.send(json.dumps(terminate_message))
206
                # Give a moment for messages to process before forceful close
207
                time.sleep(5)
208
            except Exception as e:
209
                print(f"Error sending termination message: {e}")
210
211
        # Close the WebSocket connection (will trigger on_close)
212
        if ws_app:
213
            ws_app.close()
214
215
        # Wait for WebSocket thread to finish
216
        ws_thread.join(timeout=2.0)
217
218
    except Exception as e:
219
        print(f"\nAn unexpected error occurred: {e}")
220
        stop_event.set()
221
        if ws_app:
222
            ws_app.close()
223
        ws_thread.join(timeout=2.0)
224
225
    finally:
226
        # Final cleanup (already handled in on_close, but good as a fallback)
227
        if stream and stream.is_active():
228
            stream.stop_stream()
229
        if stream:
230
            stream.close()
231
        if audio:
232
            audio.terminate()
233
        print("Cleanup complete. Exiting.")
234
235
236
if __name__ == "__main__":
237
    run()

1
# --- WebSocket Event Handlers ---
2
def on_open(ws):
3
    """Called when the WebSocket connection is established."""
4
    print("WebSocket connection opened.")
5
    print(f"Connected to: {API_ENDPOINT}")
6
    # Start sending audio data in a separate thread
7
    def stream_audio():
8
        global stream
9
        print("Starting audio streaming...")
10
        while not stop_event.is_set():
11
            try:
12
                audio_data = stream.read(FRAMES_PER_BUFFER, exception_on_overflow=False)
13
14
                # Store audio data for WAV recording
15
                with recording_lock:
16
                    recorded_frames.append(audio_data)
17
18
                # Send audio data as binary message
19
                ws.send(audio_data, websocket.ABNF.OPCODE_BINARY)
20
            except Exception as e:
21
                print(f"Error streaming audio: {e}")
22
                # If stream read fails, likely means it's closed, stop the loop
23
                break
24
        print("Audio streaming stopped.")
25
    global audio_thread
26
    audio_thread = threading.Thread(target=stream_audio)
27
    audio_thread.daemon = (
28
        True  # Allow main thread to exit even if this thread is running
29
    )
30
    audio_thread.start()
31
def on_message(ws, message):
32
    try:
33
        data = json.loads(message)
34
        msg_type = data.get('type')
35
        if msg_type == "Begin":
36
            session_id = data.get('id')
37
            expires_at = data.get('expires_at')
38
            print(f"\nSession began: ID={session_id}, ExpiresAt={datetime.fromtimestamp(expires_at)}")
39
        elif msg_type == "Turn":
40
            transcript = data.get('transcript', '')
41
            if data.get('end_of_turn'):
42
                print('\r' + ' ' * 80 + '\r', end='')
43
                print(transcript)
44
            else:
45
                print(f"\r{transcript}", end='')
46
        elif msg_type == "Termination":
47
            audio_duration = data.get('audio_duration_seconds', 0)
48
            session_duration = data.get('session_duration_seconds', 0)
49
            print(f"\nSession Terminated: Audio Duration={audio_duration}s, Session Duration={session_duration}s")
50
    except json.JSONDecodeError as e:
51
        print(f"Error decoding message: {e}")
52
    except Exception as e:
53
        print(f"Error handling message: {e}")
54
def on_error(ws, error):
55
    """Called when a WebSocket error occurs."""
56
    print(f"\nWebSocket Error: {error}")
57
    # Attempt to signal stop on error
58
    stop_event.set()
59
def on_close(ws, close_status_code, close_msg):
60
    """Called when the WebSocket connection is closed."""
61
    print(f"\nWebSocket Disconnected: Status={close_status_code}, Msg={close_msg}")
62
63
    # Save recorded audio to WAV file
64
    save_wav_file()
65
66
    # Ensure audio resources are released
67
    global stream, audio
68
    stop_event.set()  # Signal audio thread just in case it's still running
69
    if stream:
70
        if stream.is_active():
71
            stream.stop_stream()
72
        stream.close()
73
        stream = None
74
    if audio:
75
        audio.terminate()
76
        audio = None
77
    # Try to join the audio thread to ensure clean exit
78
    if audio_thread and audio_thread.is_alive():
79
        audio_thread.join(timeout=1.0)

1
   try:
2
        # Keep main thread alive until interrupted
3
        while ws_thread.is_alive():
4
            time.sleep(0.1)
5
    except KeyboardInterrupt:
6
        print("\nCtrl+C received. Stopping...")
7
        stop_event.set()  # Signal audio thread to stop
8
        # Send termination message to the server
9
        if ws_app and ws_app.sock and ws_app.sock.connected:
10
            try:
11
                terminate_message = {"type": "Terminate"}
12
                print(f"Sending termination message: {json.dumps(terminate_message)}")
13
                ws_app.send(json.dumps(terminate_message))
14
                # Give a moment for messages to process before forceful close
15
                time.sleep(5)
16
            except Exception as e:
17
                print(f"Error sending termination message: {e}")
18
        # Close the WebSocket connection (will trigger on_close)
19
        if ws_app:
20
            ws_app.close()
21
        # Wait for WebSocket thread to finish
22
        ws_thread.join(timeout=2.0)
23
    except Exception as e:
24
        print(f"\nAn unexpected error occurred: {e}")
25
        stop_event.set()
26
        if ws_app:
27
            ws_app.close()
28
        ws_thread.join(timeout=2.0)
29
    finally:
30
        # Final cleanup (already handled in on_close, but good as a fallback)
31
        if stream and stream.is_active():
32
            stream.stop_stream()
33
        if stream:
34
            stream.close()
35
        if audio:
36
            audio.terminate()
37
        print("Cleanup complete. Exiting.")