WebSocket Protocol

UG Labs uses a WebSocket-based protocol for real-time, bidirectional communication between your application and our AI services.

Connection URL

wss://pug.stg.uglabs.app/interact

Message Structure

All messages sent and received through the WebSocket are JSON objects with a common structure:

interface BaseMessage {
  kind: string;          // Message type
  uid: string;           // Unique identifier for correlation
  [key: string]: any;    // Additional fields based on message type
}

Connection Flow

1. Establish WebSocket Connection

const ws = new WebSocket('wss://pug.stg.uglabs.app/interact');

ws.onopen = () => {
  console.log('WebSocket connected');
};

2. Authentication

Before WebSocket authentication, you need to obtain an access token:

REST Authentication (performed once):

const response = await fetch('https://pug.stg.uglabs.app/api/auth/login', {
  method: 'POST',
  headers: {
    'Content-Type': 'application/json'
  },
  body: JSON.stringify({
    api_key: 'your-api-key',
    federated_id: 'user-123' // Optional
  })
});

const { access_token } = await response.json();

Then authenticate the WebSocket connection:

ws.send(JSON.stringify({
  kind: 'authenticate',
  uid: crypto.randomUUID(),
  access_token: access_token
}));

Response:

{
  "kind": "authenticated",
  "uid": "matching-uid",
  "success": true
}

3. Set Configuration

Configure the conversation behavior:

ws.send(JSON.stringify({
  kind: 'set_configuration',
  uid: crypto.randomUUID(),
  prompt: 'You are a helpful assistant.',
  utilities: {
    // Optional utilities
  }
}));

Response:

{
  "kind": "configured",
  "uid": "matching-uid"
}

4. Interact

Send messages and receive responses:

ws.send(JSON.stringify({
  kind: 'interact',
  uid: crypto.randomUUID(),
  text: 'Hello! How are you?',
  audio_output: false
}));

Streaming Response:

{
  "kind": "text",
  "uid": "matching-uid",
  "text": "Hello! I'm doing well, thanks for asking.",
  "done": false
}

{
  "kind": "text",
  "uid": "matching-uid",
  "text": " How can I help you today?",
  "done": true
}

Message Types

Request Messages

1. authenticate - Authenticate the WebSocket connection
2. set_configuration - Configure conversation parameters
3. interact - Send a text message or process audio input
4. add_audio - Stream audio chunks for voice input
5. generate_image - Request image generation

Response Messages

1. authenticated - Confirmation of successful authentication
2. configured - Confirmation of configuration
3. text - Streaming text response
4. audio - Audio data chunks (base64 encoded)
5. data - Structured data from utilities
6. image - Generated image (base64 encoded)
7. error - Error message

See Message Types for detailed documentation of each message type.

UID Correlation

Each request includes a unique uid field. Responses related to that request will include the same uid, allowing you to correlate requests and responses:

const requestUid = crypto.randomUUID();

ws.send(JSON.stringify({
  kind: 'interact',
  uid: requestUid,
  text: 'Hello'
}));

ws.onmessage = (event) => {
  const message = JSON.parse(event.data);

  if (message.uid === requestUid) {
    // This response is for our request
    console.log('Response:', message);
  }
};

Streaming Responses

Text responses are streamed in chunks for real-time display:

let fullResponse = '';

ws.onmessage = (event) => {
  const message = JSON.parse(event.data);

  if (message.kind === 'text') {
    fullResponse += message.text;

    if (message.done) {
      console.log('Complete response:', fullResponse);
      fullResponse = '';
    }
  }
};

Audio Streaming

Audio input is sent in chunks:

// Send multiple audio chunks
for (const chunk of audioChunks) {
  ws.send(JSON.stringify({
    kind: 'add_audio',
    uid: crypto.randomUUID(),
    audio: base64AudioData,
    config: {
      sample_rate: 48000,
      channels: 1,
      encoding: 'opus'
    }
  }));
}

// Then trigger processing
ws.send(JSON.stringify({
  kind: 'interact',
  uid: crypto.randomUUID(),
  audio_output: true
}));

Audio output is also streamed in chunks for real-time playback.

Error Handling

Errors are returned as error messages:

{
  "kind": "error",
  "uid": "matching-uid",
  "error": "Invalid API key",
  "code": "AUTHENTICATION_FAILED"
}

Handle errors appropriately:

ws.onmessage = (event) => {
  const message = JSON.parse(event.data);

  if (message.kind === 'error') {
    console.error('Error:', message.error);
    // Handle the error
  }
};

ws.onerror = (error) => {
  console.error('WebSocket error:', error);
};

ws.onclose = (event) => {
  console.log('WebSocket closed:', event.code, event.reason);
};

Connection Management

Heartbeat/Keep-Alive

The WebSocket connection is kept alive automatically. No manual heartbeat is required.

Reconnection

If the connection is lost, you'll need to:

1. Establish a new WebSocket connection
2. Re-authenticate
3. Re-configure if needed
4. Resume interaction

Example reconnection logic:

let ws;
let reconnectAttempts = 0;
const MAX_RECONNECT_ATTEMPTS = 5;

function connect() {
  ws = new WebSocket('wss://pug.stg.uglabs.app/interact');

  ws.onopen = () => {
    console.log('Connected');
    reconnectAttempts = 0;
    authenticate();
  };

  ws.onclose = (event) => {
    if (reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
      reconnectAttempts++;
      console.log(`Reconnecting (attempt ${reconnectAttempts})...`);
      setTimeout(connect, 1000 * reconnectAttempts);
    }
  };
}

connect();

Best Practices

1. Always include unique UIDs - Use crypto.randomUUID() for generating unique identifiers
2. Handle streaming data - Process text/audio chunks as they arrive for real-time experience
3. Implement error handling - Always handle error messages and connection failures
4. Clean up resources - Close WebSocket connections when done
5. Respect rate limits - Don't send messages too rapidly
6. Store access tokens securely - Never expose API keys in client-side code

Example: Complete Conversation

Here's a complete example of a conversation flow:

const ws = new WebSocket('wss://pug.stg.uglabs.app/interact');

// Step 1: Get access token (do this once, reuse token)
const authResponse = await fetch('https://pug.stg.uglabs.app/api/auth/login', {
  method: 'POST',
  headers: { 'Content-Type': 'application/json' },
  body: JSON.stringify({
    api_key: 'your-api-key'
  })
});
const { access_token } = await authResponse.json();

// Step 2: Authenticate WebSocket
ws.onopen = () => {
  ws.send(JSON.stringify({
    kind: 'authenticate',
    uid: crypto.randomUUID(),
    access_token: access_token
  }));
};

// Step 3: Handle authentication success
ws.onmessage = (event) => {
  const message = JSON.parse(event.data);

  if (message.kind === 'authenticated') {
    // Step 4: Configure conversation
    ws.send(JSON.stringify({
      kind: 'set_configuration',
      uid: crypto.randomUUID(),
      prompt: 'You are a helpful assistant.'
    }));
  }

  if (message.kind === 'configured') {
    // Step 5: Start interacting
    ws.send(JSON.stringify({
      kind: 'interact',
      uid: crypto.randomUUID(),
      text: 'Hello!',
      audio_output: false
    }));
  }

  if (message.kind === 'text') {
    // Step 6: Receive response
    console.log('Assistant:', message.text);
  }
};

Try It Yourself

Use our Interactive API Tester to experiment with the WebSocket protocol in real-time.

Next Steps

• Message Type Reference
• Utilities Documentation
• SDK Quickstart Guide