@livekit/agents-plugin-assemblyai 1.4.9 → 1.5.0

package/src/stt-metadata.test.ts

@@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: 2026 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { stt as sttLib } from '@livekit/agents';
+import { AudioFrame } from '@livekit/rtc-node';
+import { once } from 'node:events';
+import type { AddressInfo } from 'node:net';
+import { describe, expect, it } from 'vitest';
+import { WebSocketServer } from 'ws';
+import { STT } from './stt.js';
+
+function makeFrame(samplesPerChannel = 800, sampleRate = 16000): AudioFrame {
+  const data = new Int16Array(samplesPerChannel);
+  data.fill(1);
+  return new AudioFrame(data, sampleRate, 1, samplesPerChannel);
+}
+
+async function startWebSocketServer() {
+  const wss = new WebSocketServer({ host: '127.0.0.1', port: 0 });
+  await once(wss, 'listening');
+  const address = wss.address() as AddressInfo;
+  return { wss, baseUrl: `ws://127.0.0.1:${address.port}` };
+}
+
+async function closeWebSocketServer(wss: WebSocketServer): Promise<void> {
+  await new Promise<void>((resolve) => wss.close(() => resolve()));
+}
+
+async function waitUntil(predicate: () => boolean, timeoutMs = 1000): Promise<void> {
+  const startedAt = Date.now();
+  while (Date.now() - startedAt < timeoutMs) {
+    if (predicate()) return;
+    await new Promise((resolve) => setTimeout(resolve, 10));
+  }
+  throw new Error('timed out waiting for condition');
+}
+
+async function collectUntilEnd(stream: sttLib.SpeechStream): Promise<sttLib.SpeechEvent[]> {
+  const events: sttLib.SpeechEvent[] = [];
+  for await (const event of stream) {
+    events.push(event);
+    if (event.type === sttLib.SpeechEventType.END_OF_SPEECH) break;
+  }
+  return events;
+}
+
+describe('AssemblyAI STT metadata', () => {
+  it('maps turn confidence fields onto speech data metadata', async () => {
+    const { wss, baseUrl } = await startWebSocketServer();
+    let requestUrl = '';
+
+    wss.on('connection', (ws, req) => {
+      requestUrl = req.url ?? '';
+      ws.on('message', () => {
+        ws.send(
+          JSON.stringify({
+            type: 'Turn',
+            transcript: 'hola mundo',
+            utterance: 'hola mundo',
+            end_of_turn: true,
+            language_code: 'es',
+            language_confidence: 0.94,
+            words: [
+              { text: 'hola', start: 0, end: 200, confidence: 0.96 },
+              { text: 'mundo', start: 200, end: 500, confidence: 0.98 },
+            ],
+          }),
+        );
+      });
+    });
+
+    try {
+      const assemblyai = new STT({
+        apiKey: 'test-key',
+        baseUrl,
+        speechModel: 'u3-rt-pro',
+      });
+      const stream = assemblyai.stream({
+        connOptions: { maxRetry: 0, retryIntervalMs: 1, timeoutMs: 1000 },
+      });
+
+      await waitUntil(() => requestUrl !== '');
+
+      stream.pushFrame(makeFrame());
+      stream.endInput();
+
+      const events = await collectUntilEnd(stream);
+      stream.close();
+
+      expect(new URL(`ws://127.0.0.1${requestUrl}`).pathname).toBe('/v3/ws');
+      expect(
+        events
+          .filter((event) => event.alternatives?.[0])
+          .map((event) => ({
+            type: event.type,
+            language: event.alternatives?.[0]?.language,
+            metadata: event.alternatives?.[0]?.metadata,
+          })),
+      ).toEqual([
+        {
+          type: sttLib.SpeechEventType.INTERIM_TRANSCRIPT,
+          language: 'es',
+          metadata: {
+            assemblyai: {
+              languageConfidence: 0.94,
+            },
+          },
+        },
+        {
+          type: sttLib.SpeechEventType.PREFLIGHT_TRANSCRIPT,
+          language: 'es',
+          metadata: {
+            assemblyai: {
+              languageConfidence: 0.94,
+            },
+          },
+        },
+        {
+          type: sttLib.SpeechEventType.FINAL_TRANSCRIPT,
+          language: 'es',
+          metadata: {
+            assemblyai: {
+              languageConfidence: 0.94,
+            },
+          },
+        },
+      ]);
+    } finally {
+      await closeWebSocketServer(wss);
+    }
+  });
+});

