File tree Expand file tree Collapse file tree 3 files changed +15
-7
lines changed Expand file tree Collapse file tree 3 files changed +15
-7
lines changed Original file line number Diff line number Diff line change @@ -77,12 +77,13 @@ class StreamedAudioInput:
7777 """
7878
7979 def __init__ (self ):
80- self .queue : asyncio .Queue [npt .NDArray [np .int16 | np .float32 ]] = asyncio .Queue ()
80+ self .queue : asyncio .Queue [npt .NDArray [np .int16 | np .float32 ] | None ] = asyncio .Queue ()
8181
82- async def add_audio (self , audio : npt .NDArray [np .int16 | np .float32 ]):
82+ async def add_audio (self , audio : npt .NDArray [np .int16 | np .float32 ] | None ):
8383 """Adds more audio data to the stream.
8484
8585 Args:
86- audio: The audio data to add. Must be a numpy array of int16 or float32.
86+ audio: The audio data to add. Must be a numpy array of int16 or float32 or None.
87+ If None passed, it indicates the end of the stream.
8788 """
8889 await self .queue .put (audio )
Original file line number Diff line number Diff line change @@ -88,7 +88,7 @@ def __init__(
8888 self ._trace_include_sensitive_data = trace_include_sensitive_data
8989 self ._trace_include_sensitive_audio_data = trace_include_sensitive_audio_data
9090
91- self ._input_queue : asyncio .Queue [npt .NDArray [np .int16 | np .float32 ]] = input .queue
91+ self ._input_queue : asyncio .Queue [npt .NDArray [np .int16 | np .float32 ] | None ] = input .queue
9292 self ._output_queue : asyncio .Queue [str | ErrorSentinel | SessionCompleteSentinel ] = (
9393 asyncio .Queue ()
9494 )
@@ -245,7 +245,7 @@ async def _handle_events(self) -> None:
245245 await self ._output_queue .put (SessionCompleteSentinel ())
246246
247247 async def _stream_audio (
248- self , audio_queue : asyncio .Queue [npt .NDArray [np .int16 | np .float32 ]]
248+ self , audio_queue : asyncio .Queue [npt .NDArray [np .int16 | np .float32 ] | None ]
249249 ) -> None :
250250 assert self ._websocket is not None , "Websocket not initialized"
251251 self ._start_turn ()
Original file line number Diff line number Diff line change @@ -121,7 +121,14 @@ async def test_streamed_audio_input(self):
121121 # Verify the queue contents
122122 assert streamed_input .queue .qsize () == 2
123123 # Test non-blocking get
124- assert np .array_equal (streamed_input .queue .get_nowait (), audio1 )
124+ retrieved_audio1 = streamed_input .queue .get_nowait ()
125+ # Satisfy type checker
126+ assert retrieved_audio1 is not None
127+ assert np .array_equal (retrieved_audio1 , audio1 )
128+
125129 # Test blocking get
126- assert np .array_equal (await streamed_input .queue .get (), audio2 )
130+ retrieved_audio2 = await streamed_input .queue .get ()
131+ # Satisfy type checker
132+ assert retrieved_audio2 is not None
133+ assert np .array_equal (retrieved_audio2 , audio2 )
127134 assert streamed_input .queue .empty ()
You can’t perform that action at this time.
0 commit comments