Websocket Tutorial

WebSocket API Tutorial

Welcome to the WebSocket API tutorial! This guide will walk you through how to connect and interact with our WebSocket service step by step.

Prerequisites

Before you start, you'll need:

Connection Flow

The connection process happens in two steps:

  1. Get an access token using your login token
  2. Establish WebSocket connection using the access token

Step 1: Getting an Access Token

First, you need to exchange your login token for an access token:

async function getAccessToken(loginToken) {
  const response = await fetch("API_BASE_URL/websocket/get_access_token", {
    method: "GET",
    headers: {
      Authorization: `Bearer ${loginToken}`,
      "Content-Type": "application/json",
    },
  })

  const data = await response.json()
  return data.access_token
}

Note: In the above function, replace API_BASE_URL with the actual URL of your API.

Step 2: Establishing WebSocket Connection

Once you have the access token, you can establish the WebSocket connection:

const ws = new WebSocket(
  `wss://API_BASE_URL/websocket/ws?access_token=${accessToken}`
)

Handling WebSocket Events

Connection Events

ws.onopen = () => {
  console.log("Connected to WebSocket")
  // Enable your chat interface here
}

ws.onclose = () => {
  console.log("Disconnected from WebSocket")
  // Handle disconnection here
}

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

Sending Messages

To send a message to the API, use this format:

const message = {
  type: "conversation_with_aileen",
  message: "Your message here",
  conversation_id: "unique-conversation-id", // Use UUID
  request_id: "unique-request-id", // Use UUID for each message
}

ws.send(JSON.stringify(message))

Receiving Messages

The server sends responses in a streaming format, where each message is a chunk of the complete response. When you receive a message of type conversation_with_aileen, the response will have the following structure:

{
  "is_final": false,  // Indicates if this is the last chunk of the response
  "data": {
    "meta": {
      "conversation_id": "2617cb04-927e-4de8-a1f2-f7d8e288bea6"  // Conversation tracking ID
    },
    "follow_up_questions": null,  // Optional follow-up questions that may be present in the final chunk
    "delta": " procurando",  // The current chunk of the message
    "delta_generation_finished": false  // Indicates if the bot has finished generating the response
  },
  "request_id": "4c84baa7-8d4c-4b1d-8d7f-6fd66f0da4c2"  // ID to track this specific request
}

To properly handle these streaming responses, you need to:

  1. Track the current message state
  2. Concatenate each delta to build the complete message
  3. Update the UI as new chunks arrive
  4. Handle the final state when the message is complete

Here's an example implementation:

// State management for message tracking
const state = {
  currentRequestId: null,
  currentMessageContent: "",
  currentMessageDiv: null,
}

// Handle incoming messages
ws.onmessage = (event) => {
  const data = JSON.parse(event.data)

  if (data.type === "ping") {
    ws.send(JSON.stringify({ type: "pong" }))
    return
  }

  if (data.type === "conversation_with_aileen") {
    const response = data.response
    if (!response?.data) return

    const { delta, delta_generation_finished } = response.data
    const requestId = response.request_id

    if (delta) {
      if (requestId !== state.currentRequestId) {
        // New message starting
        state.currentRequestId = requestId
        state.currentMessageContent = delta
        state.currentMessageDiv = createNewMessageElement(
          state.currentMessageContent
        )
      } else {
        // Update existing message
        state.currentMessageContent += delta
        updateMessageElement(
          state.currentMessageDiv,
          state.currentMessageContent
        )
      }
    }

    if (delta_generation_finished) {
      // Message is complete, clean up state
      state.currentRequestId = null
      state.currentMessageDiv = null
      state.currentMessageContent = ""

      // Handle any follow-up questions if present
      if (response.data.follow_up_questions) {
        displayFollowUpQuestions(response.data.follow_up_questions)
      }
    }
  }
}

// Helper functions for UI updates
function createNewMessageElement(content) {
  const messageDiv = document.createElement("div")
  messageDiv.className = "message aileen"
  messageDiv.textContent = content
  document.getElementById("messages").appendChild(messageDiv)
  return messageDiv
}

