@livekit/agents 0.3.4 → 0.4.0

package/src/tokenize/token_stream.ts

@@ -0,0 +1,163 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { randomUUID } from 'node:crypto';
+import { AsyncIterableQueue } from '../utils.js';
+import type { TokenData } from './tokenizer.js';
+import { SentenceStream, WordStream } from './tokenizer.js';
+
+type TokenizeFunc = (x: string) => string[] | [string, number, number][];
+
+export class BufferedTokenStream implements AsyncIterableIterator<TokenData> {
+  protected queue = new AsyncIterableQueue<TokenData>();
+  protected closed = false;
+
+  #func: TokenizeFunc;
+  #minTokenLength: number;
+  #minContextLength: number;
+  #bufTokens: string[] = [];
+  #inBuf = '';
+  #outBuf = '';
+  #currentSegmentId: string;
+
+  constructor(func: TokenizeFunc, minTokenLength: number, minContextLength: number) {
+    this.#func = func;
+    this.#minTokenLength = minTokenLength;
+    this.#minContextLength = minContextLength;
+
+    this.#currentSegmentId = randomUUID();
+  }
+
+  /** Push a string of text into the token stream */
+  pushText(text: string) {
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+
+    this.#inBuf += text;
+    if (this.#inBuf.length < this.#minContextLength) return;
+
+    while (true) {
+      const tokens = this.#func(this.#inBuf);
+      if (tokens.length === 0) break;
+
+      if (this.#outBuf) this.#outBuf += ' ';
+
+      const tok = tokens.shift()!;
+      let tokText = tok as string;
+      if (tok.length > 1 && typeof tok[1] === 'number') {
+        tokText = tok[0];
+      }
+
+      this.#outBuf += tokText;
+      if (this.#outBuf.length >= this.#minTokenLength) {
+        this.queue.put({ token: this.#outBuf, segmentId: this.#currentSegmentId });
+        this.#outBuf = '';
+      }
+
+      if (typeof tok! !== 'string') {
+        this.#inBuf = this.#inBuf.slice(tok![2]);
+      } else {
+        this.#inBuf = this.#inBuf
+          .slice(Math.max(0, this.#inBuf.indexOf(tok)) + tok.length)
+          .trimStart();
+      }
+    }
+  }
+
+  /** Flush the stream, causing it to process all pending text */
+  flush() {
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+
+    if (this.#inBuf || this.#outBuf) {
+      const tokens = this.#func(this.#inBuf);
+      if (tokens) {
+        if (this.#outBuf) this.#outBuf += ' ';
+
+        if (typeof tokens[0] !== 'string') {
+          this.#outBuf += tokens.map((tok) => tok[0]).join(' ');
+        } else {
+          this.#outBuf += tokens.join(' ');
+        }
+      }
+
+      if (this.#outBuf) {
+        this.queue.put({ token: this.#outBuf, segmentId: this.#currentSegmentId });
+      }
+
+      this.#currentSegmentId = randomUUID();
+    }
+
+    this.#inBuf = '';
+    this.#outBuf = '';
+  }
+
+  /** Mark the input as ended and forbid additional pushes */
+  endInput() {
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+    this.flush();
+    this.closed = true;
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.queue.next();
+  }
+
+  /** Close both the input and output of the token stream */
+  close() {
+    this.queue.close();
+    this.closed = true;
+  }
+
+  [Symbol.asyncIterator](): BufferedTokenStream {
+    return this;
+  }
+}
+
+export class BufferedSentenceStream extends SentenceStream {
+  #stream: BufferedTokenStream;
+
+  constructor(func: TokenizeFunc, minTokenLength: number, minContextLength: number) {
+    super();
+    this.#stream = new BufferedTokenStream(func, minTokenLength, minContextLength);
+  }
+
+  pushText(text: string) {
+    this.#stream.pushText(text);
+  }
+
+  close() {
+    super.close();
+    this.#stream.close();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.#stream.next();
+  }
+}
+
+export class BufferedWordStream extends WordStream {
+  #stream: BufferedTokenStream;
+
+  constructor(func: TokenizeFunc, minTokenLength: number, minContextLength: number) {
+    super();
+    this.#stream = new BufferedTokenStream(func, minTokenLength, minContextLength);
+  }
+
+  pushText(text: string) {
+    this.#stream.pushText(text);
+  }
+
+  close() {
+    super.close();
+    this.#stream.close();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.#stream.next();
+  }
+}

package/src/tokenize/tokenizer.ts

@@ -0,0 +1,152 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { AsyncIterableQueue } from '../utils.js';
+
+// prettier-ignore
+export const PUNCTUATIONS = [
+  '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/', ':', ';', '<', '=',
+  '>', '?', '@', '[', '\\', ']', '^', '_', '`', '{', '|', '}', '~', '±', '—', '‘', '’', '“', '”',
+  '…',
+]
+
+export interface TokenData {
+  segmentId: string;
+  token: string;
+}
+
+export abstract class SentenceTokenizer {
+  abstract tokenize(text: string, language?: string): string[];
+
+  /**
+   * Returns a {@link SentenceStream} that can be used to push strings and receive smaller segments.
+   */
+  abstract stream(): SentenceStream;
+}
+
+export abstract class SentenceStream {
+  protected static readonly FLUSH_SENTINEL = Symbol('FLUSH_SENTINEL');
+  protected input = new AsyncIterableQueue<string | typeof SentenceStream.FLUSH_SENTINEL>();
+  protected queue = new AsyncIterableQueue<TokenData>();
+  #closed = false;
+
+  get closed(): boolean {
+    return this.#closed;
+  }
+
+  /** Push a string of text to the tokenizer */
+  pushText(text: string) {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.#closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.put(text);
+  }
+
+  /** Flush the tokenizer, causing it to process all pending text */
+  flush() {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.#closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.put(SentenceStream.FLUSH_SENTINEL);
+  }
+
+  /** Mark the input as ended and forbid additional pushes */
+  endInput() {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.#closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.close();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.queue.next();
+  }
+
+  /** Close both the input and output of the tokenizer stream */
+  close() {
+    this.input.close();
+    this.queue.close();
+    this.#closed = true;
+  }
+
+  [Symbol.asyncIterator](): SentenceStream {
+    return this;
+  }
+}
+
+export abstract class WordTokenizer {
+  abstract tokenize(text: string, language?: string): string[];
+
+  /**
+   * Returns a {@link WordStream} that can be used to push words and receive smaller segments.
+   */
+  abstract stream(): WordStream;
+}
+
+export abstract class WordStream {
+  protected static readonly FLUSH_SENTINEL = Symbol('FLUSH_SENTINEL');
+  protected input = new AsyncIterableQueue<string | typeof WordStream.FLUSH_SENTINEL>();
+  protected queue = new AsyncIterableQueue<TokenData>();
+  #closed = false;
+
+  get closed(): boolean {
+    return this.#closed;
+  }
+
+  /** Push a string of text to the tokenizer */
+  pushText(text: string) {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.#closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.put(text);
+  }
+
+  /** Flush the tokenizer, causing it to process all pending text */
+  flush() {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.#closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.put(WordStream.FLUSH_SENTINEL);
+  }
+
+  /** Mark the input as ended and forbid additional pushes */
+  endInput() {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.#closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.close();
+  }
+
+  next(): Promise<IteratorResult<TokenData>> {
+    return this.queue.next();
+  }
+
+  /** Close both the input and output of the tokenizer stream */
+  close() {
+    this.input.close();
+    this.queue.close();
+    this.#closed = true;
+  }
+
+  [Symbol.asyncIterator](): WordStream {
+    return this;
+  }
+}

package/src/tts/index.ts

@@ -1,23 +1,4 @@
 // SPDX-FileCopyrightText: 2024 LiveKit, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
-import { StreamAdapter, StreamAdapterWrapper } from './stream_adapter.js';
-import {
-  ChunkedStream,
-  SynthesisEvent,
-  SynthesisEventType,
-  SynthesizeStream,
-  type SynthesizedAudio,
-  TTS,
-} from './tts.js';
-
-export {
-  TTS,
-  SynthesisEvent,
-  SynthesisEventType,
-  SynthesizedAudio,
-  SynthesizeStream,
-  StreamAdapter,
-  StreamAdapterWrapper,
-  ChunkedStream,
-};
+export { type SynthesizedAudio, type TTSCapabilities, TTS, SynthesizeStream } from './tts.js';

package/src/tts/tts.ts

@@ -2,87 +2,140 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 import type { AudioFrame } from '@livekit/rtc-node';
-import { mergeFrames } from '../utils.js';
+import { AsyncIterableQueue } from '../utils.js';
 
+/** SynthesizedAudio is a packet of speech synthesis as returned by the TTS. */
 export interface SynthesizedAudio {
-  text: string;
-  data: AudioFrame;
+  /** Request ID (one segment could be made up of multiple requests) */
+  requestId: string;
+  /** Segment ID, each segment is separated by a flush */
+  segmentId: string;
+  /** Synthesized audio frame */
+  frame: AudioFrame;
+  /** Current segment of the synthesized audio */
+  deltaText?: string;
 }
 
-export enum SynthesisEventType {
-  /**
-   * Indicate the start of synthesis.
-   * Retriggered after FINISHED.
-   */
-  STARTED = 0,
-  /**
-   * Indicate that audio data is available.
-   */
-  AUDIO = 1,
-  /**
-   * Indicate the end of synthesis. Does not necessarily mean stream is done.
-   */
-  FINISHED = 2,
+/**
+ * Describes the capabilities of the TTS provider.
+ *
+ * @remarks
+ * At present, only `streaming` is supplied to this interface, and the framework only supports
+ * providers that do have a streaming endpoint.
+ */
+export interface TTSCapabilities {
+  streaming: boolean;
 }
 
-export class SynthesisEvent {
-  type: SynthesisEventType;
-  audio?: SynthesizedAudio;
-
-  constructor(type: SynthesisEventType, audio: SynthesizedAudio | undefined = undefined) {
-    this.type = type;
-    this.audio = audio;
+/**
+ * An instance of a text-to-speech adapter.
+ *
+ * @remarks
+ * This class is abstract, and as such cannot be used directly. Instead, use a provider plugin that
+ * exports its own child TTS class, which inherits this class's methods.
+ */
+export abstract class TTS {
+  #capabilities: TTSCapabilities;
+  #sampleRate: number;
+  #numChannels: number;
+
+  constructor(sampleRate: number, numChannels: number, capabilities: TTSCapabilities) {
+    this.#capabilities = capabilities;
+    this.#sampleRate = sampleRate;
+    this.#numChannels = numChannels;
   }
-}
-
-export abstract class SynthesizeStream implements IterableIterator<SynthesisEvent> {
-  abstract pushText(token?: string): void;
 
-  markSegmentEnd() {
-    this.pushText(undefined);
+  /** Returns this TTS's capabilities */
+  get capabilities(): TTSCapabilities {
+    return this.#capabilities;
   }
 
-  abstract close(wait: boolean): Promise<void>;
-  abstract next(): IteratorResult<SynthesisEvent>;
-
-  [Symbol.iterator](): SynthesizeStream {
-    return this;
+  /** Returns the sample rate of audio frames returned by this TTS */
+  get sampleRate(): number {
+    return this.#sampleRate;
   }
-}
-
-export abstract class TTS {
-  #streamingSupported: boolean;
 
-  constructor(streamingSupported: boolean) {
-    this.#streamingSupported = streamingSupported;
+  /** Returns the channel count of audio frames returned by this TTS */
+  get numChannels(): number {
+    return this.#numChannels;
   }
 
-  abstract synthesize(text: string): Promise<ChunkedStream>;
-
+  /**
+   * Returns a {@link SynthesizeStream} that can be used to push text and receive audio data
+   */
   abstract stream(): SynthesizeStream;
+}
 
-  get streamingSupported(): boolean {
-    return this.#streamingSupported;
+/**
+ * An instance of a text-to-speech stream, as an asynchronous iterable iterator.
+ *
+ * @example Looping through frames
+ * ```ts
+ * for await (const event of stream) {
+ *   await source.captureFrame(event.frame);
+ * }
+ * ```
+ *
+ * @remarks
+ * This class is abstract, and as such cannot be used directly. Instead, use a provider plugin that
+ * exports its own child SynthesizeStream class, which inherits this class's methods.
+ */
+export abstract class SynthesizeStream
+  implements AsyncIterableIterator<SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM>
+{
+  protected static readonly FLUSH_SENTINEL = Symbol('FLUSH_SENTINEL');
+  static readonly END_OF_STREAM = Symbol('END_OF_STREAM');
+  protected input = new AsyncIterableQueue<string | typeof SynthesizeStream.FLUSH_SENTINEL>();
+  protected queue = new AsyncIterableQueue<
+    SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM
+  >();
+  protected closed = false;
+
+  /** Push a string of text to the TTS */
+  pushText(text: string) {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.put(text);
   }
-}
 
-export abstract class ChunkedStream implements AsyncIterableIterator<SynthesizedAudio> {
-  async collect(): Promise<AudioFrame> {
-    const frames = [];
-    for await (const ev of this) {
-      frames.push(ev.data);
+  /** Flush the TTS, causing it to process all pending text */
+  flush() {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
     }
-    return mergeFrames(frames);
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.put(SynthesizeStream.FLUSH_SENTINEL);
   }
 
-  abstract close(): Promise<void>;
-  abstract next(): Promise<IteratorResult<SynthesizedAudio>>;
+  /** Mark the input as ended and forbid additional pushes */
+  endInput() {
+    if (this.input.closed) {
+      throw new Error('Input is closed');
+    }
+    if (this.closed) {
+      throw new Error('Stream is closed');
+    }
+    this.input.close();
+  }
 
-  [Symbol.iterator](): ChunkedStream {
-    return this;
+  next(): Promise<IteratorResult<SynthesizedAudio | typeof SynthesizeStream.END_OF_STREAM>> {
+    return this.queue.next();
+  }
+
+  /** Close both the input and output of the TTS stream */
+  close() {
+    this.input.close();
+    this.queue.close();
+    this.closed = true;
   }
 
-  [Symbol.asyncIterator](): ChunkedStream {
+  [Symbol.asyncIterator](): SynthesizeStream {
     return this;
   }
 }

package/src/utils.ts

@@ -225,39 +225,109 @@ export async function gracefullyCancel<T>(promise: CancellablePromise<T>): Promi
 }
 
 /** @internal */
-export class AsyncIterableQueue<T> implements AsyncIterable<T> {
-  private queue: Queue<T | typeof AsyncIterableQueue.QUEUE_END_MARKER>;
-  private closed = false;
-  private static readonly QUEUE_END_MARKER = Symbol('QUEUE_END_MARKER');
+export class AsyncIterableQueue<T> implements AsyncIterableIterator<T> {
+  private static readonly CLOSE_SENTINEL = Symbol('CLOSE_SENTINEL');
+  #queue = new Queue<T | typeof AsyncIterableQueue.CLOSE_SENTINEL>();
+  #closed = false;
 
-  constructor() {
-    this.queue = new Queue<T | typeof AsyncIterableQueue.QUEUE_END_MARKER>();
+  get closed(): boolean {
+    return this.#closed;
   }
 
   put(item: T): void {
-    if (this.closed) {
+    if (this.#closed) {
       throw new Error('Queue is closed');
     }
-    this.queue.put(item);
+    this.#queue.put(item);
   }
 
   close(): void {
-    this.closed = true;
-    this.queue.put(AsyncIterableQueue.QUEUE_END_MARKER);
-  }
-
-  [Symbol.asyncIterator](): AsyncIterator<T> {
-    return {
-      next: async (): Promise<IteratorResult<T>> => {
-        if (this.closed && this.queue.items.length === 0) {
-          return { value: undefined, done: true };
-        }
-        const item = await this.queue.get();
-        if (item === AsyncIterableQueue.QUEUE_END_MARKER && this.closed) {
-          return { value: undefined, done: true };
-        }
-        return { value: item as T, done: false };
-      },
-    };
+    this.#closed = true;
+    this.#queue.put(AsyncIterableQueue.CLOSE_SENTINEL);
+  }
+
+  async next(): Promise<IteratorResult<T>> {
+    if (this.#closed && this.#queue.items.length === 0) {
+      return { value: undefined, done: true };
+    }
+    const item = await this.#queue.get();
+    if (item === AsyncIterableQueue.CLOSE_SENTINEL && this.#closed) {
+      return { value: undefined, done: true };
+    }
+    return { value: item as T, done: false };
+  }
+
+  [Symbol.asyncIterator](): AsyncIterableQueue<T> {
+    return this;
+  }
+}
+
+/** @internal */
+export class ExpFilter {
+  #alpha: number;
+  #max?: number;
+  #filtered?: number = undefined;
+
+  constructor(alpha: number, max?: number) {
+    this.#alpha = alpha;
+    this.#max = max;
+  }
+
+  reset(alpha?: number) {
+    if (alpha) {
+      this.#alpha = alpha;
+    }
+    this.#filtered = undefined;
+  }
+
+  apply(exp: number, sample: number): number {
+    if (this.#filtered) {
+      const a = this.#alpha ** exp;
+      this.#filtered = a * this.#filtered + (1 - a) * sample;
+    } else {
+      this.#filtered = sample;
+    }
+
+    if (this.#max && this.#filtered > this.#max) {
+      this.#filtered = this.#max;
+    }
+
+    return this.#filtered;
+  }
+
+  get filtered(): number | undefined {
+    return this.#filtered;
+  }
+
+  set alpha(alpha: number) {
+    this.#alpha = alpha;
+  }
+}
+
+/** @internal */
+export class AudioEnergyFilter {
+  #cooldownSeconds: number;
+  #cooldown: number;
+
+  constructor(cooldownSeconds = 1) {
+    this.#cooldownSeconds = cooldownSeconds;
+    this.#cooldown = cooldownSeconds;
+  }
+
+  pushFrame(frame: AudioFrame): boolean {
+    const arr = Float32Array.from(frame.data, (x) => x / 32768);
+    const rms = (arr.map((x) => x ** 2).reduce((acc, x) => acc + x) / arr.length) ** 0.5;
+    if (rms > 0.004) {
+      this.#cooldown = this.#cooldownSeconds;
+      return true;
+    }
+
+    const durationSeconds = frame.samplesPerChannel / frame.sampleRate;
+    this.#cooldown -= durationSeconds;
+    if (this.#cooldown > 0) {
+      return true;
+    }
+
+    return false;
   }
 }