Filter profanity

1
import pyaudio
2
import websocket
3
import json
4
import threading
5
import time
6
from urllib.parse import urlencode
7
8
YOUR_API_KEY = "<YOUR_API_KEY>"
9
CONNECTION_PARAMS = {
10
    "sample_rate": 16000,
11
    "speech_model": "u3-rt-pro",
12
    "format_turns": "true",
13
    "filter_profanity": "true",
14
}
15
API_ENDPOINT_BASE_URL = "wss://streaming.assemblyai.com/v3/ws"
16
API_ENDPOINT = f"{API_ENDPOINT_BASE_URL}?{urlencode(CONNECTION_PARAMS)}"
17
18
FRAMES_PER_BUFFER = 800
19
SAMPLE_RATE = CONNECTION_PARAMS["sample_rate"]
20
CHANNELS = 1
21
FORMAT = pyaudio.paInt16
22
23
audio = None
24
stream = None
25
ws_app = None
26
audio_thread = None
27
stop_event = threading.Event()
28
29
def on_open(ws):
30
    print("WebSocket connection opened.")
31
32
    def stream_audio():
33
        global stream
34
        while not stop_event.is_set():
35
            try:
36
                audio_data = stream.read(FRAMES_PER_BUFFER, exception_on_overflow=False)
37
                ws.send(audio_data, websocket.ABNF.OPCODE_BINARY)
38
            except Exception as e:
39
                print(f"Error streaming audio: {e}")
40
                break
41
42
    global audio_thread
43
    audio_thread = threading.Thread(target=stream_audio)
44
    audio_thread.daemon = True
45
    audio_thread.start()
46
47
def on_message(ws, message):
48
    try:
49
        data = json.loads(message)
50
        msg_type = data.get("type")
51
        if msg_type == "Begin":
52
            print(f"Session began: ID={data.get('id')}")
53
        elif msg_type == "Turn":
54
            transcript = data.get("transcript", "")
55
            end_of_turn = data.get("end_of_turn", False)
56
            if end_of_turn:
57
                print(f"\r{' ' * 80}\r{transcript}")
58
        elif msg_type == "Termination":
59
            print(f"\nSession terminated: {data.get('audio_duration_seconds', 0)}s of audio")
60
    except Exception as e:
61
        print(f"Error handling message: {e}")
62
63
def on_error(ws, error):
64
    print(f"\nWebSocket Error: {error}")
65
    stop_event.set()
66
67
def on_close(ws, close_status_code, close_msg):
68
    print(f"\nWebSocket Disconnected: Status={close_status_code}")
69
    global stream, audio
70
    stop_event.set()
71
    if stream:
72
        if stream.is_active():
73
            stream.stop_stream()
74
        stream.close()
75
    if audio:
76
        audio.terminate()
77
78
def run():
79
    global audio, stream, ws_app
80
    audio = pyaudio.PyAudio()
81
    stream = audio.open(
82
        input=True,
83
        frames_per_buffer=FRAMES_PER_BUFFER,
84
        channels=CHANNELS,
85
        format=FORMAT,
86
        rate=SAMPLE_RATE,
87
    )
88
    print("Speak into your microphone. Press Ctrl+C to stop.")
89
    ws_app = websocket.WebSocketApp(
90
        API_ENDPOINT,
91
        header={"Authorization": YOUR_API_KEY},
92
        on_open=on_open,
93
        on_message=on_message,
94
        on_error=on_error,
95
        on_close=on_close,
96
    )
97
    ws_thread = threading.Thread(target=ws_app.run_forever)
98
    ws_thread.daemon = True
99
    ws_thread.start()
100
    try:
101
        while ws_thread.is_alive():
102
            time.sleep(0.1)
103
    except KeyboardInterrupt:
104
        print("\nStopping...")
105
        stop_event.set()
106
        if ws_app and ws_app.sock and ws_app.sock.connected:
107
            ws_app.send(json.dumps({"type": "Terminate"}))
108
            time.sleep(2)
109
        if ws_app:
110
            ws_app.close()
111
        ws_thread.join(timeout=2.0)
112
113
if __name__ == "__main__":
114
    run()