function updateMessageElement(element, content) {
  if (element) {
    element.textContent = content
    // Optionally parse markdown or format the content
    // element.innerHTML = marked.parse(content)
  }
}

function displayFollowUpQuestions(questions) {
  if (!questions || !questions.length) return

  const questionsContainer = document.createElement("div")
  questionsContainer.className = "follow-up-questions"

  questions.forEach((question) => {
    const questionButton = document.createElement("button")
    questionButton.textContent = question
    questionButton.onclick = () => sendMessage(question)
    questionsContainer.appendChild(questionButton)
  })

  document.getElementById("messages").appendChild(questionsContainer)
}

This implementation provides:

The delta_generation_finished flag is crucial as it indicates when the bot has finished generating its response. This is your signal to:

  1. Clean up any temporary state
  2. Finalize the message display
  3. Handle any follow-up questions
  4. Enable UI elements for the next interaction

Note: The actual UI implementation functions (createNewMessageElement, updateMessageElement, etc.) should be adapted to match your application's specific needs and styling requirements.

Heartbeat Mechanism

In many WebSocket implementations, a heartbeat is used to ensure the connection is still active. The server periodically sends a 'ping' message, and the client must respond with a 'pong' message. This mechanism helps detect connection issues early and maintain a stable connection.

Important: If the client fails to respond with a 'pong' message within 60 seconds after receiving a 'ping', the server will automatically disconnect the WebSocket connection. Make sure to implement the heartbeat response properly to maintain a stable connection.

For example:

if (data.type === "ping") {
  ws.send(JSON.stringify({ type: "pong" }))
  return
}

Implementing this heartbeat mechanism is crucial for keeping your WebSocket connection alive.

Handling Inactivity and Reconnection

The server will disconnect the WebSocket connection after a period of inactivity. To handle this gracefully, you should implement a reconnection mechanism. Here's an example of how to implement reconnection logic with exponential backoff:

// Configuration
const CONFIG = {
  maxReconnectAttempts: 5,
  initialReconnectDelay: 1000, // 1 second
}

// State management
const state = {
  ws: null,
  reconnectAttempts: 0,
  reconnectDelay: CONFIG.initialReconnectDelay,
}

// Websocket connection function
async function connect(accessToken) {
  if (state.ws) {
    state.ws.close()
  }

  state.ws = new WebSocket(
    `wss://API_BASE_URL/websocket/ws?access_token=${accessToken}`
  )

  state.ws.onclose = (event) => {
    console.log("WebSocket disconnected", event)
    updateConnectionStatus(false) // Update UI to show disconnected state

    if (state.reconnectAttempts < CONFIG.maxReconnectAttempts) {
      setTimeout(() => {
        console.log(`Reconnecting... Attempt ${state.reconnectAttempts + 1}`)
        state.reconnectAttempts++
        state.reconnectDelay *= 2 // Exponential backoff
        connect(accessToken)
      }, state.reconnectDelay)
    } else {
      console.error(
        "Connection failed after multiple attempts. Please refresh the page."
      )
      // Show error message to user
      showErrorMessage(
        "Connection failed after multiple attempts. Please refresh the page."
      )
    }
  }

  state.ws.onopen = () => {
    console.log("Connected to WebSocket")
    updateConnectionStatus(true) // Update UI to show connected state
    state.reconnectAttempts = 0
    state.reconnectDelay = CONFIG.initialReconnectDelay
  }

  state.ws.onerror = (error) => {
    console.error("WebSocket error:", error)
  }
}

This implementation includes several important features:

  1. Configuration Constants

    • maxReconnectAttempts: Maximum number of reconnection attempts (5 in this example)
    • initialReconnectDelay: Initial delay between reconnection attempts (1 second)
  2. State Management

    • Keeps track of the WebSocket instance
    • Tracks number of reconnection attempts
    • Manages reconnection delay for exponential backoff
  3. Exponential Backoff

    • Doubles the delay between each reconnection attempt
    • Helps prevent overwhelming the server with rapid reconnection attempts
  4. Connection Reset

    • Resets reconnection counters on successful connection
    • Closes existing connection before creating a new one
  5. Error Handling

    • Logs connection errors
    • Updates UI to show connection status
    • Stops reconnection attempts after maximum attempts reached

