fix typecheck

playground/streaming/agent/websocket_user_implemented_agent.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from websockets.server import serve
+from websockets.asyncio.server import serve
 
 from vocode.streaming.models.websocket_agent import (
     WebSocketAgentMessage,

vocode/streaming/agent/websocket_user_implemented_agent.py

@@ -3,7 +3,7 @@
 
 import websockets
 from loguru import logger
-from websockets.asyncio.client import WebSocketClientProtocol
+from websockets.asyncio.client import ClientConnection
 
 from vocode.streaming.agent.base_agent import (
     AgentInput,
@@ -76,7 +76,7 @@ async def _process(self) -> None:
         async with websockets.connect(socket_url) as ws:
 
             async def sender(
-                ws: WebSocketClientProtocol,
+                ws: ClientConnection,
             ) -> None:  # sends audio to websocket
                 while not self.has_ended:
                     logger.info("Waiting for data from agent request queue")
@@ -109,7 +109,7 @@ async def sender(
 
                 logger.debug("Terminating web socket agent sender")
 
-            async def receiver(ws: WebSocketClientProtocol) -> None:
+            async def receiver(ws: ClientConnection) -> None:
                 while not self.has_ended:
                     try:
                         msg = await ws.recv()

vocode/streaming/transcriber/deepgram_transcriber.py

@@ -8,7 +8,7 @@
 import websockets
 from loguru import logger
 from pydantic.v1 import BaseModel, Field
-from websockets.asyncio.client import WebSocketClientProtocol
+from websockets.asyncio.client import ClientConnection
 
 from vocode import getenv
 from vocode.streaming.models.audio import AudioEncoding
@@ -394,11 +394,13 @@ async def process(self):
         logger.info(f"Connecting to Deepgram at {deepgram_url}")
 
         try:
-            async with websockets.connect(deepgram_url, additional_headers=additional_headers) as ws:
+            async with websockets.connect(
+                deepgram_url, additional_headers=additional_headers
+            ) as ws:
                 self.connected_ts = now()
 
                 async def sender(
-                    ws: WebSocketClientProtocol,
+                    ws: ClientConnection,
                 ):  # sends audio to websocket
                     byte_rate = self.get_byte_rate()
 
@@ -417,7 +419,7 @@ async def sender(
 
                     logger.debug("Terminating Deepgram transcriber sender")
 
-                async def receiver(ws: WebSocketClientProtocol):
+                async def receiver(ws: ClientConnection):
                     buffer = ""
                     buffer_avg_confidence = 0.0
                     num_buffer_utterances = 1

vocode/streaming/transcriber/rev_ai_transcriber.py

@@ -5,7 +5,7 @@
 
 import websockets
 from loguru import logger
-from websockets.asyncio.client import WebSocketClientProtocol
+from websockets.asyncio.client import ClientConnection
 
 from vocode import getenv
 from vocode.streaming.models.transcriber import (
@@ -71,7 +71,7 @@ async def _run_loop(self):
     async def process(self):
         async with websockets.connect(self.get_rev_ai_url()) as ws:
 
-            async def sender(ws: WebSocketClientProtocol):
+            async def sender(ws: ClientConnection):
                 while not self.closed:
                     try:
                         data = await asyncio.wait_for(self._input_queue.get(), 5)
@@ -81,7 +81,7 @@ async def sender(ws: WebSocketClientProtocol):
                 await ws.close()
                 logger.debug("Terminating Rev.AI transcriber sender")
 
-            async def receiver(ws: WebSocketClientProtocol):
+            async def receiver(ws: ClientConnection):
                 buffer = ""
 
                 while not self.closed: