deepgram.clients.common.v1.abstract_async_websocket

  1# Copyright 2023-2024 Deepgram SDK contributors. All Rights Reserved.
  2# Use of this source code is governed by a MIT license that can be found in the LICENSE file.
  3# SPDX-License-Identifier: MIT
  4import asyncio
  5import json
  6import logging
  7from typing import Dict, Union, Optional, cast, Any, Callable
  8from datetime import datetime
  9import threading
 10from abc import ABC, abstractmethod
 11
 12import websockets
 13
 14try:
 15    # Websockets versions >= 13
 16    from websockets.asyncio.client import connect, ClientConnection
 17
 18    WS_ADDITIONAL_HEADERS_KEY = "additional_headers"
 19except ImportError:
 20    # Backward compatibility with websockets versions 12
 21    from websockets.legacy.client import (  # type: ignore
 22        connect,
 23        WebSocketClientProtocol as ClientConnection,
 24    )
 25
 26    WS_ADDITIONAL_HEADERS_KEY = "extra_headers"
 27
 28from ....audio import Speaker
 29from ....utils import verboselogs
 30from ....options import DeepgramClientOptions
 31from .helpers import convert_to_websocket_url, append_query_params
 32from .errors import DeepgramError
 33
 34from .websocket_response import (
 35    OpenResponse,
 36    CloseResponse,
 37    ErrorResponse,
 38)
 39from .websocket_events import WebSocketEvents
 40
 41
 42ONE_SECOND = 1
 43HALF_SECOND = 0.5
 44DEEPGRAM_INTERVAL = 5
 45PING_INTERVAL = 20
 46
 47
 48class AbstractAsyncWebSocketClient(ABC):  # pylint: disable=too-many-instance-attributes
 49    """
 50    Abstract class for using WebSockets.
 51
 52    This class provides methods to establish a WebSocket connection generically for
 53    use in all WebSocket clients.
 54    """
 55
 56    _logger: verboselogs.VerboseLogger
 57    _config: DeepgramClientOptions
 58    _endpoint: str
 59    _websocket_url: str
 60
 61    _socket: Optional[ClientConnection] = None
 62
 63    _listen_thread: Union[asyncio.Task, None]
 64    _delegate: Optional[Speaker] = None
 65
 66    _kwargs: Optional[Dict] = None
 67    _addons: Optional[Dict] = None
 68    _options: Optional[Dict] = None
 69    _headers: Optional[Dict] = None
 70
 71    def __init__(self, config: DeepgramClientOptions, endpoint: str = ""):
 72        if config is None:
 73            raise DeepgramError("Config is required")
 74        if endpoint == "":
 75            raise DeepgramError("endpoint is required")
 76
 77        self._logger = verboselogs.VerboseLogger(__name__)
 78        self._logger.addHandler(logging.StreamHandler())
 79        self._logger.setLevel(config.verbose)
 80
 81        self._config = config
 82        self._endpoint = endpoint
 83
 84        self._listen_thread = None
 85
 86        # events
 87        self._exit_event = asyncio.Event()
 88
 89        # set websocket url
 90        self._websocket_url = convert_to_websocket_url(self._config.url, self._endpoint)
 91
 92    def delegate_listening(self, delegate: Speaker) -> None:
 93        """
 94        Delegate the listening thread to the Speaker object.
 95        """
 96        self._delegate = delegate
 97
 98    # pylint: disable=too-many-branches,too-many-statements
 99    async def start(
100        self,
101        options: Optional[Any] = None,
102        addons: Optional[Dict] = None,
103        headers: Optional[Dict] = None,
104        **kwargs,
105    ) -> bool:
106        """
107        Starts the WebSocket connection for live transcription.
108        """
109        self._logger.debug("AbstractAsyncWebSocketClient.start ENTER")
110        self._logger.info("addons: %s", addons)
111        self._logger.info("headers: %s", headers)
112        self._logger.info("kwargs: %s", kwargs)
113
114        self._addons = addons
115        self._headers = headers
116
117        # set kwargs
118        if kwargs is not None:
119            self._kwargs = kwargs
120        else:
121            self._kwargs = {}
122
123        if not isinstance(options, dict):
124            self._logger.error("options is not a dict")
125            self._logger.debug("AbstractSyncWebSocketClient.start LEAVE")
126            return False
127
128        # set options
129        if options is not None:
130            self._options = options
131        else:
132            self._options = {}
133
134        combined_options = self._options.copy()
135        if self._addons is not None:
136            self._logger.info("merging addons to options")
137            combined_options.update(self._addons)
138            self._logger.info("new options: %s", combined_options)
139        self._logger.debug("combined_options: %s", combined_options)
140
141        combined_headers = self._config.headers.copy()
142        if self._headers is not None:
143            self._logger.info("merging headers to options")
144            combined_headers.update(self._headers)
145            self._logger.info("new headers: %s", combined_headers)
146        self._logger.debug("combined_headers: %s", combined_headers)
147
148        url_with_params = append_query_params(self._websocket_url, combined_options)
149
150        try:
151            ws_connect_kwargs: Dict = {
152                "ping_interval": PING_INTERVAL,
153                WS_ADDITIONAL_HEADERS_KEY: combined_headers,
154            }
155
156            self._socket = await connect(
157                url_with_params,
158                **ws_connect_kwargs,
159            )
160            self._exit_event.clear()
161
162            # debug the threads
163            for thread in threading.enumerate():
164                self._logger.debug("after running thread: %s", thread.name)
165            self._logger.debug("number of active threads: %s", threading.active_count())
166
167            # delegate the listening thread to external object
168            if self._delegate is not None:
169                self._logger.notice("_delegate is enabled. this is usually the speaker")
170                self._delegate.set_pull_callback(self._socket.recv)
171                self._delegate.set_push_callback(self._process_message)
172            else:
173                self._logger.notice("create _listening thread")
174                self._listen_thread = asyncio.create_task(self._listening())
175
176            # debug the threads
177            for thread in threading.enumerate():
178                self._logger.debug("after running thread: %s", thread.name)
179            self._logger.debug("number of active threads: %s", threading.active_count())
180
181            # push open event
182            await self._emit(
183                WebSocketEvents(WebSocketEvents.Open),
184                OpenResponse(type=WebSocketEvents.Open),
185            )
186
187            self._logger.notice("start succeeded")
188            self._logger.debug("AbstractAsyncWebSocketClient.start LEAVE")
189            return True
190        except websockets.exceptions.ConnectionClosed as e:
191            self._logger.error(
192                "ConnectionClosed in AbstractAsyncWebSocketClient.start: %s", e
193            )
194            self._logger.debug("AbstractAsyncWebSocketClient.start LEAVE")
195            if self._config.options.get("termination_exception_connect", False):
196                raise
197            return False
198        except websockets.exceptions.WebSocketException as e:
199            self._logger.error(
200                "WebSocketException in AbstractAsyncWebSocketClient.start: %s", e
201            )
202            self._logger.debug("AbstractAsyncWebSocketClient.start LEAVE")
203            if self._config.options.get("termination_exception_connect", False):
204                raise
205            return False
206        except Exception as e:  # pylint: disable=broad-except
207            self._logger.error(
208                "WebSocketException in AbstractAsyncWebSocketClient.start: %s", e
209            )
210            self._logger.debug("AbstractAsyncWebSocketClient.start LEAVE")
211            if self._config.options.get("termination_exception_connect", False):
212                raise
213            return False
214
215    async def is_connected(self) -> bool:
216        """
217        Returns the connection status of the WebSocket.
218        """
219        return self._socket is not None
220
221    # pylint: enable=too-many-branches,too-many-statements
222
223    @abstractmethod
224    def on(self, event: WebSocketEvents, handler: Callable) -> None:
225        """
226        Registers an event handler for the WebSocket connection.
227        """
228        raise NotImplementedError("no on method")
229
230    @abstractmethod
231    async def _emit(self, event: WebSocketEvents, *args, **kwargs) -> None:
232        """
233        Emits an event to the WebSocket connection.
234        """
235        raise NotImplementedError("no _emit method")
236
237    # pylint: disable=too-many-return-statements,too-many-statements,too-many-locals,too-many-branches
238    async def _listening(self) -> None:
239        """
240        Listens for messages from the WebSocket connection.
241        """
242        self._logger.debug("AbstractAsyncWebSocketClient._listening ENTER")
243
244        while True:
245            try:
246                if self._exit_event.is_set():
247                    self._logger.notice("_listening exiting gracefully")
248                    self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
249                    return
250
251                if self._socket is None:
252                    self._logger.warning("socket is empty")
253                    self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
254                    return
255
256                message = await self._socket.recv()
257
258                if message is None:
259                    self._logger.info("message is None")
260                    continue
261
262                self._logger.spam("data type: %s", type(message))
263
264                if isinstance(message, bytes):
265                    self._logger.debug("Binary data received")
266                    await self._process_binary(message)
267                else:
268                    self._logger.debug("Text data received")
269                    await self._process_text(message)
270
271                self._logger.notice("_listening Succeeded")
272                self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
273
274            except websockets.exceptions.ConnectionClosedOK as e:
275                # signal exit and close
276                await self._signal_exit()
277
278                self._logger.notice(f"_listening({e.code}) exiting gracefully")
279                self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
280                return
281
282            except websockets.exceptions.ConnectionClosed as e:
283                if e.code in [1000, 1001]:
284                    # signal exit and close
285                    await self._signal_exit()
286
287                    self._logger.notice(f"_listening({e.code}) exiting gracefully")
288                    self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
289                    return
290
291                # we need to explicitly call self._signal_exit() here because we are hanging on a recv()
292                # note: this is different than the speak websocket client
293                self._logger.error(
294                    "ConnectionClosed in AbstractAsyncWebSocketClient._listening with code %s: %s",
295                    e.code,
296                    e.reason,
297                )
298                cc_error: ErrorResponse = ErrorResponse(
299                    "ConnectionClosed in AbstractAsyncWebSocketClient._listening",
300                    f"{e}",
301                    "ConnectionClosed",
302                )
303                await self._emit(
304                    WebSocketEvents(WebSocketEvents.Error),
305                    error=cc_error,
306                    **dict(cast(Dict[Any, Any], self._kwargs)),
307                )
308
309                # signal exit and close
310                await self._signal_exit()
311
312                self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
313
314                if self._config.options.get("termination_exception_connect") is True:
315                    raise
316                return
317
318            except websockets.exceptions.WebSocketException as e:
319                self._logger.error(
320                    "WebSocketException in AbstractAsyncWebSocketClient._listening: %s",
321                    e,
322                )
323                ws_error: ErrorResponse = ErrorResponse(
324                    "WebSocketException in AbstractAsyncWebSocketClient._listening",
325                    f"{e}",
326                    "WebSocketException",
327                )
328                await self._emit(
329                    WebSocketEvents(WebSocketEvents.Error),
330                    error=ws_error,
331                    **dict(cast(Dict[Any, Any], self._kwargs)),
332                )
333
334                # signal exit and close
335                await self._signal_exit()
336
337                self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
338
339                if self._config.options.get("termination_exception_connect") is True:
340                    raise
341                return
342
343            except Exception as e:  # pylint: disable=broad-except
344                self._logger.error(
345                    "Exception in AbstractAsyncWebSocketClient._listening: %s", e
346                )
347                e_error: ErrorResponse = ErrorResponse(
348                    "Exception in AbstractAsyncWebSocketClient._listening",
349                    f"{e}",
350                    "Exception",
351                )
352                await self._emit(
353                    WebSocketEvents(WebSocketEvents.Error),
354                    error=e_error,
355                    **dict(cast(Dict[Any, Any], self._kwargs)),
356                )
357
358                # signal exit and close
359                await self._signal_exit()
360
361                self._logger.debug("AbstractAsyncWebSocketClient._listening LEAVE")
362
363                if self._config.options.get("termination_exception_connect") is True:
364                    raise
365                return
366
367    # pylint: enable=too-many-return-statements,too-many-statements,too-many-locals,too-many-branches
368
369    async def _process_message(self, message: Union[str, bytes]) -> None:
370        if isinstance(message, bytes):
371            await self._process_binary(message)
372        else:
373            await self._process_text(message)
374
375    @abstractmethod
376    async def _process_text(self, message: str) -> None:
377        raise NotImplementedError("no _process_text method")
378
379    @abstractmethod
380    async def _process_binary(self, message: bytes) -> None:
381        raise NotImplementedError("no _process_binary method")
382
383    @abstractmethod
384    async def _close_message(self) -> bool:
385        raise NotImplementedError("no _close_message method")
386
387    # pylint: disable=too-many-return-statements,too-many-branches
388
389    async def send(self, data: Union[str, bytes]) -> bool:
390        """
391        Sends data over the WebSocket connection.
392        """
393        self._logger.spam("AbstractAsyncWebSocketClient.send ENTER")
394
395        if self._exit_event.is_set():
396            self._logger.notice("send exiting gracefully")
397            self._logger.debug("AbstractAsyncWebSocketClient.send LEAVE")
398            return False
399
400        if not await self.is_connected():
401            self._logger.notice("is_connected is False")
402            self._logger.debug("AbstractAsyncWebSocketClient.send LEAVE")
403            return False
404
405        if self._socket is not None:
406            try:
407                await self._socket.send(data)
408            except websockets.exceptions.ConnectionClosedOK as e:
409                self._logger.notice(f"send() exiting gracefully: {e.code}")
410                self._logger.debug("AbstractAsyncWebSocketClient.send LEAVE")
411                if self._config.options.get("termination_exception_send") is True:
412                    raise
413                return True
414            except websockets.exceptions.ConnectionClosed as e:
415                if e.code in [1000, 1001]:
416                    self._logger.notice(f"send({e.code}) exiting gracefully")
417                    self._logger.debug("AbstractAsyncWebSocketClient.send LEAVE")
418                    if self._config.options.get("termination_exception_send") is True:
419                        raise
420                    return True
421
422                self._logger.error("send() failed - ConnectionClosed: %s", str(e))
423                self._logger.spam("AbstractAsyncWebSocketClient.send LEAVE")
424                if self._config.options.get("termination_exception_send") is True:
425                    raise
426                return False
427            except websockets.exceptions.WebSocketException as e:
428                self._logger.error("send() failed - WebSocketException: %s", str(e))
429                self._logger.spam("AbstractAsyncWebSocketClient.send LEAVE")
430                if self._config.options.get("termination_exception_send") is True:
431                    raise
432                return False
433            except Exception as e:  # pylint: disable=broad-except
434                self._logger.error("send() failed - Exception: %s", str(e))
435                self._logger.spam("AbstractAsyncWebSocketClient.send LEAVE")
436                if self._config.options.get("termination_exception_send") is True:
437                    raise
438                return False
439
440            self._logger.spam("send() succeeded")
441            self._logger.spam("AbstractAsyncWebSocketClient.send LEAVE")
442            return True
443
444        self._logger.spam("send() failed. socket is None")
445        self._logger.spam("AbstractAsyncWebSocketClient.send LEAVE")
446        return False
447
448    # pylint: enable=too-many-return-statements,too-many-branches
449
450    async def finish(self) -> bool:
451        """
452        Closes the WebSocket connection gracefully.
453        """
454        self._logger.debug("AbstractAsyncWebSocketClient.finish ENTER")
455
456        # signal exit
457        await self._signal_exit()
458
459        # stop the threads
460        self._logger.verbose("cancelling tasks...")
461        try:
462            # Before cancelling, check if the tasks were created
463            # debug the threads
464            for thread in threading.enumerate():
465                self._logger.debug("before running thread: %s", thread.name)
466            self._logger.debug("number of active threads: %s", threading.active_count())
467
468            tasks = []
469            if self._listen_thread is not None:
470                self._listen_thread.cancel()
471                tasks.append(self._listen_thread)
472                self._logger.notice("processing _listen_thread cancel...")
473
474            # Use asyncio.gather to wait for tasks to be cancelled
475            await asyncio.gather(*filter(None, tasks))
476            self._logger.notice("threads joined")
477
478            # debug the threads
479            for thread in threading.enumerate():
480                if thread is not None and thread.name is not None:
481                    self._logger.debug("after running thread: %s", thread.name)
482                else:
483                    self._logger.debug("after running thread: unknown_thread_name")
484            self._logger.debug("number of active threads: %s", threading.active_count())
485
486            self._logger.notice("finish succeeded")
487            self._logger.spam("AbstractAsyncWebSocketClient.finish LEAVE")
488            return True
489
490        except asyncio.CancelledError as e:
491            self._logger.error("tasks cancelled error: %s", e)
492            self._logger.debug("AbstractAsyncWebSocketClient.finish LEAVE")
493            return True
494
495    async def _signal_exit(self) -> None:
496        # send close event
497        self._logger.verbose("closing socket...")
498        if self._socket is not None:
499            self._logger.verbose("send Close...")
500            try:
501                # if the socket connection is closed, the following line might throw an error
502                await self._close_message()
503            except websockets.exceptions.ConnectionClosedOK as e:
504                self._logger.notice("_signal_exit  - ConnectionClosedOK: %s", e.code)
505            except websockets.exceptions.ConnectionClosed as e:
506                self._logger.error("_signal_exit  - ConnectionClosed: %s", e.code)
507            except websockets.exceptions.WebSocketException as e:
508                self._logger.error("_signal_exit - WebSocketException: %s", str(e))
509            except Exception as e:  # pylint: disable=broad-except
510                self._logger.error("_signal_exit - Exception: %s", str(e))
511
512            # push close event
513            try:
514                await self._emit(
515                    WebSocketEvents(WebSocketEvents.Close),
516                    close=CloseResponse(type=WebSocketEvents.Close),
517                    **dict(cast(Dict[Any, Any], self._kwargs)),
518                )
519            except Exception as e:  # pylint: disable=broad-except
520                self._logger.error("_emit - Exception: %s", e)
521
522            # wait for task to send
523            await asyncio.sleep(0.5)
524
525        # signal exit
526        self._exit_event.set()
527
528        # closes the WebSocket connection gracefully
529        self._logger.verbose("clean up socket...")
530        if self._socket is not None:
531            self._logger.verbose("socket.wait_closed...")
532            try:
533                await self._socket.close()
534            except websockets.exceptions.WebSocketException as e:
535                self._logger.error("socket.wait_closed failed: %s", e)
536
537        self._socket = None