Support AssemblyAI inference context carryover

.changeset/fresh-pandas-listen.md

@@ -0,0 +1,5 @@
+---
+'@livekit/agents': patch
+---
+
+Support AssemblyAI inference STT agent context carryover.

.changeset/stt-context-options.md

@@ -0,0 +1,6 @@
+---
+'@livekit/agents': patch
+'@livekit/agents-plugin-assemblyai': patch
+---
+
+Add STT context options for keyterms and chat context forwarding.

agents/src/inference/stt.test.ts

@@ -1,12 +1,14 @@
 // SPDX-FileCopyrightText: 2025 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import { beforeAll, describe, expect, it } from 'vitest';
+import { afterEach, beforeAll, describe, expect, it, vi } from 'vitest';
 import * as agents from '../index.js';
 import { normalizeLanguage } from '../language.js';
+import { AgentHandoffItem, ChatMessage } from '../llm/index.js';
 import { initializeLogger } from '../log.js';
 import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../types.js';
 import { VAD, type VADStream } from '../vad.js';
+import { createConversationItemAddedEvent } from '../voice/events.js';
 import {
   STT,
   type STTFallbackModel,
@@ -21,6 +23,10 @@ beforeAll(() => {
   initializeLogger({ level: 'silent', pretty: false });
 });
 
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
 /** Helper to create STT with required credentials. */
 function makeStt(overrides: Record<string, unknown> = {}) {
   const defaults = {
@@ -32,6 +38,16 @@ function makeStt(overrides: Record<string, unknown> = {}) {
   return new STT({ ...defaults, ...overrides });
 }
 
+function makeAssemblyStt(overrides: Record<string, unknown> = {}) {
+  return makeStt({ model: 'assemblyai/universal-3-5-pro', ...overrides });
+}
+
+function assistantItemEvent(text: string) {
+  return createConversationItemAddedEvent(
+    ChatMessage.create({ role: 'assistant', content: [text] }),
+  );
+}
+
 describe('parseSTTModelString', () => {
   it('simple model without language', () => {
     const [model, language] = parseSTTModelString('deepgram');
@@ -329,6 +345,122 @@ describe('STT diarization capabilities', () => {
   });
 });
 
+describe('STT agent_context forwarding', () => {
+  it('enables chat context for AssemblyAI U3 Pro family models', () => {
+    for (const model of ['assemblyai/u3-rt-pro', 'assemblyai/universal-3-5-pro'] as const) {
+      const stt = makeAssemblyStt({ model });
+      expect(stt.capabilities.chatContext).toBe(true);
+    }
+  });
+
+  it('disables chat context for unsupported AssemblyAI models', () => {
+    const stt = makeAssemblyStt({ model: 'assemblyai/universal-streaming' });
+
+    expect(stt.capabilities.chatContext).toBe(false);
+  });
+
+  it('disables chat context for non-AssemblyAI models', () => {
+    const stt = makeAssemblyStt({ model: 'deepgram/nova-3' });
+    expect(stt.capabilities.chatContext).toBe(false);
+  });
+
+  it('derives chatContext capability from model support only', () => {
+    const stt = makeAssemblyStt({ modelOptions: { previous_context_n_turns: 0 } });
+    expect(stt.capabilities.chatContext).toBe(true);
+  });
+
+  it('emits active carryover shaping in the session.update wire payload', () => {
+    const stt = makeAssemblyStt();
+    const stream = stt.stream();
+    const sent: string[] = [];
+    Reflect.set(stream, 'activeWs', {
+      readyState: 1,
+      send: (payload: string) => sent.push(payload),
+    });
+
+    stt._pushConversationItem(assistantItemEvent('Your room is booked for Tuesday.'));
+
+    expect(stt['opts'].modelOptions).toHaveProperty(
+      'agent_context',
+      'Your room is booked for Tuesday.',
+    );
+    expect(sent).toHaveLength(1);
+    expect(JSON.parse(sent[0]!)).toMatchObject({
+      type: 'session.update',
+      settings: {
+        extra: {
+          agent_context: 'Your room is booked for Tuesday.',
+        },
+      },
+    });
+    stream.close();
+  });
+
+  it('truncates oversize replies keeping the tail', () => {
+    const text = 'a'.repeat(2000) + 'b'.repeat(1750);
+    const stt = makeAssemblyStt();
+    stt._pushConversationItem(assistantItemEvent(text));
+    expect(stt['opts'].modelOptions).toHaveProperty('agent_context', 'b'.repeat(1750));
+  });
+
+  it('ignores non-assistant items', () => {
+    const stt = makeAssemblyStt();
+
+    stt._pushConversationItem(
+      createConversationItemAddedEvent(ChatMessage.create({ role: 'user', content: ['hi there'] })),
+    );
+    expect(stt['opts'].modelOptions).not.toHaveProperty('agent_context');
+
+    stt._pushConversationItem(
+      createConversationItemAddedEvent(ChatMessage.create({ role: 'assistant', content: [] })),
+    );
+    expect(stt['opts'].modelOptions).not.toHaveProperty('agent_context');
+  });
+
+  it('ignores agent handoff items', () => {
+    const stt = makeAssemblyStt();
+    stt._pushConversationItem(
+      createConversationItemAddedEvent(AgentHandoffItem.create({ newAgentId: 'agent-2' })),
+    );
+    expect(stt['opts'].modelOptions).not.toHaveProperty('agent_context');
+  });
+
+  it('preserves explicit agent_context and later overwrites it with carryover', () => {
+    const stt = makeAssemblyStt({
+      modelOptions: { agent_context: 'The agent asked for a booking date.' },
+    });
+    expect(stt['opts'].modelOptions).toHaveProperty(
+      'agent_context',
+      'The agent asked for a booking date.',
+    );
+
+    stt._pushConversationItem(assistantItemEvent('And your zip code?'));
+    expect(stt['opts'].modelOptions).toHaveProperty('agent_context', 'And your zip code?');
+  });
+
+  it('starts forwarding after changing from an unsupported to a supported model', () => {
+    const stt = makeAssemblyStt({ model: 'assemblyai/universal-streaming' });
+
+    stt._pushConversationItem(assistantItemEvent('ignored before transition'));
+    expect(stt['opts'].modelOptions).not.toHaveProperty('agent_context');
+
+    stt.updateOptions({ model: 'assemblyai/universal-3-5-pro' });
+    stt._pushConversationItem(assistantItemEvent('forwarded after transition'));
+
+    expect(stt['opts'].modelOptions).toHaveProperty('agent_context', 'forwarded after transition');
+  });
+
+  it('stops forwarding after changing from a supported to an unsupported model', () => {
+    const stt = makeAssemblyStt();
+    stt._pushConversationItem(assistantItemEvent('last supported context'));
+
+    stt.updateOptions({ model: 'assemblyai/universal-streaming' });
+    stt._pushConversationItem(assistantItemEvent('ignored after transition'));
+
+    expect(stt['opts'].modelOptions).toHaveProperty('agent_context', 'last supported context');
+  });
+});
+
 describe('STT session keyterms', () => {
   it('updateOptions does not bake session keyterms into the user baseline', () => {
     const stt = makeStt({ model: 'deepgram/nova-3' });

agents/src/inference/stt.ts

@@ -7,6 +7,7 @@ import type { WebSocket } from 'ws';
 import { APIError, APIStatusError } from '../_exceptions.js';
 import { AudioByteStream } from '../audio.js';
 import { type LanguageCode, areLanguagesEquivalent, normalizeLanguage } from '../language.js';
+import { ChatMessage } from '../llm/index.js';
 import { log } from '../log.js';
 import { createStreamChannel } from '../stream/stream_channel.js';
 import {
@@ -19,6 +20,7 @@ import {
 import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../types.js';
 import { type AudioBuffer, Event, Task, cancelAndWait, shortuuid, waitForAbort } from '../utils.js';
 import { type VAD, VADEventType, type VADStream } from '../vad.js';
+import type { ConversationItemAddedEvent } from '../voice/events.js';
 import { type TimedString, createTimedString } from '../voice/io.js';
 import {
   type SttServerEvent,
@@ -120,8 +122,10 @@ export interface AssemblyAIOptions {
   keyterms_prompt?: string[];
   /** Enable speaker diarization. Default: false. */
   speaker_labels?: boolean;
-  /** Context to bias recognition. Only supported with u3-rt-pro. Max 1500 chars. */
+  /** Context to bias recognition. Only supported with u3-rt-pro. Max 1750 chars. */
   agent_context?: string;
+  /** Prior turns carried as context; 0 disables carryover. Only supported with u3-rt-pro. */
+  previous_context_n_turns?: number;
   /** Isolate the primary voice. Only supported with u3-rt-pro. */
   voice_focus?: 'near-field' | 'far-field';
   /** Background suppression strength. Only supported with u3-rt-pro. */
@@ -270,6 +274,17 @@ function keytermsExtraForModel(
   return { [key]: [...new Set([...existing, ...sessionKeyterms])] };
 }
 
+const ASSEMBLYAI_CARRYOVER_MODELS = [
+  'assemblyai/u3-rt-pro',
+  'assemblyai/universal-3-5-pro',
+] as const;
+
+const ASSEMBLYAI_MAX_AGENT_CONTEXT_CHARS = 1750;
+
+function supportsChatContext(model: string | undefined): boolean {
+  return model === ASSEMBLYAI_CARRYOVER_MODELS[0] || model === ASSEMBLYAI_CARRYOVER_MODELS[1];
+}
+
 type _STTModels =
   | DeepgramModels
   | DeepgramFluxModels
@@ -412,19 +427,32 @@ export class STT<TModel extends STTModels> extends BaseSTT {
     vad?: VAD;
   }) {
     const modelOptions = (opts?.modelOptions ?? {}) as STTOptions<TModel>;
+    // Parse language from model string if provided: "provider/model:language".
+    let nextModel = opts?.model;
+    let nextLanguage = opts?.language;
+    let hasLanguageConflict = false;
+    if (typeof nextModel === 'string') {
+      const [parsedModel, parsedLanguage] = parseSTTModelString(nextModel);
+      if (parsedLanguage !== undefined) {
+        hasLanguageConflict =
+          !!nextLanguage && !areLanguagesEquivalent(nextLanguage, parsedLanguage);
+        if (!hasLanguageConflict) {
+          nextLanguage = parsedLanguage as STTLanguages;
+        }
+        nextModel = parsedModel as TModel;
+      }
+    }
     super({
       streaming: true,
       interimResults: true,
       alignedTranscript: 'word',
       diarization: diarizationEnabled(modelOptions as Record<string, unknown>),
       keyterms:
-        keytermsExtraForModel(typeof opts?.model === 'string' ? opts.model : undefined) !==
-        undefined,
+        keytermsExtraForModel(typeof nextModel === 'string' ? nextModel : undefined) !== undefined,
+      chatContext: supportsChatContext(typeof nextModel === 'string' ? nextModel : undefined),
     });
 
     const {
-      model,
-      language,
       baseURL,
       encoding = DEFAULT_ENCODING,
       sampleRate = DEFAULT_SAMPLE_RATE,
@@ -447,22 +475,11 @@ export class STT<TModel extends STTModels> extends BaseSTT {
       throw new Error('apiSecret is required: pass apiSecret or set LIVEKIT_API_SECRET');
     }
 
-    // Parse language from model string if provided: "provider/model:language"
-    let nextModel = model;
-    let nextLanguage = language;
-    if (typeof nextModel === 'string') {
-      const [parsedModel, parsedLanguage] = parseSTTModelString(nextModel);
-      if (parsedLanguage !== undefined) {
-        if (nextLanguage && !areLanguagesEquivalent(nextLanguage, parsedLanguage)) {
-          this.#logger.warn(
-            '`language` is provided via both argument and model, using the one from the argument',
-            { language: nextLanguage, model: nextModel },
-          );
-        } else {
-          nextLanguage = parsedLanguage as STTLanguages;
-        }
-        nextModel = parsedModel as TModel;
-      }
+    if (hasLanguageConflict) {
+      this.#logger.warn(
+        '`language` is provided via both argument and model, using the one from the argument',
+        { language: nextLanguage, model: opts?.model },
+      );
     }
     const normalizedFallback = fallback ? normalizeSTTFallback(fallback) : undefined;
     this.vad = resolveVADForModel(nextModel, vad);
@@ -531,6 +548,7 @@ export class STT<TModel extends STTModels> extends BaseSTT {
       this._vadPromise = undefined;
       this.updateCapabilities({
         keyterms: keytermsExtraForModel(this.opts.model) !== undefined,
+        chatContext: supportsChatContext(this.opts.model),
       });
     }
 
@@ -587,6 +605,25 @@ export class STT<TModel extends STTModels> extends BaseSTT {
     }
   }
 
+  override _pushConversationItem(ev: ConversationItemAddedEvent): void {
+    if (!this.capabilities.chatContext) {
+      return;
+    }
+
+    const chatItem = ev.item;
+    if (chatItem instanceof ChatMessage && chatItem.role === 'assistant' && chatItem.textContent) {
+      let text = chatItem.textContent;
+      if (text.length > ASSEMBLYAI_MAX_AGENT_CONTEXT_CHARS) {
+        this.#logger.debug(
+          { fromChars: text.length, toChars: ASSEMBLYAI_MAX_AGENT_CONTEXT_CHARS },
+          'truncating agent_context carryover',
+        );
+        text = text.slice(-ASSEMBLYAI_MAX_AGENT_CONTEXT_CHARS);
+      }
+      this.updateOptions({ modelOptions: { agent_context: text } as STTOptions<TModel> });
+    }
+  }
+
   stream(options?: {
     language?: STTLanguages | string;
     connOptions?: APIConnectOptions;

agents/src/stt/fallback_adapter.test.ts

@@ -4,9 +4,14 @@
 import type { EventEmitter } from 'node:events';
 import { beforeAll, describe, expect, it, vi } from 'vitest';
 import { APIConnectionError, APIError } from '../_exceptions.js';
+import { ChatMessage } from '../llm/index.js';
 import { initializeLogger } from '../log.js';
 import type { APIConnectOptions } from '../types.js';
 import { type AudioBuffer, delay } from '../utils.js';
+import {
+  type ConversationItemAddedEvent,
+  createConversationItemAddedEvent,
+} from '../voice/events.js';
 import { FallbackAdapter } from './fallback_adapter.js';
 import { STT, type SpeechEvent, SpeechEventType, SpeechStream } from './stt.js';
 import { FakeSTT, RecognizeSentinel, emptyAudioFrame } from './testing/fake_stt.js';
@@ -123,6 +128,55 @@ describe('FallbackAdapter', () => {
     expect(adapter.capabilities.streaming).toBe(true);
   });
 
+  it('forwards conversation context only to children that support it', () => {
+    class ContextRecordingSTT extends FakeSTT {
+      readonly conversationItems: ConversationItemAddedEvent[] = [];
+
+      constructor(label: string, supportsChatContext: boolean) {
+        super({ label });
+        this.updateCapabilities({ chatContext: supportsChatContext });
+      }
+
+      override _pushConversationItem(ev: ConversationItemAddedEvent): void {
+        this.conversationItems.push(ev);
+      }
+    }
+
+    const supported = new ContextRecordingSTT('supported', true);
+    const unsupported = new ContextRecordingSTT('unsupported', false);
+    const adapter = new FallbackAdapter({ sttInstances: [supported, unsupported] });
+    const event = createConversationItemAddedEvent(
+      ChatMessage.create({ role: 'assistant', content: ['hello'] }),
+    );
+
+    adapter._pushConversationItem(event);
+
+    expect(supported.conversationItems).toEqual([event]);
+    expect(unsupported.conversationItems).toEqual([]);
+  });
+
+  it('tracks dynamic child conversation-context capabilities and removes listeners on close', async () => {
+    class DynamicContextSTT extends FakeSTT {
+      setChatContext(supported: boolean): void {
+        this.updateCapabilities({ chatContext: supported });
+      }
+    }
+
+    const child = new DynamicContextSTT();
+    const adapter = new FallbackAdapter({ sttInstances: [child] });
+    expect(adapter.capabilities.chatContext).toBe(false);
+    expect(child.listenerCount('capabilities_changed')).toBe(1);
+
+    child.setChatContext(true);
+    expect(adapter.capabilities.chatContext).toBe(true);
+
+    child.setChatContext(false);
+    expect(adapter.capabilities.chatContext).toBe(false);
+
+    await adapter.close();
+    expect(child.listenerCount('capabilities_changed')).toBe(0);
+  });
+
   it('_recognize falls through to the next instance on error', async () => {
     const primary = new FakeSTT({
       label: 'primary',