To use this reconnection logic, you would call it like this:

try {
  const accessToken = await getAccessToken(loginToken)
  await connect(accessToken)
} catch (error) {
  console.error("Failed to connect:", error)
  showErrorMessage(`Failed to connect: ${error.message}`)
}

Best Practices

  1. Keep Track of Connection State

    • Monitor connection status
    • Implement reconnection logic
    • Show connection status to users
  2. Handle Messages Properly

    • Keep track of conversation IDs
    • Generate unique request IDs for each message
    • Handle streaming responses appropriately
  3. Error Handling

    • Implement proper error handling
    • Show meaningful error messages to users
    • Log errors for debugging

Example Implementation

Note: In this example, crypto.randomUUID() is used to generate unique IDs, which is available in modern browsers. If it's not supported in your environment, consider using an alternative UUID generator. Also, the functions displayMessage() and finalizeMessage() are placeholders; implement them according to your UI needs.

// Generate UUIDs for tracking
const conversationId = crypto.randomUUID()
const requestId = crypto.randomUUID()

// Prepare message
const message = {
  type: "conversation_with_aileen",
  message: "Hello, how are you?",
  conversation_id: conversationId,
  request_id: requestId,
}

// Send message
if (ws.readyState === WebSocket.OPEN) {
  ws.send(JSON.stringify(message))
}

// Track message state
let currentMessage = ""

// Handle response
ws.onmessage = (event) => {
  const data = JSON.parse(event.data)

  if (data.type === "conversation_with_aileen") {
    const response = data.response

    if (response?.data?.delta) {
      currentMessage += response.data.delta
      displayMessage(currentMessage)
    }

    if (response?.data?.delta_generation_finished) {
      finalizeMessage(currentMessage)
    }
  }
}

Troubleshooting

Common issues and solutions:

  1. Connection Failed

    • Check if your login token is valid
    • Verify the API endpoint URL
    • Check your internet connection
  2. Messages Not Sending

    • Ensure WebSocket connection is open
    • Verify message format
    • Check for console errors
  3. Not Receiving Responses

    • Verify message format
    • Check connection status
    • Ensure proper event handling

Need Help?

If you encounter any issues:

Complete Working Example

Here's a complete, working example that you can use as a starting point for your implementation:

