Skip to content

Commit f24921d

Browse files
romanlutzCopilot
andauthored
FIX Avoid blocking WAV reads in realtime target (#2305)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1c85c49c-2d6c-4050-b3e3-ce085c97c5ec
1 parent 660be89 commit f24921d

2 files changed

Lines changed: 51 additions & 8 deletions

File tree

pyrit/prompt_target/openai/openai_realtime_target.py

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -825,14 +825,7 @@ async def send_audio_async(
825825
"""
826826
connection = self._get_connection(conversation_id=conversation_id)
827827

828-
with wave.open(filename, "rb") as wav_file:
829-
# Read WAV parameters
830-
num_channels = wav_file.getnchannels()
831-
sample_width = wav_file.getsampwidth() # Should be 2 bytes for PCM16
832-
frame_rate = wav_file.getframerate()
833-
num_frames = wav_file.getnframes()
834-
835-
audio_content = wav_file.readframes(num_frames)
828+
audio_content, num_channels, sample_width, frame_rate = await asyncio.to_thread(self._read_wav_file, filename)
836829

837830
receive_tasks = asyncio.create_task(self.receive_events_async(conversation_id=conversation_id))
838831

@@ -871,3 +864,23 @@ async def _construct_message_from_response_async(self, response: Any, request: A
871864
This implementation exists to satisfy the abstract base class requirement.
872865
"""
873866
raise NotImplementedError("RealtimeTarget uses receive_events for message construction")
867+
868+
@staticmethod
869+
def _read_wav_file(filename: str) -> tuple[bytes, int, int, int]:
870+
"""
871+
Read raw audio frames and format metadata from a WAV file.
872+
873+
Args:
874+
filename (str): Path to the WAV file to read.
875+
876+
Returns:
877+
tuple[bytes, int, int, int]: The raw audio frames, number of channels,
878+
sample width in bytes, and frame rate.
879+
"""
880+
with wave.open(filename, "rb") as wav_file:
881+
return (
882+
wav_file.readframes(wav_file.getnframes()),
883+
wav_file.getnchannels(),
884+
wav_file.getsampwidth(),
885+
wav_file.getframerate(),
886+
)

tests/unit/prompt_target/target/test_realtime_target.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1216,6 +1216,36 @@ def _write_wav(
12161216
return str(path)
12171217

12181218

1219+
async def test_send_audio_async_reads_wav_off_event_loop(target, tmp_path):
1220+
connection = AsyncMock()
1221+
target._existing_conversation["conv"] = connection
1222+
target.receive_events_async = AsyncMock(
1223+
return_value=RealtimeTargetResult(audio_bytes=b"response", transcripts=["transcript"])
1224+
)
1225+
target.send_response_create_async = AsyncMock()
1226+
target.save_audio_async = AsyncMock(return_value="output.wav")
1227+
1228+
pcm = b"\x01\x02" * 8
1229+
wav_path = _write_wav(tmp_path / "input.wav", pcm=pcm)
1230+
with patch(
1231+
"pyrit.prompt_target.openai.openai_realtime_target.asyncio.to_thread",
1232+
new_callable=AsyncMock,
1233+
wraps=asyncio.to_thread,
1234+
) as to_thread_mock:
1235+
output_path, _ = await target.send_audio_async(filename=wav_path, conversation_id="conv")
1236+
1237+
assert output_path == "output.wav"
1238+
assert to_thread_mock.await_args.args[1] == wav_path
1239+
connection.conversation.item.create.assert_awaited_once_with(
1240+
item={
1241+
"type": "message",
1242+
"role": "user",
1243+
"content": [{"type": "input_audio", "audio": base64.b64encode(pcm).decode("utf-8")}],
1244+
}
1245+
)
1246+
target.save_audio_async.assert_awaited_once_with(b"response", 1, 2, 24000)
1247+
1248+
12191249
async def test_send_prompt_audio_path_calls_send_audio_async(target, tmp_path):
12201250
"""An audio_path message is routed through the atomic send_audio_async path."""
12211251
wav_path = _write_wav(tmp_path / "in.wav")

0 commit comments

Comments
 (0)