@livekit/agents 0.7.9 → 1.0.0-next.1

package/src/stt/stt.ts

@@ -1,12 +1,18 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import type { AudioFrame } from '@livekit/rtc-node';
+import { type AudioFrame, AudioResampler } from '@livekit/rtc-node';
 import type { TypedEventEmitter as TypedEmitter } from '@livekit/typed-emitter';
 import { EventEmitter } from 'node:events';
+import type { ReadableStream } from 'node:stream/web';
+import { APIConnectionError, APIError } from '../_exceptions.js';
+import { calculateAudioDuration } from '../audio.js';
+import { log } from '../log.js';
 import type { STTMetrics } from '../metrics/base.js';
+import { DeferredReadableStream } from '../stream/deferred_stream.js';
+import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../types.js';
 import type { AudioBuffer } from '../utils.js';
-import { AsyncIterableQueue } from '../utils.js';
+import { AsyncIterableQueue, delay, startSoon, toError } from '../utils.js';
 
 /** Indicates start/middle/end of speech */
 export enum SpeechEventType {
@@ -32,7 +38,6 @@ export enum SpeechEventType {
   END_OF_SPEECH = 3,
   /** Usage event, emitted periodically to indicate usage metrics. */
   RECOGNITION_USAGE = 4,
-  METRICS_COLLECTED = 5,
 }
 
 /** SpeechData contains metadata about this {@link SpeechEvent}. */
@@ -67,8 +72,17 @@ export interface STTCapabilities {
   interimResults: boolean;
 }
 
+export interface STTError {
+  type: 'stt_error';
+  timestamp: number;
+  label: string;
+  error: Error;
+  recoverable: boolean;
+}
+
 export type STTCallbacks = {
-  [SpeechEventType.METRICS_COLLECTED]: (metrics: STTMetrics) => void;
+  ['metrics_collected']: (metrics: STTMetrics) => void;
+  ['error']: (error: STTError) => void;
 };
 
 /**
@@ -97,19 +111,17 @@ export abstract class STT extends (EventEmitter as new () => TypedEmitter<STTCal
     const startTime = process.hrtime.bigint();
     const event = await this._recognize(frame);
     const duration = Number((process.hrtime.bigint() - startTime) / BigInt(1000000));
-    this.emit(SpeechEventType.METRICS_COLLECTED, {
+    this.emit('metrics_collected', {
+      type: 'stt_metrics',
       requestId: event.requestId ?? '',
       timestamp: Date.now(),
       duration,
       label: this.label,
-      audioDuration: Array.isArray(frame)
-        ? frame.reduce((sum, a) => sum + a.samplesPerChannel / a.sampleRate, 0)
-        : frame.samplesPerChannel / frame.sampleRate,
+      audioDuration: calculateAudioDuration(frame),
       streamed: false,
     });
     return event;
   }
-
   protected abstract _recognize(frame: AudioBuffer): Promise<SpeechEvent>;
 
   /**
@@ -140,35 +152,126 @@ export abstract class SpeechStream implements AsyncIterableIterator<SpeechEvent>
   protected input = new AsyncIterableQueue<AudioFrame | typeof SpeechStream.FLUSH_SENTINEL>();
   protected output = new AsyncIterableQueue<SpeechEvent>();
   protected queue = new AsyncIterableQueue<SpeechEvent>();
+  protected neededSampleRate?: number;
+  protected resampler?: AudioResampler;
   abstract label: string;
   protected closed = false;
   #stt: STT;
+  private deferredInputStream: DeferredReadableStream<AudioFrame>;
+  private logger = log();
+  private _connOptions: APIConnectOptions;
 
-  constructor(stt: STT) {
+  constructor(
+    stt: STT,
+    sampleRate?: number,
+    connectionOptions: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
+  ) {
     this.#stt = stt;
+    this._connOptions = connectionOptions;
+    this.deferredInputStream = new DeferredReadableStream<AudioFrame>();
+    this.neededSampleRate = sampleRate;
     this.monitorMetrics();
+    this.pumpInput();
+
+    // this is a hack to immitate asyncio.create_task so that mainTask
+    // is run **after** the constructor has finished. Otherwise we get
+    // runtime error when trying to access class variables in the
+    // `run` method.
+    startSoon(() => this.mainTask().then(() => this.queue.close()));
   }
 
-  protected async monitorMetrics() {
-    const startTime = process.hrtime.bigint();
+  private async mainTask() {
+    for (let i = 0; i < this._connOptions.maxRetry + 1; i++) {
+      try {
+        return await this.run();
+      } catch (error) {
+        if (error instanceof APIError) {
+          const retryInterval = this._connOptions._intervalForRetry(i);
+
+          if (this._connOptions.maxRetry === 0 || !error.retryable) {
+            this.emitError({ error, recoverable: false });
+            throw error;
+          } else if (i === this._connOptions.maxRetry) {
+            this.emitError({ error, recoverable: false });
+            throw new APIConnectionError({
+              message: `failed to recognize speech after ${this._connOptions.maxRetry + 1} attempts`,
+              options: { retryable: false },
+            });
+          } else {
+            this.emitError({ error, recoverable: true });
+            this.logger.warn(
+              { tts: this.#stt.label, attempt: i + 1, error },
+              `failed to recognize speech, retrying in ${retryInterval}s`,
+            );
+          }
+
+          if (retryInterval > 0) {
+            await delay(retryInterval);
+          }
+        } else {
+          this.emitError({ error: toError(error), recoverable: false });
+          throw error;
+        }
+      }
+    }
+  }
+
+  private emitError({ error, recoverable }: { error: Error; recoverable: boolean }) {
+    this.#stt.emit('error', {
+      type: 'stt_error',
+      timestamp: Date.now(),
+      label: this.#stt.label,
+      error,
+      recoverable,
+    });
+  }
 
+  protected async pumpInput() {
+    // TODO(AJS-35): Implement STT with webstreams API
+    const inputStream = this.deferredInputStream.stream;
+    const reader = inputStream.getReader();
+
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        this.pushFrame(value);
+      }
+    } catch (error) {
+      this.logger.error('Error in STTStream mainTask:', error);
+    } finally {
+      reader.releaseLock();
+    }
+  }
+
+  protected async monitorMetrics() {
     for await (const event of this.queue) {
       this.output.put(event);
       if (event.type !== SpeechEventType.RECOGNITION_USAGE) continue;
-      const duration = process.hrtime.bigint() - startTime;
       const metrics: STTMetrics = {
+        type: 'stt_metrics',
         timestamp: Date.now(),
         requestId: event.requestId!,
-        duration: Math.trunc(Number(duration / BigInt(1000000))),
-        label: this.label,
+        duration: 0,
+        label: this.#stt.label,
         audioDuration: event.recognitionUsage!.audioDuration,
         streamed: true,
       };
-      this.#stt.emit(SpeechEventType.METRICS_COLLECTED, metrics);
+      this.#stt.emit('metrics_collected', metrics);
     }
     this.output.close();
   }
 
+  protected abstract run(): Promise<void>;
+
+  updateInputStream(audioStream: ReadableStream<AudioFrame>) {
+    this.deferredInputStream.setSource(audioStream);
+  }
+
+  detachInputStream() {
+    this.deferredInputStream.detachSource();
+  }
+
   /** Push an audio frame to the STT */
   pushFrame(frame: AudioFrame) {
     if (this.input.closed) {
@@ -177,7 +280,21 @@ export abstract class SpeechStream implements AsyncIterableIterator<SpeechEvent>
     if (this.closed) {
       throw new Error('Stream is closed');
     }
-    this.input.put(frame);
+
+    if (this.neededSampleRate && frame.sampleRate !== this.neededSampleRate) {
+      if (!this.resampler) {
+        this.resampler = new AudioResampler(frame.sampleRate, this.neededSampleRate);
+      }
+    }
+
+    if (this.resampler) {
+      const frames = this.resampler.push(frame);
+      for (const frame of frames) {
+        this.input.put(frame);
+      }
+    } else {
+      this.input.put(frame);
+    }
   }
 
   /** Flush the STT, causing it to process all pending text */

package/src/tokenize/basic/basic.ts

@@ -12,17 +12,24 @@ interface TokenizerOptions {
   language: string;
   minSentenceLength: number;
   streamContextLength: number;
+  retainFormat: boolean;
 }
 
+const defaultTokenizerOptions: TokenizerOptions = {
+  language: 'en-US',
+  minSentenceLength: 20,
+  streamContextLength: 10,
+  retainFormat: false,
+};
+
 export class SentenceTokenizer extends tokenizer.SentenceTokenizer {
   #config: TokenizerOptions;
 
-  constructor(language = 'en-US', minSentenceLength = 20, streamContextLength = 10) {
+  constructor(options?: Partial<TokenizerOptions>) {
     super();
     this.#config = {
-      language,
-      minSentenceLength,
-      streamContextLength,
+      ...defaultTokenizerOptions,
+      ...options,
     };
   }
 
@@ -34,7 +41,8 @@ export class SentenceTokenizer extends tokenizer.SentenceTokenizer {
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   stream(language?: string): tokenizer.SentenceStream {
     return new BufferedSentenceStream(
-      (text: string) => splitSentences(text, this.#config.minSentenceLength),
+      (text: string) =>
+        splitSentences(text, this.#config.minSentenceLength, this.#config.retainFormat),
       this.#config.minSentenceLength,
       this.#config.streamContextLength,
     );

package/src/tokenize/basic/sentence.ts

@@ -5,7 +5,11 @@
 /**
  * Split the text into sentences.
  */
-export const splitSentences = (text: string, minLength = 20): [string, number, number][] => {
+export const splitSentences = (
+  text: string,
+  minLength = 20,
+  retainFormat: boolean = false,
+): [string, number, number][] => {
   const alphabets = /([A-Za-z])/g;
   const prefixes = /(Mr|St|Mrs|Ms|Dr)[.]/g;
   const suffixes = /(Inc|Ltd|Jr|Sr|Co)/g;
@@ -16,7 +20,12 @@ export const splitSentences = (text: string, minLength = 20): [string, number, n
   const digits = /([0-9])/g;
   const dots = /\.{2,}/g;
 
-  text = text.replaceAll('\n', ' ');
+  if (retainFormat) {
+    text = text.replaceAll('\n', '<nel><stop>');
+  } else {
+    text = text.replaceAll('\n', ' ');
+  }
+
   text = text.replaceAll(prefixes, '$1<prd>');
   text = text.replaceAll(websites, '<prd>$2');
   text = text.replaceAll(new RegExp(`${digits.source}[.]${digits.source}`, 'g'), '$1<prd>$2');
@@ -47,6 +56,10 @@ export const splitSentences = (text: string, minLength = 20): [string, number, n
   text = text.replaceAll('!', '!<stop>');
   text = text.replaceAll('<prd>', '.');
 
+  if (retainFormat) {
+    text = text.replaceAll('<nel>', '\n');
+  }
+
   const split = text.split('<stop>');
   text = text.replaceAll('<stop>', '');
 
@@ -54,21 +67,22 @@ export const splitSentences = (text: string, minLength = 20): [string, number, n
   let buf = '';
   let start = 0;
   let end = 0;
+  const prePad = retainFormat ? '' : ' ';
   for (const match of split) {
-    const sentence = match.trim();
+    const sentence = retainFormat ? match : match.trim();
     if (!sentence) continue;
 
-    buf += ' ' + sentence;
+    buf += prePad + sentence;
     end += match.length;
     if (buf.length > minLength) {
-      sentences.push([buf.slice(1), start, end]);
+      sentences.push([buf.slice(prePad.length), start, end]);
       start = end;
       buf = '';
     }
   }
 
   if (buf) {
-    sentences.push([buf.slice(1), start, text.length - 1]);
+    sentences.push([buf.slice(prePad.length), start, text.length - 1]);
   }
 
   return sentences;

package/src/tokenize/token_stream.ts

@@ -1,8 +1,7 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import { randomUUID } from 'node:crypto';
-import { AsyncIterableQueue } from '../utils.js';
+import { AsyncIterableQueue, shortuuid } from '../utils.js';
 import type { TokenData } from './tokenizer.js';
 import { SentenceStream, WordStream } from './tokenizer.js';
 
@@ -25,7 +24,7 @@ export class BufferedTokenStream implements AsyncIterableIterator<TokenData> {
     this.#minTokenLength = minTokenLength;
     this.#minContextLength = minContextLength;
 
-    this.#currentSegmentId = randomUUID();
+    this.#currentSegmentId = shortuuid();
   }
 
   /** Push a string of text into the token stream */
@@ -90,7 +89,7 @@ export class BufferedTokenStream implements AsyncIterableIterator<TokenData> {
         this.queue.put({ token: this.#outBuf, segmentId: this.#currentSegmentId });
       }
 
-      this.#currentSegmentId = randomUUID();
+      this.#currentSegmentId = shortuuid();
     }
 
     this.#inBuf = '';
@@ -142,6 +141,10 @@ export class BufferedSentenceStream extends SentenceStream {
     this.#stream.close();
   }
 
+  endInput() {
+    this.#stream.endInput();
+  }
+
   next(): Promise<IteratorResult<TokenData>> {
     return this.#stream.next();
   }

package/src/transcription.ts

@@ -4,11 +4,10 @@
 import { TranscriptionSegment } from '@livekit/protocol';
 import { AudioFrame } from '@livekit/rtc-node';
 import type { TypedEventEmitter as TypedEmitter } from '@livekit/typed-emitter';
-import { randomUUID } from 'node:crypto';
 import { EventEmitter } from 'node:events';
 import { basic } from './tokenize/index.js';
 import type { SentenceStream, SentenceTokenizer } from './tokenize/tokenizer.js';
-import { AsyncIterableQueue, Future } from './utils.js';
+import { AsyncIterableQueue, Future, shortuuid } from './utils.js';
 
 // standard speech rate in hyphens/ms
 const STANDARD_SPEECH_RATE = 3830;
@@ -215,7 +214,7 @@ export class TextAudioSynchronizer extends (EventEmitter as new () => TypedEmitt
       realSpeed = this.#calcHyphens(textData.pushedText).length / audioData.pushedDuration;
     }
 
-    const segId = 'SG_' + randomUUID();
+    const segId = shortuuid('SG_');
     const words = this.#opts.splitWords(sentence);
     const processedWords: string[] = [];
 

package/src/tts/index.ts

@@ -6,7 +6,6 @@ export {
   type TTSCapabilities,
   type TTSCallbacks,
   TTS,
-  TTSEvent,
   SynthesizeStream,
   ChunkedStream,
 } from './tts.js';

package/src/tts/stream_adapter.ts

@@ -2,8 +2,9 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 import type { SentenceStream, SentenceTokenizer } from '../tokenize/index.js';
+import { Task } from '../utils.js';
 import type { ChunkedStream } from './tts.js';
-import { SynthesizeStream, TTS, TTSEvent } from './tts.js';
+import { SynthesizeStream, TTS } from './tts.js';
 
 export class StreamAdapter extends TTS {
   #tts: TTS;
@@ -17,8 +18,8 @@ export class StreamAdapter extends TTS {
     this.label = this.#tts.label;
     this.label = `tts.StreamAdapter<${this.#tts.label}>`;
 
-    this.#tts.on(TTSEvent.METRICS_COLLECTED, (metrics) => {
-      this.emit(TTSEvent.METRICS_COLLECTED, metrics);
+    this.#tts.on('metrics_collected', (metrics) => {
+      this.emit('metrics_collected', metrics);
     });
   }
 
@@ -41,17 +42,13 @@ export class StreamAdapterWrapper extends SynthesizeStream {
     this.#tts = tts;
     this.#sentenceStream = sentenceTokenizer.stream();
     this.label = `tts.StreamAdapterWrapper<${this.#tts.label}>`;
-
-    this.#run();
-  }
-
-  async monitorMetrics() {
-    return; // do nothing
   }
 
-  async #run() {
+  protected async run() {
     const forwardInput = async () => {
       for await (const input of this.input) {
+        if (this.abortController.signal.aborted) break;
+
         if (input === SynthesizeStream.FLUSH_SENTINEL) {
           this.#sentenceStream.flush();
         } else {
@@ -62,15 +59,44 @@ export class StreamAdapterWrapper extends SynthesizeStream {
       this.#sentenceStream.close();
     };
 
-    const synthesize = async () => {
+    const synthesizeSentenceStream = async () => {
+      let task: Task<void> | undefined;
+      const tokenCompletionTasks: Task<void>[] = [];
+
       for await (const ev of this.#sentenceStream) {
-        for await (const audio of this.#tts.synthesize(ev.token)) {
-          this.output.put(audio);
-        }
+        if (this.abortController.signal.aborted) break;
+
+        // this will enable non-blocking synthesis of the stream of tokens
+        task = Task.from(
+          (controller) => synthesize(ev.token, task, controller),
+          this.abortController,
+        );
+
+        tokenCompletionTasks.push(task);
+      }
+
+      await Promise.all(tokenCompletionTasks.map((t) => t.result));
+      this.queue.put(SynthesizeStream.END_OF_STREAM);
+    };
+
+    const synthesize = async (
+      token: string,
+      prevTask: Task<void> | undefined,
+      controller: AbortController,
+    ) => {
+      const audioStream = this.#tts.synthesize(token);
+
+      // wait for previous audio transcription to complete before starting
+      // to queuing audio frames of the current token
+      await prevTask?.result;
+      if (controller.signal.aborted) return;
+
+      for await (const audio of audioStream) {
+        if (controller.signal.aborted) break;
+        this.queue.put(audio);
       }
-      this.output.put(SynthesizeStream.END_OF_STREAM);
     };
 
-    Promise.all([forwardInput(), synthesize()]);
+    await Promise.all([forwardInput(), synthesizeSentenceStream()]);
   }
 }