<!DOCTYPE html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>Aileen Chat Example</title>
    <style>
      body {
        font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto,
          sans-serif;
        max-width: 800px;
        margin: 0 auto;
        padding: 20px;
      }

      #chat-container {
        border: 1px solid #ccc;
        border-radius: 8px;
        padding: 20px;
        margin-top: 20px;
      }

      #messages {
        height: 400px;
        overflow-y: auto;
        margin-bottom: 20px;
        padding: 10px;
        background: #f9f9f9;
        border-radius: 4px;
      }

      .message {
        margin: 8px 0;
        padding: 8px 12px;
        border-radius: 4px;
      }

      .message.user {
        background: #e3f2fd;
        margin-left: 20%;
        margin-right: 0;
      }

      .message.aileen {
        background: #f5f5f5;
        margin-right: 20%;
        margin-left: 0;
      }

      .follow-up-questions {
        margin-top: 10px;
        display: flex;
        flex-direction: column;
        gap: 8px;
      }

      .follow-up-questions button {
        background: #e3f2fd;
        border: 1px solid #bbdefb;
        padding: 8px 12px;
        border-radius: 4px;
        cursor: pointer;
        text-align: left;
      }

      .follow-up-questions button:hover {
        background: #bbdefb;
      }

      #message-form {
        display: flex;
        gap: 10px;
      }

      #message-input {
        flex: 1;
        padding: 8px;
        border: 1px solid #ccc;
        border-radius: 4px;
      }

      button {
        padding: 8px 16px;
        background: #2196f3;
        color: white;
        border: none;
        border-radius: 4px;
        cursor: pointer;
      }

      button:hover {
        background: #1976d2;
      }

      button:disabled {
        background: #ccc;
        cursor: not-allowed;
      }

      #connection-status {
        margin-bottom: 10px;
        padding: 8px;
        border-radius: 4px;
      }

      #connection-status.connected {
        background: #c8e6c9;
        color: #2e7d32;
      }

      #connection-status.disconnected {
        background: #ffcdd2;
        color: #c62828;
      }

      #auth-container {
        margin-bottom: 20px;
      }

      .token-input {
        width: 100%;
        padding: 8px;
        margin-bottom: 10px;
        border: 1px solid #ccc;
        border-radius: 4px;
      }
    </style>
  </head>
  <body>
    <div id="auth-container">
      <input
        type="password"
        id="token-input"
        class="token-input"
        placeholder="Enter your login token"
      />
      <button id="connect-btn">Connect</button>
    </div>

    <div id="chat-container">
      <div id="connection-status" class="disconnected">Disconnected</div>
      <div id="messages"></div>
      <form id="message-form">
        <input
          type="text"
          id="message-input"
          placeholder="Type your message..."
          required
          disabled
        />
        <button type="submit" disabled>Send</button>
      </form>
    </div>

    <script>
      // Configuration
      const CONFIG = {
        maxReconnectAttempts: 5,
        initialReconnectDelay: 1000,
        baseUrl: "API_BASE_URL", // Replace with your API base URL
      }

      // State management
      const state = {
        ws: null,
        reconnectAttempts: 0,
        reconnectDelay: CONFIG.initialReconnectDelay,
        conversationId: crypto.randomUUID(),
        currentRequestId: null,
        currentMessageContent: "",
        currentMessageDiv: null,
      }

      // DOM Elements
      const elements = {
        messages: document.getElementById("messages"),
        messageForm: document.getElementById("message-form"),
        messageInput: document.getElementById("message-input"),
        connectionStatus: document.getElementById("connection-status"),
        tokenInput: document.getElementById("token-input"),
        connectBtn: document.getElementById("connect-btn"),
        sendBtn: document.querySelector("#message-form button"),
      }

      // UI Helper Functions
      const ui = {
        updateConnectionStatus(isConnected) {
          elements.connectionStatus.textContent = isConnected
            ? "Connected"
            : "Disconnected"
          elements.connectionStatus.className = isConnected
            ? "connected"
            : "disconnected"
          elements.messageInput.disabled = !isConnected
          elements.sendBtn.disabled = !isConnected
        },

        addMessage(text, type = "user") {
          const messageDiv = document.createElement("div")
          messageDiv.className = `message ${type}`
          messageDiv.textContent = text
          elements.messages.appendChild(messageDiv)
          elements.messages.scrollTop = elements.messages.scrollHeight
          return messageDiv
        },

        updateMessageContent(messageDiv, content) {
          if (messageDiv) {
            messageDiv.textContent = content
            elements.messages.scrollTop = elements.messages.scrollHeight
          }
        },

        displayFollowUpQuestions(questions) {
          if (!questions?.length) return

          const container = document.createElement("div")
          container.className = "follow-up-questions"

          questions.forEach((question) => {
            const button = document.createElement("button")
            button.textContent = question
            button.onclick = () => {
              elements.messageInput.value = question
              elements.messageForm.dispatchEvent(new Event("submit"))
            }
            container.appendChild(button)
          })

          elements.messages.appendChild(container)
          elements.messages.scrollTop = elements.messages.scrollHeight
        },
      }

      // WebSocket Functions
      const websocket = {
        async connect(token) {
          if (state.ws) {
            state.ws.close()
          }

          state.ws = new WebSocket(
            `${CONFIG.baseUrl.replace(
              "http",
              "ws"
            )}/websocket/ws?access_token=${token}`
          )

          state.ws.onopen = () => {
            console.log("Connected to WebSocket")
            ui.updateConnectionStatus(true)
            state.reconnectAttempts = 0
            state.reconnectDelay = CONFIG.initialReconnectDelay
          }

          state.ws.onmessage = this.handleMessage
          state.ws.onclose = this.handleClose
          state.ws.onerror = (error) => console.error("WebSocket error:", error)
        },

        handleMessage(event) {
          const data = JSON.parse(event.data)

          if (data.type === "ping") {
            state.ws.send(JSON.stringify({ type: "pong" }))
            return
          }

          if (data.type === "conversation_with_aileen") {
            const response = data.response
            if (!response?.data) return

            const { delta, delta_generation_finished } = response.data
            const requestId = response.request_id

            if (delta) {
              if (requestId !== state.currentRequestId) {
                state.currentRequestId = requestId
                state.currentMessageContent = delta
                state.currentMessageDiv = ui.addMessage(
                  state.currentMessageContent,
                  "aileen"
                )
              } else {
                state.currentMessageContent += delta
                ui.updateMessageContent(
                  state.currentMessageDiv,
                  state.currentMessageContent
                )
              }
            }

            if (delta_generation_finished) {
              if (response.data.follow_up_questions) {
                ui.displayFollowUpQuestions(response.data.follow_up_questions)
              }
              state.currentRequestId = null
              state.currentMessageDiv = null
              state.currentMessageContent = ""
            }
          }
        },

        handleClose(event) {
          console.log("WebSocket disconnected", event)
          ui.updateConnectionStatus(false)

          if (state.reconnectAttempts < CONFIG.maxReconnectAttempts) {
            setTimeout(() => {
              console.log(
                `Reconnecting... Attempt ${state.reconnectAttempts + 1}`
              )
              state.reconnectAttempts++
              state.reconnectDelay *= 2 // Exponential backoff
              auth
                .getAccessToken(elements.tokenInput.value.trim())
                .then((token) => websocket.connect(token))
                .catch(console.error)
            }, state.reconnectDelay)
          }
        },
      }

      // Authentication Functions
      const auth = {
        async getAccessToken(loginToken) {
          try {
            const response = await fetch(
              `${CONFIG.baseUrl}/websocket/get_access_token`,
              {
                method: "GET",
                headers: {
                  Authorization: `Bearer ${loginToken}`,
                  "Content-Type": "application/json",
                },
              }
            )

            if (!response.ok) {
              throw new Error(`Failed to get access token: ${response.status}`)
            }

            const data = await response.json()
            return data.access_token
          } catch (error) {
            console.error("Error getting access token:", error)
            throw error
          }
        },
      }

      // Event Listeners
      elements.connectBtn.addEventListener("click", async () => {
        const loginToken = elements.tokenInput.value.trim()
        if (!loginToken) {
          alert("Please enter your login token")
          return
        }

        try {
          elements.connectBtn.disabled = true
          const accessToken = await auth.getAccessToken(loginToken)
          await websocket.connect(accessToken)
        } catch (error) {
          alert(`Failed to connect: ${error.message}`)
          elements.connectBtn.disabled = false
        }
      })

      elements.messageForm.addEventListener("submit", (e) => {
        e.preventDefault()
        const message = elements.messageInput.value.trim()
        if (message && state.ws?.readyState === WebSocket.OPEN) {
          const payload = {
            type: "conversation_with_aileen",
            message,
            conversation_id: state.conversationId,
            request_id: crypto.randomUUID(),
          }
          state.ws.send(JSON.stringify(payload))
          ui.addMessage(message, "user")
          elements.messageInput.value = ""
        }
      })
    </script>
  </body>
</html>

This example provides a complete, working chat interface with:

To use this example:

  1. Replace API_BASE_URL with your actual API base URL
  2. Host the HTML file on a web server
  3. Open in a browser and enter your login token
  4. Start chatting with Aileen

The interface includes:

This example can be used as a foundation for building more complex chat interfaces with additional features and customizations.