First, you will add the necessary libraries:
npm install @aws-sdk/client-bedrock-runtime @aws-sdk/credential-providers @smithy/node-http-handler @smithy/typesThen let's add some constants and types to define your configurations, these will be used throughout the application to provide and share information:
export const DefaultInferenceConfiguration = {
maxTokens: 1024,
topP: 0.9,
temperature: 0.7,
};
export const DefaultAudioInputConfiguration = {
audioType: "SPEECH" as AudioType,
encoding: "base64",
mediaType: "audio/lpcm" as AudioMediaType,
sampleRateHertz: 16000,
sampleSizeBits: 16,
channelCount: 1,
};
export const DefaultToolSchema = JSON.stringify({
"$schema": "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {},
required: [],
});
export const DefaultTextConfiguration = {
mediaType: "text/plain" as TextMediaType,
};
export const DefaultSystemPrompt = `
You are a podcast co-host. The user and you will engage in a spoken dialog exchanging
opinions and conversations about AWS. You should ask follow up questions and keep your responses short,
generally two or three sentences for chatty scenarios.
`;
export const DefaultAudioOutputConfiguration = {
...DefaultAudioInputConfiguration,
sampleRateHertz: 24000,
voiceId: "matthew",
};
export interface InferenceConfig {
readonly maxTokens: number;
readonly topP: number;
readonly temperature: number;
}
export type ContentType = "AUDIO" | "TEXT" | "TOOL";
export type AudioType = "SPEECH";
export type AudioMediaType = "audio/wav" | "audio/lpcm" | "audio/mulaw" | "audio/mpeg";
export type TextMediaType = "text/plain" | "application/json";
export interface AudioConfiguration {
readonly audioType: AudioType;
readonly mediaType: AudioMediaType;
readonly sampleRateHertz: number;
readonly sampleSizeBits: number;
readonly channelCount: number;
readonly encoding: string;
readonly voiceId?: string;
}
export interface TextConfiguration {
readonly mediaType: TextMediaType;
}
export interface ToolConfiguration {
readonly toolUseId: string;
readonly type: "TEXT";
readonly textInputConfiguration: {
readonly mediaType: "text/plain";
};
}You will start by implementing the BedrockService class. It is designed to:
- Create and manage sessions with Amazon Bedrock's Nova Sonic model
- Handle audio streaming to the model
- Register event handlers for different types of responses
- Properly close sessions and clean up resources
Create it under services/bedrock/BedrockService.ts file:
import { S2SBidirectionalStreamClient, StreamSession } from './S2SBidirectionalStreamClient';
import { fromIni } from '@aws-sdk/credential-providers';
export class BedrockService {
private client: S2SBidirectionalStreamClient;
private sessions: Map<string, StreamSession> = new Map();
constructor() {
// Initialize the AWS Bedrock client with configuration
this.client = new S2SBidirectionalStreamClient({
requestHandlerConfig: {
maxConcurrentStreams: 10,
},
clientConfig: {
region: process.env.AWS_REGION || 'us-east-1',
credentials: fromIni({ profile: process.env.AWS_PROFILE || 'default' })
}
});
}
/**
* Creates a new session with AWS Bedrock and initializes it with the system prompt
*/
async createSession(sessionId: string, systemPrompt: string): Promise<void> {
try {
// Create a new streaming session
const session = this.client.createStreamSession(sessionId);
this.sessions.set(sessionId, session);
// Initialize the session with AWS Bedrock
await this.client.initiateSession(sessionId);
// Set up the prompt sequence
await session.setupPromptStart();
await session.setupSystemPrompt(undefined, systemPrompt);
await session.setupStartAudio();
console.log(`Session ${sessionId} created successfully`);
} catch (error) {
console.error(`Error creating session ${sessionId}:`, error);
throw error;
}
}
/**
* Streams audio data to AWS Bedrock for processing
*/
async streamAudio(sessionId: string, audioData: Buffer): Promise<void> {
const session = this.sessions.get(sessionId);
if (!session) {
throw new Error(`Session ${sessionId} not found`);
}
await session.streamAudio(audioData);
}
/**
* Registers an event handler for a specific session and event type
*/
registerEventHandler(
sessionId: string,
eventType: string,
handler: (data: any) => void
): void {
const session = this.sessions.get(sessionId);
if (!session) {
throw new Error(`Session ${sessionId} not found`);
}
session.onEvent(eventType, handler);
}
/**
* Properly closes a session and cleans up resources
*/
async closeSession(sessionId: string): Promise<void> {
const session = this.sessions.get(sessionId);
if (!session) return;
try {
// Follow the proper shutdown sequence
await session.endAudioContent();
await session.endPrompt();
await session.close();
// Remove the session from our tracking
this.sessions.delete(sessionId);
console.log(`Session ${sessionId} closed successfully`);
} catch (error) {
console.error(`Error closing session ${sessionId}:`, error);
// Even if there's an error, remove the session from tracking
this.sessions.delete(sessionId);
}
}
/**
* Returns all active session IDs
*/
getActiveSessions(): string[] {
return Array.from(this.sessions.keys());
}
/**
* Checks if a session is active
*/
isSessionActive(sessionId: string): boolean {
return this.sessions.has(sessionId);
}
}Next step is to implement the AudioManager class which handles:
- Initializing the Web Audio API context
- Requesting microphone access from the browser
- Recording audio from the microphone and converting it to the format required by AWS Bedrock
- Playing back audio responses received from the server
- Proper cleanup of audio resources
Create it under services/audio/AudioManager.ts:
export class AudioManager {
private audioContext: AudioContext | null = null;
private audioStream: MediaStream | null = null;
private processor: ScriptProcessorNode | null = null;
private sourceNode: MediaStreamAudioSourceNode | null = null;
private audioPlayer: any = null; // Will be initialized later
private isRecording: boolean = false;
/**
* Initialize the audio context and player
*/
async initialize(): Promise<void> {
try {
// Create audio context with 16kHz sample rate (matching Bedrock's requirements)
this.audioContext = new AudioContext({ sampleRate: 16000 });
// Initialize audio player (implementation will be in a separate class)
this.audioPlayer = new AudioPlayer();
await this.audioPlayer.start();
console.log('AudioManager initialized successfully');
} catch (error) {
console.error('Error initializing AudioManager:', error);
throw error;
}
}
/**
* Request microphone access from the browser
*/
async requestMicrophoneAccess(): Promise<boolean> {
try {
this.audioStream = await navigator.mediaDevices.getUserMedia({
audio: {
echoCancellation: true,
noiseSuppression: true,
autoGainControl: true
}
});
return true;
} catch (error) {
console.error("Error accessing microphone:", error);
return false;
}
}
/**
* Start recording audio from the microphone
* @param onAudioData Callback function that receives audio data
*/
startRecording(onAudioData: (data: string) => void): void {
if (!this.audioContext || !this.audioStream) {
console.error('AudioContext or AudioStream not initialized');
return;
}
try {
this.isRecording = true;
// Create audio source from microphone stream
this.sourceNode = this.audioContext.createMediaStreamSource(this.audioStream);
// Create script processor for audio processing
this.processor = this.audioContext.createScriptProcessor(512, 1, 1);
// Process audio data
this.processor.onaudioprocess = (e) => {
if (!this.isRecording) return;
const inputData = e.inputBuffer.getChannelData(0);
// Convert to 16-bit PCM (required by Bedrock)
const pcmData = new Int16Array(inputData.length);
for (let i = 0; i < inputData.length; i++) {
pcmData[i] = Math.max(-1, Math.min(1, inputData[i])) * 0x7FFF;
}
// Convert to base64 for transmission
const base64Data = this.arrayBufferToBase64(pcmData.buffer);
// Send to callback
onAudioData(base64Data);
};
// Connect the audio nodes
this.sourceNode.connect(this.processor);
this.processor.connect(this.audioContext.destination);
console.log('Recording started');
} catch (error) {
console.error('Error starting recording:', error);
this.isRecording = false;
}
}
/**
* Stop recording audio
*/
stopRecording(): void {
this.isRecording = false;
// Clean up audio processing nodes
if (this.processor) {
this.processor.disconnect();
this.processor = null;
}
if (this.sourceNode) {
this.sourceNode.disconnect();
this.sourceNode = null;
}
console.log('Recording stopped');
}
/**
* Play audio data received from the server
* @param audioData Base64 encoded audio data
*/
playAudio(audioData: string): void {
if (!this.audioPlayer) {
console.error('AudioPlayer not initialized');
return;
}
try {
// Convert base64 to Float32Array for playback
const audioBuffer = this.base64ToFloat32Array(audioData);
this.audioPlayer.playAudio(audioBuffer);
} catch (error) {
console.error('Error playing audio:', error);
}
}
/**
* Clean up resources when component unmounts
*/
cleanup(): void {
this.stopRecording();
if (this.audioContext) {
this.audioContext.close();
this.audioContext = null;
}
if (this.audioStream) {
this.audioStream.getTracks().forEach(track => track.stop());
this.audioStream = null;
}
if (this.audioPlayer) {
this.audioPlayer.stop();
this.audioPlayer = null;
}
console.log('AudioManager cleaned up');
}
/**
* Convert ArrayBuffer to base64 string
*/
private arrayBufferToBase64(buffer: ArrayBuffer): string {
const binary = [];
const bytes = new Uint8Array(buffer);
for (let i = 0; i < bytes.byteLength; i++) {
binary.push(String.fromCharCode(bytes[i]));
}
return btoa(binary.join(''));
}
/**
* Convert base64 string to Float32Array for audio playback
*/
private base64ToFloat32Array(base64String: string): Float32Array {
const binaryString = window.atob(base64String);
const bytes = new Uint8Array(binaryString.length);
for (let i = 0; i < binaryString.length; i++) {
bytes[i] = binaryString.charCodeAt(i);
}
const int16Array = new Int16Array(bytes.buffer);
const float32Array = new Float32Array(int16Array.length);
for (let i = 0; i < int16Array.length; i++) {
float32Array[i] = int16Array[i] / 32768.0;
}
return float32Array;
}
/**
* Get the current audio context
*/
getAudioContext(): AudioContext | null {
return this.audioContext;
}
/**
* Check if recording is active
*/
isActive(): boolean {
return this.isRecording;
}
}
/**
* Audio Player class for handling audio playback
* This is a simplified version that would be expanded in the actual implementation
*/
class AudioPlayer {
private audioContext: AudioContext | null = null;
private workletNode: AudioWorkletNode | null = null;
private initialized: boolean = false;
async start(): Promise<void> {
this.audioContext = new AudioContext({ sampleRate: 24000 });
// Load the audio worklet
await this.audioContext.audioWorklet.addModule('/audio-player-processor.js');
this.workletNode = new AudioWorkletNode(this.audioContext, 'audio-player-processor');
this.workletNode.connect(this.audioContext.destination);
this.initialized = true;
}
playAudio(samples: Float32Array): void {
if (!this.initialized || !this.workletNode) {
console.error("AudioPlayer not initialized");
return;
}
this.workletNode.port.postMessage({
type: "audio",
audioData: samples,
});
}
stop(): void {
if (this.audioContext) {
this.audioContext.close();
}
if (this.workletNode) {
this.workletNode.disconnect();
}
this.initialized = false;
this.audioContext = null;
this.workletNode = null;
}
}Lastly, you will add a conversation manager as your core functionality. The ConversationManager class:
- Manages the state of the conversation, including message history and speaking/thinking indicators
- Provides methods to update the state and notify listeners
- Handles message concatenation when multiple messages come from the same role
- Includes a utility to create a React Context for easy integration with React components
Create it under services/conversation/ConversationManager.ts
export interface Message {
role: string;
content: string;
timestamp: number;
}
export interface ConversationState {
messages: Message[];
isUserSpeaking: boolean;
isAssistantThinking: boolean;
isAssistantSpeaking: boolean;
isConnected: boolean;
}
/**
* ConversationManager handles the state of the conversation,
* including message history and speaking/thinking indicators
*/
export class ConversationManager {
private state: ConversationState;
private listeners: ((state: ConversationState) => void)[] = [];
constructor() {
// Initialize with default state
this.state = {
messages: [],
isUserSpeaking: false,
isAssistantThinking: false,
isAssistantSpeaking: false,
isConnected: false
};
}
/**
* Subscribe to state changes
* @param listener Function to call when state changes
* @returns Function to unsubscribe
*/
subscribe(listener: (state: ConversationState) => void): () => void {
this.listeners.push(listener);
// Call the listener immediately with current state
listener({...this.state});
// Return unsubscribe function
return () => {
this.listeners = this.listeners.filter(l => l !== listener);
};
}
/**
* Notify all listeners of state changes
*/
private notifyListeners(): void {
const stateCopy = {...this.state};
this.listeners.forEach(listener => listener(stateCopy));
}
/**
* Add a new message to the conversation
* @param role Role of the message sender (USER or ASSISTANT)
* @param content Content of the message
*/
addMessage(role: string, content: string): void {
// Check if the last message is from the same role
const lastMessage = this.state.messages[this.state.messages.length - 1];
if (lastMessage && lastMessage.role === role) {
// Update the last message
const updatedMessages = [...this.state.messages];
updatedMessages[updatedMessages.length - 1] = {
...lastMessage,
content: lastMessage.content + " " + content,
timestamp: Date.now(),
};
this.state.messages = updatedMessages;
} else {
// Add a new message
this.state.messages = [
...this.state.messages,
{
role,
content,
timestamp: Date.now(),
},
];
}
this.notifyListeners();
}
/**
* Set the user speaking state
* @param isSpeaking Whether the user is speaking
*/
setUserSpeaking(isSpeaking: boolean): void {
this.state.isUserSpeaking = isSpeaking;
this.notifyListeners();
}
/**
* Set the assistant thinking state
* @param isThinking Whether the assistant is thinking
*/
setAssistantThinking(isThinking: boolean): void {
this.state.isAssistantThinking = isThinking;
this.notifyListeners();
}
/**
* Set the assistant speaking state
* @param isSpeaking Whether the assistant is speaking
*/
setAssistantSpeaking(isSpeaking: boolean): void {
this.state.isAssistantSpeaking = isSpeaking;
this.notifyListeners();
}
/**
* Set the connection state
* @param isConnected Whether the WebSocket is connected
*/
setConnected(isConnected: boolean): void {
this.state.isConnected = isConnected;
this.notifyListeners();
}
/**
* Clear the conversation history
*/
clearConversation(): void {
this.state.messages = [];
this.notifyListeners();
}
/**
* Get the current state
* @returns Copy of the current state
*/
getState(): ConversationState {
return {...this.state};
}
/**
* End the current turn in the conversation
*/
endTurn(): void {
// This could be expanded to add metadata to messages
// or perform other actions when a turn ends
this.notifyListeners();
}
/**
* Mark the conversation as ended
*/
endConversation(): void {
// This could add a system message or perform cleanup
this.notifyListeners();
}
}
/**
* Create a React Context wrapper for the ConversationManager
* This would be implemented in a separate file in the actual application
*/
export function createConversationContext() {
const conversationManager = new ConversationManager();
return {
useConversation: () => {
// This would be implemented using React's useState and useEffect
// to subscribe to the ConversationManager
return {
state: conversationManager.getState(),
addMessage: conversationManager.addMessage.bind(conversationManager),
setUserSpeaking: conversationManager.setUserSpeaking.bind(conversationManager),
setAssistantThinking: conversationManager.setAssistantThinking.bind(conversationManager),
setAssistantSpeaking: conversationManager.setAssistantSpeaking.bind(conversationManager),
setConnected: conversationManager.setConnected.bind(conversationManager),
clearConversation: conversationManager.clearConversation.bind(conversationManager),
};
}
};
}To provide a complete Socket.IO integration for both the server and client sides, with proper context providers for React components to access the socket functionality and conversation state. The code handles:
- Socket.IO server setup with Next.js integration
- Client-side socket connection management
- Event handling for all the bidirectional communication
- React context providers for easy access to socket and conversation state
- Proper cleanup and error handling
- First create a server.js file with Socket.IO implementation:
import {
BedrockRuntimeClient,
BedrockRuntimeClientConfig,
InvokeModelWithBidirectionalStreamCommand,
InvokeModelWithBidirectionalStreamInput,
} from "@aws-sdk/client-bedrock-runtime";
import {
NodeHttp2Handler,
NodeHttp2HandlerOptions,
} from "@smithy/node-http-handler";
import { Provider } from "@smithy/types";
import { Buffer } from "node:buffer";
import { randomUUID } from "node:crypto";
import { InferenceConfig } from "./types";
import { Subject } from "rxjs";
import { take } from "rxjs/operators";
import { firstValueFrom } from "rxjs";
import {
DefaultAudioInputConfiguration,
DefaultAudioOutputConfiguration,
DefaultSystemPrompt,
DefaultTextConfiguration,
DefaultToolSchema,
} from "./consts";
export interface S2SBidirectionalStreamClientConfig {
requestHandlerConfig?:
| NodeHttp2HandlerOptions
| Provider<NodeHttp2HandlerOptions | void>;
clientConfig: Partial<BedrockRuntimeClientConfig>;
inferenceConfig?: InferenceConfig;
}
export class StreamSession {
private audioBufferQueue: Buffer[] = [];
private maxQueueSize = 200; // Maximum number of audio chunks to queue
private isProcessingAudio = false;
private isActive = true;
constructor(
private sessionId: string,
private client: S2SBidirectionalStreamClient
) {}
// Register event handlers for this specific session
public onEvent(
eventType: string,
handler: (data: any) => void
): StreamSession {
this.client.registerEventHandler(this.sessionId, eventType, handler);
return this; // For chaining
}
public async setupPromptStart(): Promise<void> {
this.client.setupPromptStartEvent(this.sessionId);
}
public async setupSystemPrompt(
textConfig: typeof DefaultTextConfiguration = DefaultTextConfiguration,
systemPromptContent: string = DefaultSystemPrompt
): Promise<void> {
this.client.setupSystemPromptEvent(
this.sessionId,
textConfig,
systemPromptContent
);
}
public async setupStartAudio(
audioConfig: typeof DefaultAudioInputConfiguration = DefaultAudioInputConfiguration
): Promise<void> {
this.client.setupStartAudioEvent(this.sessionId, audioConfig);
}
// Stream audio for this session
public async streamAudio(audioData: Buffer): Promise<void> {
// Check queue size to avoid memory issues
if (this.audioBufferQueue.length >= this.maxQueueSize) {
// Queue is full, drop oldest chunk
this.audioBufferQueue.shift();
console.log("Audio queue full, dropping oldest chunk");
}
// Queue the audio chunk for streaming
this.audioBufferQueue.push(audioData);
this.processAudioQueue();
}
// Process audio queue for continuous streaming
private async processAudioQueue() {
if (
this.isProcessingAudio ||
this.audioBufferQueue.length === 0 ||
!this.isActive
)
return;
this.isProcessingAudio = true;
try {
// Process all chunks in the queue, up to a reasonable limit
let processedChunks = 0;
const maxChunksPerBatch = 5; // Process max 5 chunks at a time to avoid overload
while (
this.audioBufferQueue.length > 0 &&
processedChunks < maxChunksPerBatch &&
this.isActive
) {
const audioChunk = this.audioBufferQueue.shift();
if (audioChunk) {
await this.client.streamAudioChunk(this.sessionId, audioChunk);
processedChunks++;
}
}
} finally {
this.isProcessingAudio = false;
// If there are still items in the queue, schedule the next processing using setTimeout
if (this.audioBufferQueue.length > 0 && this.isActive) {
setTimeout(() => this.processAudioQueue(), 0);
}
}
}
// Get session ID
public getSessionId(): string {
return this.sessionId;
}
public async endAudioContent(): Promise<void> {
if (!this.isActive) return;
await this.client.sendContentEnd(this.sessionId);
}
public async endPrompt(): Promise<void> {
if (!this.isActive) return;
await this.client.sendPromptEnd(this.sessionId);
}
public async close(): Promise<void> {
if (!this.isActive) return;
this.isActive = false;
this.audioBufferQueue = []; // Clear any pending audio
await this.client.sendSessionEnd(this.sessionId);
console.log(`Session ${this.sessionId} close completed`);
}
}
// Session data type
interface SessionData {
queue: Array<any>;
queueSignal: Subject<void>;
closeSignal: Subject<void>;
responseSubject: Subject<any>;
toolUseContent: any;
toolUseId: string;
toolName: string;
responseHandlers: Map<string, (data: any) => void>;
promptName: string;
inferenceConfig: InferenceConfig;
isActive: boolean;
isPromptStartSent: boolean;
isAudioContentStartSent: boolean;
audioContentId: string;
}
export class S2SBidirectionalStreamClient {
private bedrockRuntimeClient: BedrockRuntimeClient;
private inferenceConfig: InferenceConfig;
private activeSessions: Map<string, SessionData> = new Map();
private sessionLastActivity: Map<string, number> = new Map();
private sessionCleanupInProgress = new Set<string>();
constructor(config: S2SBidirectionalStreamClientConfig) {
const nodeClientHandler = new NodeHttp2Handler({
requestTimeout: 300000,
sessionTimeout: 300000,
disableConcurrentStreams: false,
maxConcurrentStreams: 20,
...config.requestHandlerConfig,
});
if (!config.clientConfig.credentials) {
throw new Error("No credentials provided");
}
this.bedrockRuntimeClient = new BedrockRuntimeClient({
...config.clientConfig,
credentials: config.clientConfig.credentials,
region: config.clientConfig.region || "us-east-1",
requestHandler: nodeClientHandler,
});
this.inferenceConfig = config.inferenceConfig ?? {
maxTokens: 1024,
topP: 0.9,
temperature: 0.7,
};
}
public isSessionActive(sessionId: string): boolean {
const session = this.activeSessions.get(sessionId);
return !!session && session.isActive;
}
public getActiveSessions(): string[] {
return Array.from(this.activeSessions.keys());
}
public getLastActivityTime(sessionId: string): number {
return this.sessionLastActivity.get(sessionId) || 0;
}
private updateSessionActivity(sessionId: string): void {
this.sessionLastActivity.set(sessionId, Date.now());
}
public isCleanupInProgress(sessionId: string): boolean {
return this.sessionCleanupInProgress.has(sessionId);
}
// Create a new streaming session
public createStreamSession(
sessionId: string = randomUUID(),
config?: S2SBidirectionalStreamClientConfig
): StreamSession {
if (this.activeSessions.has(sessionId)) {
throw new Error(`Stream session with ID ${sessionId} already exists`);
}
const session: SessionData = {
queue: [],
queueSignal: new Subject<void>(),
closeSignal: new Subject<void>(),
responseSubject: new Subject<any>(),
toolUseContent: null,
toolUseId: "",
toolName: "",
responseHandlers: new Map(),
promptName: randomUUID(),
inferenceConfig: config?.inferenceConfig ?? this.inferenceConfig,
isActive: true,
isPromptStartSent: false,
isAudioContentStartSent: false,
audioContentId: randomUUID(),
};
this.activeSessions.set(sessionId, session);
return new StreamSession(sessionId, this);
}
private async processToolUse(toolName: string): Promise<Object> {
const tool = toolName.toLowerCase();
switch (tool) {
case "gettimetool":
const pstTime = new Date().toLocaleString("en-US", {
timeZone: "America/Los_Angeles",
});
return {
timezone: "PST",
formattedTime: new Date(pstTime).toLocaleTimeString("en-US", {
hour12: true,
hour: "2-digit",
minute: "2-digit",
}),
};
case "getdatetool":
const date = new Date().toLocaleString("en-US", {
timeZone: "America/Los_Angeles",
});
const pstDate = new Date(date);
return {
date: pstDate.toISOString().split("T")[0],
year: pstDate.getFullYear(),
month: pstDate.getMonth() + 1,
day: pstDate.getDate(),
dayOfWeek: pstDate
.toLocaleString("en-US", { weekday: "long" })
.toUpperCase(),
timezone: "PST",
};
default:
console.log(`Tool ${tool} not supported`);
throw new Error(`Tool ${tool} not supported`);
}
}
// Stream audio for a specific session
public async initiateSession(sessionId: string): Promise<void> {
const session = this.activeSessions.get(sessionId);
if (!session) {
throw new Error(`Stream session ${sessionId} not found`);
}
try {
// Set up initial events for this session
this.setupSessionStartEvent(sessionId);
// Create the bidirectional stream with session-specific async iterator
const asyncIterable = this.createSessionAsyncIterable(sessionId);
console.log(`Starting bidirectional stream for session ${sessionId}...`);
const response = await this.bedrockRuntimeClient.send(
new InvokeModelWithBidirectionalStreamCommand({
modelId: "amazon.nova-sonic-v1:0",
body: asyncIterable,
})
);
console.log(
`Stream established for session ${sessionId}, processing responses...`
);
// Process responses for this session
await this.processResponseStream(sessionId, response);
} catch (error) {
console.error(`Error in session ${sessionId}:`, error);
this.dispatchEventForSession(sessionId, "error", {
source: "bidirectionalStream",
error,
});
// Make sure to clean up if there's an error
if (session.isActive) {
this.closeSession(sessionId);
}
}
}
// Dispatch events to handlers for a specific session
private dispatchEventForSession(
sessionId: string,
eventType: string,
data: any
): void {
const session = this.activeSessions.get(sessionId);
if (!session) return;
const handler = session.responseHandlers.get(eventType);
if (handler) {
try {
handler(data);
} catch (e) {
console.error(
`Error in ${eventType} handler for session ${sessionId}:`,
e
);
}
}
// Also dispatch to "any" handlers
const anyHandler = session.responseHandlers.get("any");
if (anyHandler) {
try {
anyHandler({ type: eventType, data });
} catch (e) {
console.error(`Error in 'any' handler for session ${sessionId}:`, e);
}
}
}
private createSessionAsyncIterable(
sessionId: string
): AsyncIterable<InvokeModelWithBidirectionalStreamInput> {
if (!this.isSessionActive(sessionId)) {
console.log(
`Cannot create async iterable: Session ${sessionId} not active`
);
return {
[Symbol.asyncIterator]: () => ({
next: async () => ({ value: undefined, done: true }),
}),
};
}
const session = this.activeSessions.get(sessionId);
if (!session) {
throw new Error(
`Cannot create async iterable: Session ${sessionId} not found`
);
}
let eventCount = 0;
return {
[Symbol.asyncIterator]: () => {
console.log(
`AsyncIterable iterator requested for session ${sessionId}`
);
return {
next: async (): Promise<
IteratorResult<InvokeModelWithBidirectionalStreamInput>
> => {
try {
// Check if session is still active
if (!session.isActive || !this.activeSessions.has(sessionId)) {
console.log(
`Iterator closing for session ${sessionId}, done=true`
);
return { value: undefined, done: true };
}
// Wait for items in the queue or close signal
if (session.queue.length === 0) {
try {
await Promise.race([
firstValueFrom(session.queueSignal.pipe(take(1))),
firstValueFrom(session.closeSignal.pipe(take(1))).then(
() => {
throw new Error("Stream closed");
}
),
]);
} catch (error) {
if (error instanceof Error) {
if (
error.message === "Stream closed" ||
!session.isActive
) {
// This is an expected condition when closing the session
if (this.activeSessions.has(sessionId)) {
console.log(`Session \${sessionId} closed during wait`);
}
return { value: undefined, done: true };
}
} else {
console.error(`Error on event close`, error);
}
}
}
// If queue is still empty or session is inactive, we're done
if (session.queue.length === 0 || !session.isActive) {
console.log(`Queue empty or session inactive: ${sessionId}`);
return { value: undefined, done: true };
}
// Get next item from the session's queue
const nextEvent = session.queue.shift();
eventCount++;
console.log(
`Sending event #${eventCount} for session ${sessionId}: ${JSON.stringify(
nextEvent
).substring(0, 100)}...`
);
return {
value: {
chunk: {
bytes: new TextEncoder().encode(JSON.stringify(nextEvent)),
},
},
done: false,
};
} catch (error) {
console.error(`Error in session ${sessionId} iterator:`, error);
session.isActive = false;
return { value: undefined, done: true };
}
},
return: async (): Promise<
IteratorResult<InvokeModelWithBidirectionalStreamInput>
> => {
console.log(`Iterator return() called for session ${sessionId}`);
session.isActive = false;
return { value: undefined, done: true };
},
throw: async (
error: any
): Promise<
IteratorResult<InvokeModelWithBidirectionalStreamInput>
> => {
console.log(
`Iterator throw() called for session ${sessionId} with error:`,
error
);
session.isActive = false;
throw error;
},
};
},
};
}
// Process the response stream from AWS Bedrock
private async processResponseStream(
sessionId: string,
response: any
): Promise<void> {
const session = this.activeSessions.get(sessionId);
if (!session) return;
try {
for await (const event of response.body) {
if (!session.isActive) {
console.log(
`Session ${sessionId} is no longer active, stopping response processing`
);
break;
}
if (event.chunk?.bytes) {
try {
this.updateSessionActivity(sessionId);
const textResponse = new TextDecoder().decode(event.chunk.bytes);
try {
const jsonResponse = JSON.parse(textResponse);
if (jsonResponse.event?.contentStart) {
this.dispatchEvent(
sessionId,
"contentStart",
jsonResponse.event.contentStart
);
} else if (jsonResponse.event?.textOutput) {
this.dispatchEvent(
sessionId,
"textOutput",
jsonResponse.event.textOutput
);
} else if (jsonResponse.event?.audioOutput) {
this.dispatchEvent(
sessionId,
"audioOutput",
jsonResponse.event.audioOutput
);
} else if (jsonResponse.event?.toolUse) {
this.dispatchEvent(
sessionId,
"toolUse",
jsonResponse.event.toolUse
);
// Store tool use information for later
session.toolUseContent = jsonResponse.event.toolUse;
session.toolUseId = jsonResponse.event.toolUse.toolUseId;
session.toolName = jsonResponse.event.toolUse.toolName;
} else if (
jsonResponse.event?.contentEnd &&
jsonResponse.event?.contentEnd?.type === "TOOL"
) {
// Process tool use
console.log(`Processing tool use for session ${sessionId}`);
this.dispatchEvent(sessionId, "toolEnd", {
toolUseContent: session.toolUseContent,
toolUseId: session.toolUseId,
toolName: session.toolName,
});
console.log("calling tooluse");
// Call external model
const externalModelResult = await this.processToolUse(
session.toolName
);
// Send tool result
this.sendToolResult(
sessionId,
session.toolUseId,
externalModelResult
);
// Also dispatch event about tool result
this.dispatchEvent(sessionId, "toolResult", {
toolUseId: session.toolUseId,
result: externalModelResult,
});
} else {
// Handle other events
const eventKeys = Object.keys(jsonResponse.event || {});
console.log(`Event keys for session ${sessionId}:`, eventKeys);
console.log(`Handling other events`);
if (eventKeys.length > 0) {
this.dispatchEvent(
sessionId,
eventKeys[0],
jsonResponse.event[eventKeys[0]]
);
} else if (Object.keys(jsonResponse).length > 0) {
this.dispatchEvent(sessionId, "unknown", jsonResponse);
}
}
} catch (e) {
console.log(
`Raw text response for session ${sessionId} (parse error):`,
textResponse
);
}
} catch (e) {
console.error(
`Error processing response chunk for session ${sessionId}:`,
e
);
}
} else if (event.modelStreamErrorException) {
console.error(
`Model stream error for session ${sessionId}:`,
event.modelStreamErrorException
);
this.dispatchEvent(sessionId, "error", {
type: "modelStreamErrorException",
details: event.modelStreamErrorException,
});
} else if (event.internalServerException) {
console.error(
`Internal server error for session ${sessionId}:`,
event.internalServerException
);
this.dispatchEvent(sessionId, "error", {
type: "internalServerException",
details: event.internalServerException,
});
}
}
console.log(
`Response stream processing complete for session ${sessionId}`
);
this.dispatchEvent(sessionId, "streamComplete", {
timestamp: new Date().toISOString(),
});
} catch (error) {
console.error(
`Error processing response stream for session ${sessionId}:`,
error
);
this.dispatchEvent(sessionId, "error", {
source: "responseStream",
message: "Error processing response stream",
details: error instanceof Error ? error.message : String(error),
});
}
}
// Add an event to a session's queue
private addEventToSessionQueue(sessionId: string, event: any): void {
const session = this.activeSessions.get(sessionId);
if (!session || !session.isActive) return;
this.updateSessionActivity(sessionId);
session.queue.push(event);
session.queueSignal.next();
}
// Set up initial events for a session
private setupSessionStartEvent(sessionId: string): void {
console.log(`Setting up initial events for session ${sessionId}...`);
const session = this.activeSessions.get(sessionId);
if (!session) return;
// Session start event
this.addEventToSessionQueue(sessionId, {
event: {
sessionStart: {
inferenceConfiguration: session.inferenceConfig,
},
},
});
}
public setupPromptStartEvent(sessionId: string): void {
console.log(`Setting up prompt start event for session ${sessionId}...`);
const session = this.activeSessions.get(sessionId);
if (!session) return;
// Prompt start event
this.addEventToSessionQueue(sessionId, {
event: {
promptStart: {
promptName: session.promptName,
textOutputConfiguration: {
mediaType: "text/plain",
},
audioOutputConfiguration: DefaultAudioOutputConfiguration,
toolUseOutputConfiguration: {
mediaType: "application/json",
},
toolConfiguration: {
tools: [
{
toolSpec: {
name: "getDateTool",
description: "get information about the current date",
inputSchema: {
json: DefaultToolSchema,
},
},
},
{
toolSpec: {
name: "getTimeTool",
description: "get information about the current time",
inputSchema: {
json: DefaultToolSchema,
},
},
},
],
},
},
},
});
session.isPromptStartSent = true;
}
public setupSystemPromptEvent(
sessionId: string,
textConfig: typeof DefaultTextConfiguration = DefaultTextConfiguration,
systemPromptContent: string = DefaultSystemPrompt
): void {
console.log(`Setting up systemPrompt events for session ${sessionId}...`);
const session = this.activeSessions.get(sessionId);
if (!session) return;
// Text content start
const textPromptID = randomUUID();
this.addEventToSessionQueue(sessionId, {
event: {
contentStart: {
promptName: session.promptName,
contentName: textPromptID,
type: "TEXT",
interactive: true,
textInputConfiguration: textConfig,
},
},
});
// Text input content
this.addEventToSessionQueue(sessionId, {
event: {
textInput: {
promptName: session.promptName,
contentName: textPromptID,
content: systemPromptContent,
role: "SYSTEM",
},
},
});
// Text content end
this.addEventToSessionQueue(sessionId, {
event: {
contentEnd: {
promptName: session.promptName,
contentName: textPromptID,
},
},
});
}
public setupStartAudioEvent(
sessionId: string,
audioConfig: typeof DefaultAudioInputConfiguration = DefaultAudioInputConfiguration
): void {
console.log(
`Setting up startAudioContent event for session ${sessionId}...`
);
const session = this.activeSessions.get(sessionId);
if (!session) return;
console.log(`Using audio content ID: ${session.audioContentId}`);
// Audio content start
this.addEventToSessionQueue(sessionId, {
event: {
contentStart: {
promptName: session.promptName,
contentName: session.audioContentId,
type: "AUDIO",
interactive: true,
audioInputConfiguration: audioConfig,
},
},
});
session.isAudioContentStartSent = true;
console.log(`Initial events setup complete for session ${sessionId}`);
}
// Stream an audio chunk for a session
public async streamAudioChunk(
sessionId: string,
audioData: Buffer
): Promise<void> {
const session = this.activeSessions.get(sessionId);
if (!session || !session.isActive || !session.audioContentId) {
throw new Error(`Invalid session ${sessionId} for audio streaming`);
}
// Convert audio to base64
const base64Data = audioData.toString("base64");
this.addEventToSessionQueue(sessionId, {
event: {
audioInput: {
promptName: session.promptName,
contentName: session.audioContentId,
content: base64Data,
role: "USER",
},
},
});
}
// Send tool result back to the model
private async sendToolResult(
sessionId: string,
toolUseId: string,
result: any
): Promise<void> {
const session = this.activeSessions.get(sessionId);
console.log("inside tool result");
if (!session || !session.isActive) return;
console.log(
`Sending tool result for session ${sessionId}, tool use ID: ${toolUseId}`
);
const contentId = randomUUID();
// Tool content start
this.addEventToSessionQueue(sessionId, {
event: {
contentStart: {
promptName: session.promptName,
contentName: contentId,
interactive: false,
type: "TOOL",
toolResultInputConfiguration: {
toolUseId: toolUseId,
type: "TEXT",
textInputConfiguration: {
mediaType: "text/plain",
},
},
},
},
});
// Tool content input
const resultContent =
typeof result === "string" ? result : JSON.stringify(result);
this.addEventToSessionQueue(sessionId, {
event: {
toolResult: {
promptName: session.promptName,
contentName: contentId,
content: resultContent,
role: "TOOL",
},
},
});
// Tool content end
this.addEventToSessionQueue(sessionId, {
event: {
contentEnd: {
promptName: session.promptName,
contentName: contentId,
},
},
});
console.log(`Tool result sent for session ${sessionId}`);
}
public async sendContentEnd(sessionId: string): Promise<void> {
const session = this.activeSessions.get(sessionId);
if (!session || !session.isAudioContentStartSent) return;
await this.addEventToSessionQueue(sessionId, {
event: {
contentEnd: {
promptName: session.promptName,
contentName: session.audioContentId,
},
},
});
// Wait to ensure it's processed
await new Promise((resolve) => setTimeout(resolve, 500));
}
public async sendPromptEnd(sessionId: string): Promise<void> {
const session = this.activeSessions.get(sessionId);
if (!session || !session.isPromptStartSent) return;
await this.addEventToSessionQueue(sessionId, {
event: {
promptEnd: {
promptName: session.promptName,
},
},
});
// Wait to ensure it's processed
await new Promise((resolve) => setTimeout(resolve, 300));
}
public async sendSessionEnd(sessionId: string): Promise<void> {
const session = this.activeSessions.get(sessionId);
if (!session) return;
await this.addEventToSessionQueue(sessionId, {
event: {
sessionEnd: {},
},
});
// Wait to ensure it's processed
await new Promise((resolve) => setTimeout(resolve, 300));
// Now it's safe to clean up
session.isActive = false;
session.closeSignal.next();
session.closeSignal.complete();
this.activeSessions.delete(sessionId);
this.sessionLastActivity.delete(sessionId);
console.log(`Session ${sessionId} closed and removed from active sessions`);
}
// Register an event handler for a session
public registerEventHandler(
sessionId: string,
eventType: string,
handler: (data: any) => void
): void {
const session = this.activeSessions.get(sessionId);
if (!session) {
throw new Error(`Session ${sessionId} not found`);
}
session.responseHandlers.set(eventType, handler);
}
// Dispatch an event to registered handlers
private dispatchEvent(sessionId: string, eventType: string, data: any): void {
const session = this.activeSessions.get(sessionId);
if (!session) return;
const handler = session.responseHandlers.get(eventType);
if (handler) {
try {
handler(data);
} catch (e) {
console.error(
`Error in ${eventType} handler for session ${sessionId}:`,
e
);
}
}
// Also dispatch to "any" handlers
const anyHandler = session.responseHandlers.get("any");
if (anyHandler) {
try {
anyHandler({ type: eventType, data });
} catch (e) {
console.error(`Error in 'any' handler for session ${sessionId}:`, e);
}
}
}
public async closeSession(sessionId: string): Promise<void> {
if (this.sessionCleanupInProgress.has(sessionId)) {
console.log(
`Cleanup already in progress for session ${sessionId}, skipping`
);
return;
}
this.sessionCleanupInProgress.add(sessionId);
try {
console.log(`Starting close process for session ${sessionId}`);
await this.sendContentEnd(sessionId);
await this.sendPromptEnd(sessionId);
await this.sendSessionEnd(sessionId);
console.log(`Session ${sessionId} cleanup complete`);
} catch (error) {
console.error(
`Error during closing sequence for session ${sessionId}:`,
error
);
// Ensure cleanup happens even if there's an error
const session = this.activeSessions.get(sessionId);
if (session) {
session.isActive = false;
this.activeSessions.delete(sessionId);
this.sessionLastActivity.delete(sessionId);
}
} finally {
// Always clean up the tracking set
this.sessionCleanupInProgress.delete(sessionId);
}
}
// Same for forceCloseSession:
public forceCloseSession(sessionId: string): void {
if (
this.sessionCleanupInProgress.has(sessionId) ||
!this.activeSessions.has(sessionId)
) {
console.log(
`Session ${sessionId} already being cleaned up or not active`
);
return;
}
this.sessionCleanupInProgress.add(sessionId);
try {
const session = this.activeSessions.get(sessionId);
if (!session) return;
console.log(`Force closing session ${sessionId}`);
// Immediately mark as inactive and clean up resources
session.isActive = false;
session.closeSignal.next();
session.closeSignal.complete();
this.activeSessions.delete(sessionId);
this.sessionLastActivity.delete(sessionId);
console.log(`Session ${sessionId} force closed`);
} finally {
this.sessionCleanupInProgress.delete(sessionId);
}
}
}