Support AssemblyAI inference context carryover

.changeset/fresh-pandas-listen.md

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

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 { initializeLogger } from '../log.js';
+import { AgentHandoffItem, ChatMessage } from '../llm/index.js';
+import { initializeLogger, log } 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,111 @@ describe('STT diarization capabilities', () => {
   });
 });
 
+describe('STT agent_context carryover', () => {
+  it('agentContextCarryover defaults to enabled on 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('agentContextCarryover defaults off for unsupported models without warning', () => {
+    const warnSpy = vi.spyOn(log(), 'warn');
+    const stt = makeAssemblyStt({ model: 'assemblyai/universal-streaming' });
+
+    expect(stt.capabilities.chatContext).toBe(false);
+    expect(warnSpy).not.toHaveBeenCalled();
+  });
+
+  it('agentContextCarryover is off for non-AssemblyAI models', () => {
+    const stt = makeAssemblyStt({ model: 'deepgram/nova-3' });
+    expect(stt.capabilities.chatContext).toBe(false);
+  });
+
+  it('explicit true on unsupported model warns and is ignored', () => {
+    const warnSpy = vi.spyOn(log(), 'warn');
+    const stt = makeAssemblyStt({
+      model: 'assemblyai/universal-streaming',
+      agentContextCarryover: true,
+    });
+
+    expect(stt.capabilities.chatContext).toBe(false);
+    expect(warnSpy).toHaveBeenCalledWith(
+      { model: 'assemblyai/universal-streaming' },
+      'agentContextCarryover is enabled but model does not support it; ignoring',
+    );
+  });
+
+  it('explicit false disables carryover on a supported model', () => {
+    const stt = makeAssemblyStt({ agentContextCarryover: false });
+    expect(stt.capabilities.chatContext).toBe(false);
+  });
+
+  it('previous_context_n_turns=0 disables the default carryover', () => {
+    const stt = makeAssemblyStt({ modelOptions: { previous_context_n_turns: 0 } });
+    expect(stt.capabilities.chatContext).toBe(false);
+  });
+
+  it('explicit true wins over previous_context_n_turns=0', () => {
+    const stt = makeAssemblyStt({
+      modelOptions: { previous_context_n_turns: 0 },
+      agentContextCarryover: true,
+    });
+    expect(stt.capabilities.chatContext).toBe(true);
+  });
+
+  it('forwards short assistant replies verbatim', () => {
+    const stt = makeAssemblyStt();
+    stt._pushConversationItem(assistantItemEvent('Your room is booked for Tuesday.'));
+    expect(stt['opts'].modelOptions).toHaveProperty(
+      'agent_context',
+      'Your room is booked for Tuesday.',
+    );
+  });
+
+  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?');
+  });
+});
+
 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,28 @@ 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 supportsAgentContextCarryover(model: string | undefined): boolean {
+  return ASSEMBLYAI_CARRYOVER_MODELS.includes(
+    model as (typeof ASSEMBLYAI_CARRYOVER_MODELS)[number],
+  );
+}
+
+function agentContextCarryoverEnabled(opts: {
+  model: string | undefined;
+  modelOptions: Record<string, unknown>;
+  agentContextCarryover?: boolean;
+}): boolean {
+  if (!supportsAgentContextCarryover(opts.model)) return false;
+  return opts.agentContextCarryover ?? opts.modelOptions.previous_context_n_turns !== 0;
+}
+
 type _STTModels =
   | DeepgramModels
   | DeepgramFluxModels
@@ -370,6 +396,7 @@ export interface InferenceSTTOptions<TModel extends STTModels> {
   apiKey: string;
   apiSecret: string;
   modelOptions: STTOptions<TModel>;
+  agentContextCarryover?: boolean;
   fallback?: STTFallbackModel[];
   connOptions?: APIConnectOptions;
 }
@@ -408,23 +435,49 @@ export class STT<TModel extends STTModels> extends BaseSTT {
     apiSecret?: string;
     modelOptions?: STTOptions<TModel>;
     fallback?: STTFallbackModelType | STTFallbackModelType[];
+    agentContextCarryover?: boolean;
     connOptions?: APIConnectOptions;
     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;
+      }
+    }
+    const carryoverEnabled = agentContextCarryoverEnabled({
+      model: typeof nextModel === 'string' ? nextModel : undefined,
+      modelOptions: modelOptions as Record<string, unknown>,
+      agentContextCarryover: opts?.agentContextCarryover,
+    });
     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: carryoverEnabled,
     });
 
+    if (opts?.agentContextCarryover && !supportsAgentContextCarryover(nextModel)) {
+      log().warn(
+        { model: nextModel },
+        'agentContextCarryover is enabled but model does not support it; ignoring',
+      );
+    }
+
     const {
-      model,
-      language,
       baseURL,
       encoding = DEFAULT_ENCODING,
       sampleRate = DEFAULT_SAMPLE_RATE,
@@ -447,22 +500,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);
@@ -476,6 +518,7 @@ export class STT<TModel extends STTModels> extends BaseSTT {
       apiKey: lkApiKey,
       apiSecret: lkApiSecret,
       modelOptions,
+      agentContextCarryover: opts?.agentContextCarryover,
       fallback: normalizedFallback,
       connOptions: connOptions ?? DEFAULT_API_CONNECT_OPTIONS,
     };
@@ -531,12 +574,22 @@ export class STT<TModel extends STTModels> extends BaseSTT {
       this._vadPromise = undefined;
       this.updateCapabilities({
         keyterms: keytermsExtraForModel(this.opts.model) !== undefined,
+        chatContext: agentContextCarryoverEnabled({
+          model: this.opts.model,
+          modelOptions: this.opts.modelOptions as Record<string, unknown>,
+          agentContextCarryover: this.opts.agentContextCarryover,
+        }),
       });
     }
 
     if (nextOpts.modelOptions) {
       this.updateCapabilities({
         diarization: diarizationEnabled(this.opts.modelOptions as Record<string, unknown>),
+        chatContext: agentContextCarryoverEnabled({
+          model: this.opts.model,
+          modelOptions: this.opts.modelOptions as Record<string, unknown>,
+          agentContextCarryover: this.opts.agentContextCarryover,
+        }),
       });
       // re-apply the active session keyterms on top of the update sent to live streams,
       // so a user extra update doesn't drop them. `this.opts.modelOptions` must stay a
@@ -587,6 +640,21 @@ export class STT<TModel extends STTModels> extends BaseSTT {
     }
   }
 
+  override _pushConversationItem(ev: ConversationItemAddedEvent): void {
+    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;