fix websocket clients

Author: twitchard

src/hume/empathic_voice/chat/socket_client.py

@@ -10,6 +10,11 @@
 import websockets.protocol
 from json.decoder import JSONDecodeError
 
+try:
+    from websockets.legacy.client import connect as websockets_client_connect  # type: ignore
+except ImportError:
+    from websockets import connect as websockets_client_connect  # type: ignore
+
 from hume.core.websocket import (
     OnErrorHandlerType,
     OnMessageHandlerType,
@@ -356,7 +361,7 @@ async def connect(
         ws_uri = await self._construct_ws_uri(options)
 
         try:
-            async with websockets.connect(
+            async with websockets_client_connect(
                 ws_uri,
                 extra_headers=exclude_auth_headers(self.client_wrapper.get_headers()),
                 max_size=self.DEFAULT_MAX_PAYLOAD_SIZE_BYTES,
@@ -448,7 +453,7 @@ async def connect_with_callbacks(
         background_task: typing.Optional[asyncio.Task[None]] = None
 
         try:
-            async with websockets.connect(
+            async with websockets_client_connect(
                 ws_uri,
                 extra_headers=exclude_auth_headers(self.client_wrapper.get_headers()),
                 max_size=self.DEFAULT_MAX_PAYLOAD_SIZE_BYTES,

src/hume/expression_measurement/stream/socket_client.py

@@ -9,6 +9,11 @@
 import websockets
 import websockets.protocol
 
+try:
+    from websockets.legacy.client import connect as websockets_client_connect  # type: ignore
+except ImportError:
+    from websockets import connect as websockets_client_connect  # type: ignore
+
 from hume.core.api_error import ApiError
 
 from .stream.types.config import Config
@@ -216,7 +221,7 @@ async def connect(
 
         base = self.client_wrapper.get_environment().base.replace('https://', 'wss://').replace('http://', 'ws://')
         try:
-            async with websockets.connect(  # type: ignore[attr-defined]
+            async with websockets_client_connect(
                 f"{base}/v0/stream/models",
                 extra_headers={
                     **exclude_auth_headers(self.client_wrapper.get_headers()),