package/src/stt.test.ts

@@ -3,9 +3,102 @@
 // SPDX-License-Identifier: Apache-2.0
 import { VAD } from '@livekit/agents-plugin-silero';
 import { stt } from '@livekit/agents-plugins-test';
-import { describe, it } from 'vitest';
+import { once } from 'node:events';
+import type { AddressInfo } from 'node:net';
+import { describe, expect, it } from 'vitest';
+import { WebSocketServer } from 'ws';
 import { STT } from './stt.js';
 
+async function startWebSocketServer() {
+  const wss = new WebSocketServer({ host: '127.0.0.1', port: 0 });
+  await once(wss, 'listening');
+  const address = wss.address() as AddressInfo;
+  return { wss, baseUrl: `ws://127.0.0.1:${address.port}` };
+}
+
+async function closeWebSocketServer(wss: WebSocketServer): Promise<void> {
+  for (const client of wss.clients) client.close();
+  await new Promise<void>((resolve) => wss.close(() => resolve()));
+}
+
+async function waitUntil(predicate: () => boolean, timeoutMs = 1000): Promise<void> {
+  const startedAt = Date.now();
+  while (Date.now() - startedAt < timeoutMs) {
+    if (predicate()) return;
+    await new Promise((resolve) => setTimeout(resolve, 10));
+  }
+  throw new Error('timed out waiting for condition');
+}
+
+describe('AssemblyAI options', () => {
+  it('accepts u3-rt-pro-beta-1', () => {
+    const stt = new STT({ apiKey: 'test-key', speechModel: 'u3-rt-pro-beta-1' });
+
+    expect(stt.model).toBe('u3-rt-pro-beta-1');
+  });
+
+  it('accepts u3-pro parameters for u3-rt-pro-beta-1', () => {
+    expect(
+      () =>
+        new STT({
+          apiKey: 'test-key',
+          speechModel: 'u3-rt-pro-beta-1',
+          prompt: 'medical dictation',
+          agentContext: "The agent asked for the patient's name.",
+          previousContextNTurns: 10,
+        }),
+    ).not.toThrow();
+  });
+
+  it('requires a u3-rt-pro model for agentContext', () => {
+    expect(
+      () =>
+        new STT({
+          apiKey: 'test-key',
+          speechModel: 'universal-streaming-english',
+          agentContext: 'hello',
+        }),
+    ).toThrow(/agentContext/);
+  });
+
+  it('requires a u3-rt-pro model for previousContextNTurns', () => {
+    expect(
+      () =>
+        new STT({
+          apiKey: 'test-key',
+          speechModel: 'universal-streaming-english',
+          previousContextNTurns: 5,
+        }),
+    ).toThrow(/previousContextNTurns/);
+  });
+
+  it('forwards inactivity timeout to the streaming query', async () => {
+    const { wss, baseUrl } = await startWebSocketServer();
+    let requestUrl = '';
+
+    wss.on('connection', (_ws, req) => {
+      requestUrl = req.url ?? '';
+    });
+
+    try {
+      const stream = new STT({
+        apiKey: 'test-key',
+        baseUrl,
+        inactivityTimeout: 45,
+      }).stream();
+
+      await waitUntil(() => requestUrl !== '');
+      stream.close();
+
+      const url = new URL(`ws://127.0.0.1${requestUrl}`);
+      expect(url.pathname).toBe('/v3/ws');
+      expect(url.searchParams.get('inactivity_timeout')).toBe('45');
+    } finally {
+      await closeWebSocketServer(wss);
+    }
+  });
+});
+
 const hasAssemblyAIApiKey = Boolean(process.env.ASSEMBLYAI_API_KEY);
 
 if (hasAssemblyAIApiKey) {

package/src/stt.ts

@@ -20,6 +20,7 @@ import type { RawData } from 'ws';
 import { WebSocket } from 'ws';
 import type { STTEncoding, STTModels, VoiceFocus } from './models.js';
 
+// Speech models in the Universal-3 Pro family, which share the same parameter support.
 const U3_PRO_MODELS = ['u3-rt-pro', 'u3-rt-pro-beta-1', 'universal-3-5-pro'] as const;
 
 function isU3ProModel(model: STTModels): boolean {
@@ -40,6 +41,7 @@ interface StreamEventMessage {
   end_of_turn_confidence?: number;
   turn_is_formatted?: boolean;
   language_code?: string;
+  language_confidence?: number;
   speaker_label?: string;
   words?: Array<{
     text?: string;
@@ -53,6 +55,18 @@ interface StreamEventMessage {
   session_duration_seconds?: number;
 }
 
+function speechDataMetadata(data: StreamEventMessage): stt.SpeechData['metadata'] | undefined {
+  const assemblyai: Record<string, number> = {};
+
+  if (typeof data.language_confidence === 'number') {
+    assemblyai.languageConfidence = data.language_confidence;
+  }
+
+  if (Object.keys(assemblyai).length === 0) return undefined;
+
+  return { assemblyai };
+}
+
 export interface STTOptions {
   apiKey?: string;
   sampleRate: number;
@@ -65,6 +79,11 @@ export interface STTOptions {
   encoding: STTEncoding;
   speechModel: STTModels;
   languageDetection?: boolean;
+  /**
+   * Session inactivity timeout in seconds. AssemblyAI accepts integer values
+   * from 5 to 3600; when unset, no inactivity timeout is applied.
+   */
+  inactivityTimeout?: number;
   endOfTurnConfidenceThreshold?: number;
   /** Minimum silence (ms) before a confident end-of-turn is finalized. */
   minTurnSilence?: number;
@@ -72,8 +91,12 @@ export interface STTOptions {
   maxTurnSilence?: number;
   formatTurns?: boolean;
   keytermsPrompt?: string[];
-  /** Only supported with the `u3-rt-pro` model. */
+  /** Only supported with the Universal-3 Pro model family. */
   prompt?: string;
+  /** Only supported with the Universal-3 Pro model family. */
+  agentContext?: string;
+  /** Only supported with the Universal-3 Pro model family. Set at connection time only. */
+  previousContextNTurns?: number;
   vadThreshold?: number;
   /**
    * Enable speaker diarization. Note: AssemblyAI will return per-word speaker
@@ -91,8 +114,8 @@ export interface STTOptions {
   /** Background audio suppression aggressiveness, from 0.0 to 1.0. Connect-time only. */
   voiceFocusThreshold?: number;
   /**
-   * Accuracy/latency preset for u3-rt-pro: `min_latency`, `balanced`, or `max_accuracy`.
-   * Explicit silence, partials, or VAD options still take precedence over mode defaults.
+   * Accuracy/latency preset for the Universal-3 Pro model family: `min_latency`, `balanced`,
+   * or `max_accuracy`. Explicit turn-silence values still take precedence over mode defaults.
    */
   mode?: 'min_latency' | 'balanced' | 'max_accuracy';
   baseUrl: string;
@@ -128,13 +151,20 @@ export class STT extends stt.STT {
     });
 
     if (opts.speechModel === 'u3-pro') {
-      log().warn("'u3-pro' is deprecated, use 'u3-rt-pro' instead.");
-      opts.speechModel = 'u3-rt-pro';
+      log().warn("'u3-pro' is deprecated, use 'universal-3-5-pro' instead.");
+      opts.speechModel = 'universal-3-5-pro';
     }
 
     const speechModel = opts.speechModel ?? defaultSTTOptions.speechModel;
     if (!isU3ProModel(speechModel)) {
-      for (const param of ['prompt', 'voiceFocus', 'voiceFocusThreshold', 'mode'] as const) {
+      for (const param of [
+        'prompt',
+        'agentContext',
+        'previousContextNTurns',
+        'voiceFocus',
+        'voiceFocusThreshold',
+        'mode',
+      ] as const) {
         if (opts[param] !== undefined) {
           throw new Error(
             `The '${param}' parameter is only supported with the ${U3_PRO_MODELS.join(', ')} models.`,
@@ -150,8 +180,8 @@ export class STT extends stt.STT {
       );
     }
 
-    // Minimize latency; matches LK's end-of-turn detector well.
-    const minTurnSilence = opts.minTurnSilence ?? 100;
+    // Minimize latency by default, but let AssemblyAI's mode preset control silence tuning.
+    const minTurnSilence = opts.minTurnSilence ?? (opts.mode === undefined ? 100 : undefined);
 
     this.#opts = {
       ...defaultSTTOptions,
@@ -226,6 +256,7 @@ export class SpeechStream extends stt.SpeechStream {
 
     const configMsg: Record<string, unknown> = { type: 'UpdateConfiguration' };
     if (opts.prompt !== undefined) configMsg.prompt = opts.prompt;
+    if (opts.agentContext !== undefined) configMsg.agent_context = opts.agentContext;
     if (opts.keytermsPrompt !== undefined) configMsg.keyterms_prompt = opts.keytermsPrompt;
     if (opts.maxTurnSilence !== undefined) configMsg.max_turn_silence = opts.maxTurnSilence;
     if (opts.minTurnSilence !== undefined) configMsg.min_turn_silence = opts.minTurnSilence;
@@ -284,11 +315,13 @@ export class SpeechStream extends stt.SpeechStream {
   }
 
   async #connectWS(): Promise<WebSocket> {
-    // u3-rt-pro family models default both min and max silence to 100ms when unset.
+    // Universal-3 Pro family models default both min and max silence to 100ms when unset.
+    // When a mode preset is selected, leave them unset unless explicitly provided so the
+    // server's per-mode silence tuning is not overridden by the latency-optimized default.
     let minSilence = this.#opts.minTurnSilence;
     let maxSilence = this.#opts.maxTurnSilence;
     if (isU3ProModel(this.#opts.speechModel)) {
-      if (minSilence === undefined) minSilence = 100;
+      if (minSilence === undefined && this.#opts.mode === undefined) minSilence = 100;
       if (maxSilence === undefined) maxSilence = minSilence;
     }
 
@@ -310,7 +343,10 @@ export class SpeechStream extends stt.SpeechStream {
           ? JSON.stringify(this.#opts.keytermsPrompt)
           : undefined,
       language_detection: languageDetection,
+      inactivity_timeout: this.#opts.inactivityTimeout,
       prompt: this.#opts.prompt,
+      agent_context: this.#opts.agentContext,
+      previous_context_n_turns: this.#opts.previousContextNTurns,
       vad_threshold: this.#opts.vadThreshold,
       speaker_labels: this.#opts.speakerLabels,
       max_speakers: this.#opts.maxSpeakers,
@@ -508,6 +544,7 @@ export class SpeechStream extends stt.SpeechStream {
     const utterance = data.utterance ?? '';
     const transcript = data.transcript ?? '';
     const language = normalizeLanguage(data.language_code ?? 'en');
+    const metadata = speechDataMetadata(data);
 
     // Word timestamps are in milliseconds:
     // https://www.assemblyai.com/docs/api-reference/streaming-api/streaming-api#receive.receiveTurn.words
@@ -542,6 +579,7 @@ export class SpeechStream extends stt.SpeechStream {
             endTime,
             confidence,
             words: timedWords,
+            ...(metadata ? { metadata } : {}),
           },
         ],
       });
@@ -569,6 +607,7 @@ export class SpeechStream extends stt.SpeechStream {
             endTime,
             confidence: utteranceConfidence,
             words: utteranceWords,
+            ...(metadata ? { metadata } : {}),
           },
         ],
       });
@@ -589,6 +628,7 @@ export class SpeechStream extends stt.SpeechStream {
             endTime,
             confidence,
             words: timedWords,
+            ...(metadata ? { metadata } : {}),
           },
         ],
       });