Add STT and utils

Author: nbsp

agents/build.ts

@@ -5,7 +5,7 @@
 import dts from 'bun-plugin-dts';
 
 await Bun.build({
-  entrypoints: ['./src/index.ts'],
+  entrypoints: ['./src/index.ts', './src/tts/index.ts', './src/stt/index.ts'],
   outdir: './dist',
   target: 'node',
   sourcemap: 'external',

agents/src/stt/index.ts

@@ -0,0 +1,6 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+export { STT, SpeechEvent, SpeechEventType, SpeechStream } from './stt';
+export { StreamAdapter, StreamAdapterWrapper } from './stream_adapter';

agents/src/stt/stream_adapter.ts

@@ -0,0 +1,105 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import { STT, SpeechEvent, SpeechEventType, SpeechStream } from './stt';
+import { VADEventType, VADStream } from '../vad';
+import { AudioFrame } from '@livekit/rtc-node';
+import { AudioBuffer, mergeFrames } from '../utils';
+
+export class StreamAdapterWrapper extends SpeechStream {
+  closed: boolean;
+  stt: STT;
+  vadStream: VADStream;
+  eventQueue: (SpeechEvent | undefined)[];
+  language: string | undefined;
+  task: {
+    run: Promise<void>;
+    cancel: () => void;
+  };
+
+  constructor(stt: STT, vadStream: VADStream, language: string | undefined = undefined) {
+    super();
+    this.closed = false;
+    this.stt = stt;
+    this.vadStream = vadStream;
+    this.eventQueue = [];
+    this.language = language;
+    this.task = {
+      run: new Promise((_, reject) => {
+        this.run(reject);
+      }),
+      cancel: () => {},
+    };
+  }
+
+  async run(reject: (arg: Error) => void) {
+    this.task.cancel = () => {
+      this.closed = true;
+      reject(new Error('cancelled'));
+    };
+
+    for (const event of this.vadStream) {
+      if (event.type == VADEventType.START_OF_SPEECH) {
+        const startEvent = new SpeechEvent(SpeechEventType.START_OF_SPEECH);
+        this.eventQueue.push(startEvent);
+      } else if (event.type == VADEventType.END_OF_SPEECH) {
+        const mergedFrames = mergeFrames(event.speech);
+        const endEvent = await this.stt.recognize(mergedFrames, this.language);
+        this.eventQueue.push(endEvent);
+      }
+    }
+
+    this.eventQueue.push(undefined);
+  }
+
+  pushFrame(frame: AudioFrame) {
+    if (this.closed) {
+      throw new TypeError('cannot push frame to closed stream');
+    }
+
+    this.vadStream.pushFrame(frame);
+  }
+
+  async close(wait: boolean = true): Promise<void> {
+    this.closed = true;
+
+    if (!wait) {
+      this.task.cancel();
+    }
+
+    await this.vadStream.close(wait);
+    await this.task.run;
+  }
+
+  next(): IteratorResult<SpeechEvent> {
+    const item = this.eventQueue.shift();
+    if (item) {
+      return { done: false, value: item };
+    } else {
+      return { done: true, value: undefined };
+    }
+  }
+}
+
+export class StreamAdapter extends STT {
+  stt: STT;
+  vadStream: VADStream;
+
+  constructor(stt: STT, vadStream: VADStream) {
+    super(true);
+    this.stt = stt;
+    this.vadStream = vadStream;
+  }
+
+  async recognize(
+    buffer: AudioBuffer,
+    language: string | undefined = undefined,
+  ): Promise<SpeechEvent> {
+    return await this.stt.recognize(buffer, language);
+  }
+
+  stream(language: string | undefined = undefined) {
+    return new StreamAdapterWrapper(this.stt, this.vadStream, language);
+  }
+}

agents/src/stt/stt.ts

@@ -0,0 +1,59 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import { AudioFrame } from '@livekit/rtc-node';
+import { AudioBuffer } from '../utils';
+
+export enum SpeechEventType {
+  START_OF_SPEECH = 0,
+  INTERIM_TRANSCRIPT = 1,
+  FINAL_TRANSCRIPT = 2,
+  END_OF_SPEECH = 3,
+}
+
+export interface SpeechData {
+  language: string;
+  text: string;
+  startTime: number;
+  endTime: number;
+  confidence: number;
+}
+
+export class SpeechEvent {
+  type: SpeechEventType;
+  alternatives: SpeechData[];
+
+  constructor(type: SpeechEventType, alternatives: SpeechData[] = []) {
+    this.type = type;
+    this.alternatives = alternatives;
+  }
+}
+
+export abstract class SpeechStream implements IterableIterator<SpeechEvent> {
+  abstract pushFrame(token: AudioFrame): void;
+
+  abstract close(wait: boolean): Promise<void>;
+
+  abstract next(): IteratorResult<SpeechEvent>;
+
+  [Symbol.iterator](): SpeechStream {
+    return this;
+  }
+}
+
+export abstract class STT {
+  #streamingSupported: boolean;
+
+  constructor(streamingSupported: boolean) {
+    this.#streamingSupported = streamingSupported;
+  }
+
+  abstract recognize(buffer: AudioBuffer, language: string | undefined): Promise<SpeechEvent>;
+
+  abstract stream(language: string | undefined): SpeechStream;
+
+  get streamingSupported(): boolean {
+    return this.#streamingSupported;
+  }
+}

agents/src/utils.ts

@@ -0,0 +1,38 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+import { AudioFrame } from '@livekit/rtc-node';
+
+export type AudioBuffer = AudioFrame[] | AudioFrame;
+
+export const mergeFrames = (buffer: AudioBuffer): AudioFrame => {
+  if (Array.isArray(buffer)) {
+    buffer = buffer as AudioFrame[];
+    if (buffer.length == 0) {
+      throw new TypeError('buffer is empty');
+    }
+
+    const sampleRate = buffer[0].sampleRate;
+    const channels = buffer[0].channels;
+    let samplesPerChannel = 0;
+    let data = new Uint16Array();
+
+    for (const frame of buffer) {
+      if (frame.sampleRate !== sampleRate) {
+        throw new TypeError('sample rate mismatch');
+      }
+
+      if (frame.channels !== channels) {
+        throw new TypeError('channel count mismatch');
+      }
+
+      data = new Uint16Array([...data, ...frame.data]);
+      samplesPerChannel += frame.samplesPerChannel;
+    }
+
+    return new AudioFrame(data, sampleRate, channels, samplesPerChannel);
+  }
+
+  return buffer;
+};