Fix support for Python WebSocket >=14 while preserving backward compa…

…tibility Python websockets did a breaking change with version 14. "extra_headers" was renamed to "additionnal_headers".
deepgram · Jan 28, 2025 · 6f8e3e2 · 6f8e3e2
1 parent ad276e1
commit 6f8e3e2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 7 deletions.
diff --git a/deepgram/clients/common/v1/abstract_async_websocket.py b/deepgram/clients/common/v1/abstract_async_websocket.py
@@ -10,7 +10,20 @@
 from abc import ABC, abstractmethod
 
 import websockets
-from websockets.client import WebSocketClientProtocol
+
+try:
+    # Websockets versions >= 13
+    from websockets.asyncio.client import connect, ClientConnection
+
+    WS_ADDITIONAL_HEADERS_KEY = "additional_headers"
+except ImportError:
+    # Backward compatibility with websockets versions 12
+    from websockets.legacy.client import (
+        connect,
+        WebSocketClientProtocol as ClientConnection,
+    )
+
+    WS_ADDITIONAL_HEADERS_KEY = "extra_headers"
 
 from ....audio import Speaker
 from ....utils import verboselogs
@@ -25,6 +38,7 @@
 )
 from .websocket_events import WebSocketEvents
 
+
 ONE_SECOND = 1
 HALF_SECOND = 0.5
 DEEPGRAM_INTERVAL = 5
@@ -44,7 +58,7 @@ class AbstractAsyncWebSocketClient(ABC):  # pylint: disable=too-many-instance-at
     _endpoint: str
     _websocket_url: str
 
-    _socket: Optional[WebSocketClientProtocol] = None
+    _socket: Optional[ClientConnection] = None
 
     _listen_thread: Union[asyncio.Task, None]
     _delegate: Optional[Speaker] = None
@@ -134,10 +148,14 @@ async def start(
         url_with_params = append_query_params(self._websocket_url, combined_options)
 
         try:
-            self._socket = await websockets.connect(
+            ws_connect_kwargs: Dict = {
+                "ping_interval": PING_INTERVAL,
+                WS_ADDITIONAL_HEADERS_KEY: combined_headers,
+            }
+
+            self._socket = await connect(
                 url_with_params,
-                extra_headers=combined_headers,
-                ping_interval=PING_INTERVAL,
+                **ws_connect_kwargs,
             )
             self._exit_event.clear()
 
@@ -169,7 +187,7 @@ async def start(
             self._logger.notice("start succeeded")
             self._logger.debug("AbstractAsyncWebSocketClient.start LEAVE")
             return True
-        except websockets.ConnectionClosed as e:
+        except websockets.exceptions.ConnectionClosed as e:
             self._logger.error(
                 "ConnectionClosed in AbstractAsyncWebSocketClient.start: %s", e
             )

diff --git a/deepgram/clients/common/v1/abstract_sync_websocket.py b/deepgram/clients/common/v1/abstract_sync_websocket.py
@@ -184,7 +184,7 @@ def start(
             self._logger.notice("start succeeded")
             self._logger.debug("AbstractSyncWebSocketClient.start LEAVE")
             return True
-        except websockets.ConnectionClosed as e:
+        except websockets.exceptions.ConnectionClosed as e:
             self._logger.error(
                 "ConnectionClosed in AbstractSyncWebSocketClient.start: %s", e
             )