"""
Amazon Nova Sonic via TrueFoundry Gateway (HTTP/2 bidirectional EventStream)
pip install aws_sdk_bedrock_runtime smithy-aws-core pyaudio
"""
import asyncio
import base64
import json
import os
import uuid
import pyaudio
from aws_sdk_bedrock_runtime.client import (
BedrockRuntimeClient,
InvokeModelWithBidirectionalStreamOperationInput,
)
from aws_sdk_bedrock_runtime.config import Config
from aws_sdk_bedrock_runtime.models import (
BidirectionalInputPayloadPart,
InvokeModelWithBidirectionalStreamInputChunk,
)
API_KEY = "your-truefoundry-api-key"
GATEWAY_BASE = "https://{GATEWAY_STREAMING_HOST}" # e.g. https://gateway.example.com:8443
PROVIDER_ACCOUNT = "tfy-ai-bedrock" # Bedrock account name on TrueFoundry
MODEL_ID = "amazon.nova-2-sonic-v1:0"
REGION = "us-east-1"
os.environ["AWS_BEARER_TOKEN_BEDROCK"] = API_KEY
INPUT_SAMPLE_RATE = 16000
OUTPUT_SAMPLE_RATE = 24000
CHANNELS = 1
FORMAT = pyaudio.paInt16
CHUNK_SIZE = 1024
class GatewayNovaSonic:
def __init__(self):
self.client = None
self.stream = None
self.response_task = None
self.is_active = False
self.prompt_name = str(uuid.uuid4())
self.content_name = str(uuid.uuid4())
self.audio_content_name = str(uuid.uuid4())
self.audio_queue: asyncio.Queue[bytes] = asyncio.Queue()
self.role = None
self.display_assistant_text = False
def _initialize_client(self):
# SDK appends /model/{model_id}/invoke-with-bidirectional-stream
config = Config(
endpoint_uri=f"{GATEWAY_BASE.rstrip('/')}/bedrock-stream/{PROVIDER_ACCOUNT}",
region=REGION,
)
self.client = BedrockRuntimeClient(config=config)
async def send_event(self, event_json: str):
event = InvokeModelWithBidirectionalStreamInputChunk(
value=BidirectionalInputPayloadPart(bytes_=event_json.encode("utf-8"))
)
await self.stream.input_stream.send(event)
async def start_session(self):
if not self.client:
self._initialize_client()
self.stream = await self.client.invoke_model_with_bidirectional_stream(
InvokeModelWithBidirectionalStreamOperationInput(model_id=MODEL_ID)
)
self.is_active = True
await self.send_event(
json.dumps(
{
"event": {
"sessionStart": {
"inferenceConfiguration": {
"maxTokens": 1024,
"topP": 0.9,
"temperature": 0.7,
}
}
}
}
)
)
await self.send_event(
json.dumps(
{
"event": {
"promptStart": {
"promptName": self.prompt_name,
"textOutputConfiguration": {"mediaType": "text/plain"},
"audioOutputConfiguration": {
"mediaType": "audio/lpcm",
"sampleRateHertz": OUTPUT_SAMPLE_RATE,
"sampleSizeBits": 16,
"channelCount": 1,
"voiceId": "matthew",
"encoding": "base64",
"audioType": "SPEECH",
},
}
}
}
)
)
await self.send_event(
json.dumps(
{
"event": {
"contentStart": {
"promptName": self.prompt_name,
"contentName": self.content_name,
"type": "TEXT",
"interactive": False,
"role": "SYSTEM",
"textInputConfiguration": {"mediaType": "text/plain"},
}
}
}
)
)
await self.send_event(
json.dumps(
{
"event": {
"textInput": {
"promptName": self.prompt_name,
"contentName": self.content_name,
"content": "You are a friendly assistant. Keep spoken replies short.",
}
}
}
)
)
await self.send_event(
json.dumps(
{
"event": {
"contentEnd": {
"promptName": self.prompt_name,
"contentName": self.content_name,
}
}
}
)
)
self.response_task = asyncio.create_task(self._process_responses())
print("Connected to TrueFoundry Nova Sonic streaming endpoint")
async def start_audio_input(self):
await self.send_event(
json.dumps(
{
"event": {
"contentStart": {
"promptName": self.prompt_name,
"contentName": self.audio_content_name,
"type": "AUDIO",
"interactive": True,
"role": "USER",
"audioInputConfiguration": {
"mediaType": "audio/lpcm",
"sampleRateHertz": INPUT_SAMPLE_RATE,
"sampleSizeBits": 16,
"channelCount": 1,
"audioType": "SPEECH",
"encoding": "base64",
},
}
}
}
)
)
async def send_audio_chunk(self, audio_bytes: bytes):
if not self.is_active:
return
await self.send_event(
json.dumps(
{
"event": {
"audioInput": {
"promptName": self.prompt_name,
"contentName": self.audio_content_name,
"content": base64.b64encode(audio_bytes).decode("utf-8"),
}
}
}
)
)
async def end_audio_input(self):
await self.send_event(
json.dumps(
{
"event": {
"contentEnd": {
"promptName": self.prompt_name,
"contentName": self.audio_content_name,
}
}
}
)
)
async def end_session(self):
if not self.is_active:
return
await self.send_event(json.dumps({"event": {"promptEnd": {"promptName": self.prompt_name}}}))
await self.send_event(json.dumps({"event": {"sessionEnd": {}}}))
await self.stream.input_stream.close()
async def _process_responses(self):
try:
while self.is_active:
output = await self.stream.await_output()
result = await output[1].receive()
if not (result.value and result.value.bytes_):
continue
data = json.loads(result.value.bytes_.decode("utf-8"))
event = data.get("event") or {}
if "contentStart" in event:
self.role = event["contentStart"].get("role")
fields = event["contentStart"].get("additionalModelFields")
if fields:
stage = json.loads(fields).get("generationStage")
self.display_assistant_text = stage == "SPECULATIVE"
else:
self.display_assistant_text = False
elif "textOutput" in event:
text = event["textOutput"]["content"]
if self.role == "ASSISTANT" and self.display_assistant_text:
print(f"Assistant: {text}")
elif self.role == "USER":
print(f"User: {text}")
elif "audioOutput" in event:
await self.audio_queue.put(base64.b64decode(event["audioOutput"]["content"]))
except Exception as exc:
print(f"Error processing responses: {exc}")
async def play_audio(self):
p = pyaudio.PyAudio()
stream = p.open(format=FORMAT, channels=CHANNELS, rate=OUTPUT_SAMPLE_RATE, output=True)
try:
while self.is_active:
stream.write(await self.audio_queue.get())
finally:
stream.stop_stream()
stream.close()
p.terminate()
async def capture_audio(self):
p = pyaudio.PyAudio()
stream = p.open(
format=FORMAT,
channels=CHANNELS,
rate=INPUT_SAMPLE_RATE,
input=True,
frames_per_buffer=CHUNK_SIZE,
)
print("Speak into your microphone. Press Enter to stop...")
await self.start_audio_input()
try:
while self.is_active:
await self.send_audio_chunk(stream.read(CHUNK_SIZE, exception_on_overflow=False))
await asyncio.sleep(0.01)
finally:
stream.stop_stream()
stream.close()
p.terminate()
await self.end_audio_input()
async def main():
client = GatewayNovaSonic()
await client.start_session()
playback = asyncio.create_task(client.play_audio())
capture = asyncio.create_task(client.capture_audio())
await asyncio.get_event_loop().run_in_executor(None, input)
for task in (playback, capture):
if not task.done():
task.cancel()
await asyncio.gather(playback, capture, return_exceptions=True)
await client.end_session()
client.is_active = False
if client.response_task and not client.response_task.done():
client.response_task.cancel()
print("Session ended")
if __name__ == "__main__":
asyncio.run(main())
