|
| 1 | +# --- |
| 2 | +# jupyter: |
| 3 | +# jupytext: |
| 4 | +# cell_metadata_filter: -all |
| 5 | +# text_representation: |
| 6 | +# extension: .py |
| 7 | +# format_name: percent |
| 8 | +# format_version: '1.3' |
| 9 | +# jupytext_version: 1.18.1 |
| 10 | +# --- |
| 11 | + |
| 12 | +# %% [markdown] |
| 13 | +# # Barge-In Attack (Streaming Audio) |
| 14 | +# |
| 15 | +# `BargeInAttack` streams user audio to a `RealtimeTarget` and uses server-side voice-activity |
| 16 | +# detection (VAD) to detect turn boundaries. When the user speaks while the assistant is still |
| 17 | +# responding, server VAD cancels the in-flight response (barge-in). Interrupted turns are |
| 18 | +# persisted with `prompt_metadata["interrupted"] = True`. |
| 19 | +# |
| 20 | +# Audio converters are applied per turn after VAD commits. The raw audio drives interruption |
| 21 | +# timing while the model responds to the converted version. |
| 22 | +# |
| 23 | +# > **Note:** Memory must be initialized via `initialize_pyrit_async`. See the |
| 24 | +# > [Memory Configuration Guide](../../memory/0_memory.md). |
| 25 | + |
| 26 | +# %% [markdown] |
| 27 | +# ## Setup |
| 28 | +# |
| 29 | +# `BargeInAttack` requires a `RealtimeTarget` with `server_vad=True` (or a `ServerVadConfig` |
| 30 | +# for custom tuning). |
| 31 | + |
| 32 | +# %% |
| 33 | +import asyncio |
| 34 | +import wave |
| 35 | +from pathlib import Path |
| 36 | + |
| 37 | +from pyrit.executor.attack import ( |
| 38 | + AttackConverterConfig, |
| 39 | + BargeInAttack, |
| 40 | + BargeInAttackContext, |
| 41 | + ConsoleAttackResultPrinter, |
| 42 | +) |
| 43 | +from pyrit.executor.attack.core import AttackParameters |
| 44 | +from pyrit.memory import CentralMemory |
| 45 | +from pyrit.prompt_converter import AudioFrequencyConverter |
| 46 | +from pyrit.prompt_normalizer import PromptConverterConfiguration |
| 47 | +from pyrit.prompt_target import RealtimeTarget |
| 48 | +from pyrit.setup import IN_MEMORY, initialize_pyrit_async |
| 49 | + |
| 50 | +await initialize_pyrit_async(memory_db_type=IN_MEMORY) # type: ignore |
| 51 | + |
| 52 | +# %% [markdown] |
| 53 | +# ## Shared setup |
| 54 | +# |
| 55 | +# Both sections use a pre-recorded 24 kHz mono PCM16 question about photosynthesis. The |
| 56 | +# format matches what the OpenAI Realtime API expects. Any async generator yielding 24 kHz |
| 57 | +# PCM16 bytes works as a chunk source (live mic, TTS, etc.). |
| 58 | + |
| 59 | +# %% |
| 60 | +CHUNK_MS = 100 |
| 61 | +CHUNK_SIZE = CHUNK_MS * 48 # PCM16 @ 24 kHz mono = 48 bytes per millisecond. |
| 62 | +SILENCE_CHUNK = b"\x00" * CHUNK_SIZE |
| 63 | +audio_path = Path("../../../../assets/photosynthesis_question.wav").resolve() |
| 64 | + |
| 65 | + |
| 66 | +def _load_pcm(path: Path) -> bytes: |
| 67 | + """Read a WAV at 24 kHz / mono / PCM16 into raw PCM bytes.""" |
| 68 | + with wave.open(str(path), "rb") as wav: |
| 69 | + assert wav.getframerate() == 24000 and wav.getnchannels() == 1 and wav.getsampwidth() == 2 |
| 70 | + return wav.readframes(wav.getnframes()) |
| 71 | + |
| 72 | + |
| 73 | +async def _yield_chunks(pcm: bytes, real_time: bool = True): |
| 74 | + """Yield PCM in 100ms slices, optionally pacing at real-time.""" |
| 75 | + for offset in range(0, len(pcm), CHUNK_SIZE): |
| 76 | + yield pcm[offset : offset + CHUNK_SIZE] |
| 77 | + if real_time: |
| 78 | + await asyncio.sleep(CHUNK_MS / 1000) |
| 79 | + |
| 80 | + |
| 81 | +question_pcm_24k = _load_pcm(audio_path) |
| 82 | +print(f"Loaded question: {len(question_pcm_24k) / 48 / 1000:.2f}s @ 24 kHz") |
| 83 | + |
| 84 | +converters = PromptConverterConfiguration.from_converters(converters=[AudioFrequencyConverter(shift_value=200)]) |
| 85 | + |
| 86 | + |
| 87 | +# %% [markdown] |
| 88 | +# ## Section 1: Single-turn streaming with a converter |
| 89 | +# |
| 90 | +# Streams one user statement, applies a frequency-shift converter after VAD commits the turn, |
| 91 | +# and gets the model's response. Exercises the full pipeline (chunk push, convert-on-commit, |
| 92 | +# item swap, response trigger, memory persistence) without barge-in. |
| 93 | + |
| 94 | + |
| 95 | +# %% |
| 96 | +async def single_turn_source(): |
| 97 | + async for chunk in _yield_chunks(question_pcm_24k): |
| 98 | + yield chunk |
| 99 | + # Trailing silence helps server VAD recognize end-of-turn. |
| 100 | + for _ in range(25): # 2.5s trailing silence, above the 1.5s VAD threshold |
| 101 | + yield SILENCE_CHUNK |
| 102 | + await asyncio.sleep(CHUNK_MS / 1000) |
| 103 | + |
| 104 | + |
| 105 | +target = RealtimeTarget() |
| 106 | +attack = BargeInAttack( |
| 107 | + objective_target=target, |
| 108 | + attack_converter_config=AttackConverterConfig(request_converters=converters), |
| 109 | +) |
| 110 | + |
| 111 | +context = BargeInAttackContext( |
| 112 | + params=AttackParameters(objective="Observe a single converted user turn end-to-end"), |
| 113 | + audio_chunks=single_turn_source(), |
| 114 | +) |
| 115 | + |
| 116 | +result = await attack.execute_with_context_async(context=context) # type: ignore |
| 117 | +print(f"executed_turns: {result.executed_turns}") |
| 118 | +await ConsoleAttackResultPrinter(width=200).write_async(result=result) # type: ignore |
| 119 | +await target.cleanup_target_async() # type: ignore |
| 120 | + |
| 121 | +# %% [markdown] |
| 122 | +# ## Section 2: Barge-in (interrupting the assistant mid-response) |
| 123 | +# |
| 124 | +# Plays the question twice with timing arranged so turn 2's speech arrives during turn 1's |
| 125 | +# response. Server VAD detects the new speech, cancels turn 1's response, and resolves it |
| 126 | +# with `interrupted=True`. |
| 127 | + |
| 128 | +# %% [markdown] |
| 129 | +# ### Reading the barge-in output |
| 130 | +# |
| 131 | +# After running the next cell, if barge-in fired successfully: |
| 132 | +# - `executed_turns: 2` (two VAD-detected user turns) |
| 133 | +# - First assistant turn shows `[INTERRUPTED]` with a truncated transcript |
| 134 | +# - Second assistant turn completes normally |
| 135 | +# |
| 136 | +# If you don't see `[INTERRUPTED]`, decrease `TURN1_RESPONSE_WAIT_S` so turn 2's audio |
| 137 | +# arrives earlier in turn 1's response window. |
| 138 | + |
| 139 | +# %% |
| 140 | +TURN1_RESPONSE_WAIT_S = 0.2 # how long to let the model start speaking before barging in |
| 141 | + |
| 142 | + |
| 143 | +async def barge_in_source(): |
| 144 | + # Turn 1: speak the question, then 1.5s of silence so VAD commits. |
| 145 | + async for chunk in _yield_chunks(question_pcm_24k): |
| 146 | + yield chunk |
| 147 | + for _ in range(25): # 2.5s trailing silence |
| 148 | + yield SILENCE_CHUNK |
| 149 | + await asyncio.sleep(CHUNK_MS / 1000) |
| 150 | + |
| 151 | + # Let the model get partway into its response before we interrupt. |
| 152 | + for _ in range(int(TURN1_RESPONSE_WAIT_S * 10)): |
| 153 | + yield SILENCE_CHUNK |
| 154 | + await asyncio.sleep(CHUNK_MS / 1000) |
| 155 | + |
| 156 | + # Turn 2: speak the question again. VAD's speech_started fires while turn 1's response |
| 157 | + # is still streaming → server cancels + truncates turn 1. |
| 158 | + async for chunk in _yield_chunks(question_pcm_24k): |
| 159 | + yield chunk |
| 160 | + for _ in range(25): # 2.5s trailing silence |
| 161 | + yield SILENCE_CHUNK |
| 162 | + await asyncio.sleep(CHUNK_MS / 1000) |
| 163 | + |
| 164 | + |
| 165 | +target2 = RealtimeTarget() |
| 166 | +attack2 = BargeInAttack( |
| 167 | + objective_target=target2, |
| 168 | + attack_converter_config=AttackConverterConfig(request_converters=converters), |
| 169 | +) |
| 170 | + |
| 171 | +barge_in_context = BargeInAttackContext( |
| 172 | + params=AttackParameters(objective="Demonstrate barge-in by interrupting a benign answer"), |
| 173 | + audio_chunks=barge_in_source(), |
| 174 | +) |
| 175 | + |
| 176 | +barge_in_result = await attack2.execute_with_context_async(context=barge_in_context) # type: ignore |
| 177 | +print(f"executed_turns: {barge_in_result.executed_turns}") |
| 178 | + |
| 179 | +# Inspect memory to verify the barge-in landed in metadata. |
| 180 | +memory = CentralMemory.get_memory_instance() |
| 181 | +turns = memory.get_conversation(conversation_id=barge_in_result.conversation_id) |
| 182 | +print(f"\nPersisted pieces ({len(turns)} messages):") |
| 183 | +for message in turns: |
| 184 | + for piece in message.message_pieces: |
| 185 | + interrupted = piece.prompt_metadata.get("interrupted") |
| 186 | + marker = " [INTERRUPTED]" if interrupted else "" |
| 187 | + val = piece.converted_value |
| 188 | + if piece.converted_value_data_type == "audio_path": |
| 189 | + val = Path(val).name |
| 190 | + value_preview = (val[:80] + "...") if len(val) > 80 else val |
| 191 | + print(f" {piece._role} {piece.converted_value_data_type}{marker}: {value_preview}") |
| 192 | + |
| 193 | +await ConsoleAttackResultPrinter(width=200).write_async(result=barge_in_result) # type: ignore |
| 194 | +await target2.cleanup_target_async() # type: ignore |
| 195 | + |
| 196 | +# %% [markdown] |
| 197 | +# ## Alternate chunk sources |
| 198 | +# |
| 199 | +# The chunk source is the main strategy hook: |
| 200 | +# |
| 201 | +# - **Pre-recorded WAV** (this notebook): most common starting point |
| 202 | +# - **TTS converter**: generate audio from text prompts dynamically |
| 203 | +# - **Live microphone**: use `sounddevice` or similar; yield what the mic produces |
| 204 | +# |
| 205 | +# For feedback-driven attacks — for example, scoring each assistant turn and choosing |
| 206 | +# to barge in with follow-up audio only when the response shows incomplete refusal — |
| 207 | +# subclass `BargeInAttack` and override `_perform_async` to interleave turn observation |
| 208 | +# with chunk generation. |
0 commit comments