Fix support for Python WebSocket >=14 while preserving backward compatibility (#499)

* Fix support for Python WebSocket >=14 while preserving backward compatibility

Python websockets did a breaking change with version 14.
"extra_headers" was renamed to "additionnal_headers".

* Unpin websockets dependency to support at least versions 12, 13 and 14

Author: fviard

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 (  # type: ignore
+        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
             )

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
             )

pyproject.toml

@@ -31,7 +31,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.10"
 httpx = "^0.25.2"
-websockets = "^12.0"
+websockets = ">=12.0"
 typing-extensions = "^4.9.0"
 dataclasses-json = "^0.6.3"
 aiohttp = "^3.9.1"

requirements.txt

@@ -1,7 +1,7 @@
 # pip install -r requirements.txt
 
 # standard libs
-websockets==12.*
+websockets>=12.0
 httpx==0.*
 dataclasses-json==0.*
 dataclasses==0.*

setup.py

@@ -28,7 +28,7 @@
     packages=find_packages(exclude=["tests"]),
     install_requires=[
         "httpx>=0.25.2",
-        "websockets>=12.0,<14.0",
+        "websockets>=12.0",
         "dataclasses-json>=0.6.3",
         "typing_extensions>=4.9.0",
         "aiohttp>=3.9.1",