@livekit/agents 1.1.0 → 1.2.1

package/src/inference/interruption/utils.ts

@@ -0,0 +1,137 @@
+// SPDX-FileCopyrightText: 2026 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { FRAME_DURATION_IN_S, MIN_INTERRUPTION_DURATION_IN_S } from './defaults.js';
+
+/**
+ * A bounded cache that automatically evicts the oldest entries when the cache exceeds max size.
+ * Uses FIFO eviction strategy.
+ */
+export class BoundedCache<K, V extends object> {
+  private cache: Map<K, V> = new Map();
+  private readonly maxLen: number;
+
+  constructor(maxLen: number = 10) {
+    this.maxLen = maxLen;
+  }
+
+  set(key: K, value: V): void {
+    this.cache.set(key, value);
+    if (this.cache.size > this.maxLen) {
+      // Remove the oldest entry (first inserted)
+      const firstKey = this.cache.keys().next().value as K;
+      this.cache.delete(firstKey);
+    }
+  }
+
+  /**
+   * Update existing value fields if present and defined.
+   * Mirrors python BoundedDict.update_value behavior.
+   */
+  updateValue(key: K, fields: Partial<V>): V | undefined {
+    const value = this.cache.get(key);
+    if (!value) return value;
+
+    for (const [fieldName, fieldValue] of Object.entries(fields) as [keyof V, V[keyof V]][]) {
+      if (fieldValue === undefined) continue;
+      // Runtime field update parity with python's hasattr + setattr.
+      if (fieldName in (value as object)) {
+        (value as Record<string, unknown>)[String(fieldName)] = fieldValue;
+      }
+    }
+    return value;
+  }
+
+  /**
+   * Set a new value with factory when missing; otherwise update in place.
+   * Mirrors python BoundedDict.set_or_update behavior.
+   */
+  setOrUpdate(key: K, factory: () => V, fields: Partial<V>): V {
+    if (!this.cache.has(key)) {
+      this.set(key, factory());
+    }
+    const result = this.updateValue(key, fields);
+    if (!result) {
+      throw new Error('setOrUpdate invariant failed: entry should exist after set');
+    }
+    return result;
+  }
+
+  get(key: K): V | undefined {
+    return this.cache.get(key);
+  }
+
+  has(key: K): boolean {
+    return this.cache.has(key);
+  }
+
+  delete(key: K): boolean {
+    return this.cache.delete(key);
+  }
+
+  /**
+   * Pop an entry if it satisfies the predicate.
+   * - No predicate: pop oldest (FIFO)
+   * - With predicate: search in reverse order and pop first match
+   */
+  pop(predicate?: (value: V) => boolean): V | undefined {
+    if (predicate === undefined) {
+      const first = this.cache.entries().next().value as [K, V] | undefined;
+      if (!first) return undefined;
+      const [key, value] = first;
+      this.cache.delete(key);
+      return value;
+    }
+
+    const keys = Array.from(this.cache.keys());
+    for (let i = keys.length - 1; i >= 0; i--) {
+      const key = keys[i]!;
+      const value = this.cache.get(key)!;
+      if (predicate(value)) {
+        this.cache.delete(key);
+        return value;
+      }
+    }
+    return undefined;
+  }
+
+  clear(): void {
+    this.cache.clear();
+  }
+
+  get size(): number {
+    return this.cache.size;
+  }
+
+  values(): IterableIterator<V> {
+    return this.cache.values();
+  }
+
+  keys(): IterableIterator<K> {
+    return this.cache.keys();
+  }
+
+  entries(): IterableIterator<[K, V]> {
+    return this.cache.entries();
+  }
+}
+
+/**
+ * Estimate probability by finding the n-th maximum value in the probabilities array.
+ * The n-th position is determined by the window size (25ms per frame).
+ * Returns 0 if there are insufficient probabilities.
+ */
+export function estimateProbability(
+  probabilities: number[],
+  windowSizeInS: number = MIN_INTERRUPTION_DURATION_IN_S,
+): number {
+  const nTh = Math.ceil(windowSizeInS / FRAME_DURATION_IN_S);
+  if (probabilities.length < nTh) {
+    return 0;
+  }
+
+  // Find the n-th maximum value by sorting in descending order
+  // Create a copy to avoid mutating the original array
+  const sorted = [...probabilities].sort((a, b) => b - a);
+  return sorted[nTh - 1]!;
+}

package/src/inference/interruption/ws_transport.ts

@@ -0,0 +1,416 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { Throws } from '@livekit/throws-transformer/throws';
+import { TransformStream } from 'stream/web';
+import WebSocket from 'ws';
+import { z } from 'zod';
+import { APIConnectionError, APIStatusError, APITimeoutError } from '../../_exceptions.js';
+import { log } from '../../log.js';
+import TypedPromise from '../../typed_promise.js';
+import { createAccessToken } from '../utils.js';
+import { InterruptionCacheEntry } from './interruption_cache_entry.js';
+import type { OverlappingSpeechEvent } from './types.js';
+import type { BoundedCache } from './utils.js';
+
+// WebSocket message types
+const MSG_SESSION_CREATE = 'session.create';
+const MSG_SESSION_CLOSE = 'session.close';
+const MSG_SESSION_CREATED = 'session.created';
+const MSG_SESSION_CLOSED = 'session.closed';
+const MSG_INTERRUPTION_DETECTED = 'bargein_detected';
+const MSG_INFERENCE_DONE = 'inference_done';
+const MSG_ERROR = 'error';
+
+export interface WsTransportOptions {
+  baseUrl: string;
+  apiKey: string;
+  apiSecret: string;
+  sampleRate: number;
+  threshold: number;
+  minFrames: number;
+  timeout: number;
+  maxRetries?: number;
+}
+
+export interface WsTransportState {
+  overlapSpeechStarted: boolean;
+  overlapSpeechStartedAt: number | undefined;
+  cache: BoundedCache<number, InterruptionCacheEntry>;
+}
+
+const wsMessageSchema = z.discriminatedUnion('type', [
+  z.object({
+    type: z.literal(MSG_SESSION_CREATED),
+  }),
+  z.object({
+    type: z.literal(MSG_SESSION_CLOSED),
+  }),
+  z.object({
+    type: z.literal(MSG_INTERRUPTION_DETECTED),
+    created_at: z.number(),
+    probabilities: z.array(z.number()).default([]),
+    prediction_duration: z.number().default(0),
+  }),
+  z.object({
+    type: z.literal(MSG_INFERENCE_DONE),
+    created_at: z.number(),
+    probabilities: z.array(z.number()).default([]),
+    prediction_duration: z.number().default(0),
+    is_bargein: z.boolean().optional(),
+  }),
+  z.object({
+    type: z.literal(MSG_ERROR),
+    message: z.string(),
+    code: z.number().optional(),
+    session_id: z.string().optional(),
+  }),
+]);
+
+type WsMessage = z.infer<typeof wsMessageSchema>;
+
+/**
+ * Creates a WebSocket connection and waits for it to open.
+ */
+async function connectWebSocket(
+  options: WsTransportOptions,
+): Promise<Throws<WebSocket, APIStatusError | APITimeoutError | APIConnectionError>> {
+  const baseUrl = options.baseUrl.replace(/^http/, 'ws');
+  const token = await createAccessToken(options.apiKey, options.apiSecret);
+  const url = `${baseUrl}/bargein`;
+
+  const ws = new WebSocket(url, {
+    headers: { Authorization: `Bearer ${token}` },
+  });
+
+  await new TypedPromise<void, APIStatusError | APITimeoutError | APIConnectionError>(
+    (resolve, reject) => {
+      const timeout = setTimeout(() => {
+        ws.terminate();
+        reject(
+          new APITimeoutError({
+            message: 'WebSocket connection timeout',
+            options: { retryable: false },
+          }),
+        );
+      }, options.timeout);
+      ws.once('open', () => {
+        clearTimeout(timeout);
+        resolve();
+      });
+      ws.once('unexpected-response', (_req, res) => {
+        clearTimeout(timeout);
+        ws.terminate();
+        const statusCode = res.statusCode ?? -1;
+        reject(
+          new APIStatusError({
+            message: `WebSocket connection rejected with status ${statusCode}`,
+            options: { statusCode, retryable: false },
+          }),
+        );
+      });
+      ws.once('error', (err: Error) => {
+        clearTimeout(timeout);
+        ws.terminate();
+        reject(new APIConnectionError({ message: `WebSocket connection error: ${err.message}` }));
+      });
+    },
+  );
+
+  return ws;
+}
+
+export interface WsTransportResult {
+  transport: TransformStream<Int16Array | OverlappingSpeechEvent, OverlappingSpeechEvent>;
+  reconnect: () => Promise<void>;
+}
+
+/**
+ * Creates a WebSocket transport TransformStream for interruption detection.
+ *
+ * This transport receives Int16Array audio slices and outputs InterruptionEvents.
+ * It maintains a persistent WebSocket connection with automatic retry on failure.
+ * Returns both the transport and a reconnect function for option updates.
+ */
+export function createWsTransport(
+  options: WsTransportOptions,
+  getState: () => WsTransportState,
+  setState: (partial: Partial<WsTransportState>) => void,
+  updateUserSpeakingSpan?: (entry: InterruptionCacheEntry) => void,
+  onRequestSent?: () => void,
+  getAndResetNumRequests?: () => number,
+): WsTransportResult {
+  const logger = log();
+  let ws: WebSocket | null = null;
+  let outputController: TransformStreamDefaultController<OverlappingSpeechEvent> | null = null;
+
+  function setupMessageHandler(socket: WebSocket): void {
+    socket.on('message', (data: WebSocket.Data) => {
+      try {
+        const message = wsMessageSchema.parse(JSON.parse(data.toString()));
+        handleMessage(message);
+      } catch {
+        logger.warn({ data: data.toString() }, 'Failed to parse WebSocket message');
+      }
+    });
+
+    socket.on('error', (err: Error) => {
+      outputController?.error(
+        new APIConnectionError({ message: `WebSocket error: ${err.message}` }),
+      );
+    });
+
+    socket.on('close', (code: number, reason: Buffer) => {
+      logger.debug({ code, reason: reason.toString() }, 'WebSocket closed');
+    });
+  }
+
+  async function ensureConnection(): Promise<
+    Throws<void, APIStatusError | APITimeoutError | APIConnectionError>
+  > {
+    if (ws && ws.readyState === WebSocket.OPEN) return;
+
+    ws = await connectWebSocket(options);
+    setupMessageHandler(ws);
+
+    const sessionCreateMsg = JSON.stringify({
+      type: MSG_SESSION_CREATE,
+      settings: {
+        sample_rate: options.sampleRate,
+        num_channels: 1,
+        threshold: options.threshold,
+        min_frames: options.minFrames,
+        encoding: 's16le',
+      },
+    });
+    ws.send(sessionCreateMsg);
+  }
+
+  function handleMessage(message: WsMessage): void {
+    const state = getState();
+
+    switch (message.type) {
+      case MSG_SESSION_CREATED:
+        logger.debug('WebSocket session created');
+        break;
+
+      case MSG_INTERRUPTION_DETECTED: {
+        const createdAt = message.created_at;
+        const overlapSpeechStartedAt = state.overlapSpeechStartedAt;
+        if (state.overlapSpeechStarted && overlapSpeechStartedAt !== undefined) {
+          const existing = state.cache.get(createdAt);
+
+          const totalDurationInS =
+            existing?.requestStartedAt !== undefined
+              ? (performance.now() - existing.requestStartedAt) / 1000
+              : (performance.now() - createdAt) / 1000;
+
+          const entry = state.cache.setOrUpdate(
+            createdAt,
+            () => new InterruptionCacheEntry({ createdAt }),
+            {
+              speechInput: existing?.speechInput,
+              requestStartedAt: existing?.requestStartedAt,
+              totalDurationInS,
+              probabilities: message.probabilities,
+              isInterruption: true,
+              predictionDurationInS: message.prediction_duration,
+              detectionDelayInS: (Date.now() - overlapSpeechStartedAt) / 1000,
+            },
+          );
+
+          if (updateUserSpeakingSpan) {
+            updateUserSpeakingSpan(entry);
+          }
+
+          logger.debug(
+            {
+              totalDuration: entry.totalDurationInS,
+              predictionDuration: entry.predictionDurationInS,
+              detectionDelay: entry.detectionDelayInS,
+              probability: entry.probability,
+            },
+            'interruption detected',
+          );
+
+          const event: OverlappingSpeechEvent = {
+            type: 'overlapping_speech',
+            detectedAt: Date.now(),
+            isInterruption: true,
+            totalDurationInS: entry.totalDurationInS,
+            predictionDurationInS: entry.predictionDurationInS,
+            overlapStartedAt: overlapSpeechStartedAt,
+            speechInput: entry.speechInput,
+            probabilities: entry.probabilities,
+            detectionDelayInS: entry.detectionDelayInS,
+            probability: entry.probability,
+            numRequests: getAndResetNumRequests?.() ?? 0,
+          };
+
+          outputController?.enqueue(event);
+          setState({ overlapSpeechStarted: false });
+        }
+        break;
+      }
+
+      case MSG_INFERENCE_DONE: {
+        const createdAt = message.created_at;
+        const overlapSpeechStartedAt = state.overlapSpeechStartedAt;
+        if (state.overlapSpeechStarted && overlapSpeechStartedAt !== undefined) {
+          const existing = state.cache.get(createdAt);
+          const totalDurationInS =
+            existing?.requestStartedAt !== undefined
+              ? (performance.now() - existing.requestStartedAt) / 1000
+              : (performance.now() - createdAt) / 1000;
+          const entry = state.cache.setOrUpdate(
+            createdAt,
+            () => new InterruptionCacheEntry({ createdAt }),
+            {
+              speechInput: existing?.speechInput,
+              requestStartedAt: existing?.requestStartedAt,
+              totalDurationInS,
+              predictionDurationInS: message.prediction_duration,
+              probabilities: message.probabilities,
+              isInterruption: message.is_bargein ?? false,
+              detectionDelayInS: (Date.now() - overlapSpeechStartedAt) / 1000,
+            },
+          );
+
+          logger.debug(
+            {
+              totalDurationInS: entry.totalDurationInS,
+              predictionDurationInS: entry.predictionDurationInS,
+            },
+            'interruption inference done',
+          );
+        }
+        break;
+      }
+
+      case MSG_SESSION_CLOSED:
+        logger.debug('WebSocket session closed');
+        break;
+
+      case MSG_ERROR:
+        outputController?.error(
+          new APIStatusError({
+            message: `LiveKit Adaptive Interruption error: ${message.message}`,
+            options: { statusCode: message.code ?? -1 },
+          }),
+        );
+        break;
+    }
+  }
+
+  function sendAudioData(audioSlice: Int16Array): void {
+    if (!ws || ws.readyState !== WebSocket.OPEN) {
+      throw new APIConnectionError({ message: 'WebSocket not connected' });
+    }
+
+    const state = getState();
+    const createdAt = Math.floor(performance.now());
+
+    state.cache.set(
+      createdAt,
+      new InterruptionCacheEntry({
+        createdAt,
+        requestStartedAt: performance.now(),
+        speechInput: audioSlice,
+      }),
+    );
+
+    const header = new ArrayBuffer(8);
+    const view = new DataView(header);
+    view.setUint32(0, createdAt >>> 0, true);
+    view.setUint32(4, Math.floor(createdAt / 0x100000000) >>> 0, true);
+
+    const audioBytes = new Uint8Array(
+      audioSlice.buffer,
+      audioSlice.byteOffset,
+      audioSlice.byteLength,
+    );
+    const combined = new Uint8Array(8 + audioBytes.length);
+    combined.set(new Uint8Array(header), 0);
+    combined.set(audioBytes, 8);
+
+    ws.send(combined);
+    onRequestSent?.();
+  }
+
+  function close(): void {
+    if (ws?.readyState === WebSocket.OPEN) {
+      const closeMsg = JSON.stringify({ type: MSG_SESSION_CLOSE });
+      try {
+        ws.send(closeMsg);
+      } catch (e: unknown) {
+        logger.error(e, 'failed to send close message');
+      }
+    }
+    ws?.close(1000); // signal normal websocket closure
+    ws = null;
+  }
+
+  /**
+   * Reconnect the WebSocket with updated options.
+   * This is called when options are updated via updateOptions().
+   */
+  async function reconnect(): Promise<void> {
+    close();
+  }
+
+  const transport = new TransformStream<
+    Int16Array | OverlappingSpeechEvent,
+    OverlappingSpeechEvent
+  >(
+    {
+      async start(controller) {
+        outputController = controller;
+        await ensureConnection().catch((e) => {
+          controller.error(e);
+        });
+      },
+
+      transform(chunk, controller) {
+        if (!(chunk instanceof Int16Array)) {
+          controller.enqueue(chunk);
+          return;
+        }
+
+        // Only forwards buffered audio while overlap speech is actively on.
+        const state = getState();
+        if (!state.overlapSpeechStartedAt || !state.overlapSpeechStarted) return;
+
+        if (options.timeout > 0) {
+          const now = performance.now();
+          for (const [, entry] of state.cache.entries()) {
+            if (entry.totalDurationInS !== 0) continue;
+            if (now - entry.createdAt > options.timeout) {
+              controller.error(
+                new APIStatusError({
+                  message: `interruption inference timed out after ${((now - entry.createdAt) / 1000).toFixed(1)}s (ws)`,
+                  options: { statusCode: 408, retryable: false },
+                }),
+              );
+              return;
+            }
+            break;
+          }
+        }
+
+        try {
+          sendAudioData(chunk);
+        } catch (err) {
+          controller.error(err);
+        }
+      },
+
+      flush() {
+        close();
+      },
+    },
+    { highWaterMark: 2 },
+    { highWaterMark: 2 },
+  );
+
+  return { transport, reconnect };
+}