@livekit/agents 0.1.0

package/src/tts/index.ts

@@ -0,0 +1,23 @@
+// 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,
+};

package/src/tts/stream_adapter.ts

@@ -0,0 +1,92 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { SentenceStream, SentenceTokenizer } from '../tokenize.js';
+import { ChunkedStream, SynthesisEvent, SynthesisEventType, SynthesizeStream, TTS } from './tts.js';
+
+export class StreamAdapterWrapper extends SynthesizeStream {
+  closed: boolean;
+  tts: TTS;
+  sentenceStream: SentenceStream;
+  eventQueue: (SynthesisEvent | undefined)[];
+  task: {
+    run: Promise<void>;
+    cancel: () => void;
+  };
+
+  constructor(tts: TTS, sentenceStream: SentenceStream) {
+    super();
+    this.closed = false;
+    this.tts = tts;
+    this.sentenceStream = sentenceStream;
+    this.eventQueue = [];
+    this.task = {
+      run: new Promise((_, reject) => {
+        this.run(reject);
+      }),
+      cancel: () => {},
+    };
+  }
+
+  async run(reject: (arg: Error) => void) {
+    while (!this.closed) {
+      this.task.cancel = () => {
+        this.closed = true;
+        reject(new Error('cancelled'));
+      };
+      for await (const sentence of this.sentenceStream) {
+        const audio = await this.tts.synthesize(sentence.text).then((data) => data.next());
+        if (!audio.done) {
+          this.eventQueue.push(new SynthesisEvent(SynthesisEventType.STARTED));
+          this.eventQueue.push(new SynthesisEvent(SynthesisEventType.AUDIO, audio.value));
+          this.eventQueue.push(new SynthesisEvent(SynthesisEventType.FINISHED));
+        }
+      }
+    }
+  }
+
+  pushText(token: string) {
+    this.sentenceStream.pushText(token);
+  }
+
+  async flush() {
+    await this.sentenceStream.flush();
+  }
+
+  next(): IteratorResult<SynthesisEvent> {
+    const event = this.eventQueue.shift();
+    if (event) {
+      return { done: false, value: event };
+    } else {
+      return { done: true, value: undefined };
+    }
+  }
+
+  async close(): Promise<void> {
+    this.task.cancel();
+    try {
+      await this.task.run;
+    } finally {
+      this.eventQueue.push(undefined);
+    }
+  }
+}
+
+export class StreamAdapter extends TTS {
+  tts: TTS;
+  tokenizer: SentenceTokenizer;
+
+  constructor(tts: TTS, tokenizer: SentenceTokenizer) {
+    super(true);
+    this.tts = tts;
+    this.tokenizer = tokenizer;
+  }
+
+  synthesize(text: string): Promise<ChunkedStream> {
+    return this.tts.synthesize(text);
+  }
+
+  stream() {
+    return new StreamAdapterWrapper(this.tts, this.tokenizer.stream(undefined));
+  }
+}

package/src/tts/tts.ts

@@ -0,0 +1,78 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { AudioFrame } from '@livekit/rtc-node';
+import { mergeFrames } from '../utils.js';
+
+export interface SynthesizedAudio {
+  text: string;
+  data: AudioFrame;
+}
+
+export enum SynthesisEventType {
+  STARTED = 0,
+  AUDIO = 1,
+  FINISHED = 2,
+}
+
+export class SynthesisEvent {
+  type: SynthesisEventType;
+  audio?: SynthesizedAudio;
+
+  constructor(type: SynthesisEventType, audio: SynthesizedAudio | undefined = undefined) {
+    this.type = type;
+    this.audio = audio;
+  }
+}
+
+export abstract class SynthesizeStream implements IterableIterator<SynthesisEvent> {
+  abstract pushText(token?: string): void;
+
+  markSegmentEnd() {
+    this.pushText(undefined);
+  }
+
+  abstract close(wait: boolean): Promise<void>;
+  abstract next(): IteratorResult<SynthesisEvent>;
+
+  [Symbol.iterator](): SynthesizeStream {
+    return this;
+  }
+}
+
+export abstract class TTS {
+  #streamingSupported: boolean;
+
+  constructor(streamingSupported: boolean) {
+    this.#streamingSupported = streamingSupported;
+  }
+
+  abstract synthesize(text: string): Promise<ChunkedStream>;
+
+  abstract stream(): SynthesizeStream;
+
+  get streamingSupported(): boolean {
+    return this.#streamingSupported;
+  }
+}
+
+export abstract class ChunkedStream implements AsyncIterableIterator<SynthesizedAudio> {
+  async collect(): Promise<AudioFrame> {
+    const frames = [];
+    for await (const ev of this) {
+      frames.push(ev.data);
+    }
+    return mergeFrames(frames);
+  }
+
+  abstract close(): Promise<void>;
+  abstract next(): Promise<IteratorResult<SynthesizedAudio>>;
+
+  [Symbol.iterator](): ChunkedStream {
+    return this;
+  }
+
+  [Symbol.asyncIterator](): ChunkedStream {
+    return this;
+  }
+}

package/src/utils.ts

@@ -0,0 +1,37 @@
+// 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;
+};

package/src/vad.ts

@@ -0,0 +1,42 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { AudioFrame } from '@livekit/rtc-node';
+
+export enum VADEventType {
+  START_OF_SPEECH = 1,
+  SPEAKING = 2,
+  END_OF_SPEECH = 3,
+}
+
+export interface VADEvent {
+  type: VADEventType;
+  samplesIndex: number;
+  duration: number;
+  speech: AudioFrame[];
+}
+
+export abstract class VAD {
+  abstract stream({
+    minSpeakingDuration,
+    minSilenceDuration,
+    paddingDuration,
+    sampleRate,
+    maxBufferedSpeech,
+  }: {
+    minSpeakingDuration: number;
+    minSilenceDuration: number;
+    paddingDuration: number;
+    sampleRate: number;
+    maxBufferedSpeech: number;
+  }): VADStream;
+}
+
+export abstract class VADStream implements IterableIterator<VADEvent> {
+  abstract pushFrame(frame: AudioFrame): void;
+  abstract close(wait: boolean): Promise<void>;
+  abstract next(): IteratorResult<VADEvent>;
+  [Symbol.iterator](): VADStream {
+    return this;
+  }
+}

package/src/version.ts

@@ -0,0 +1,5 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+export const version = '0.1.0';

package/src/worker.ts

@@ -0,0 +1,384 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import {
+  type AvailabilityRequest,
+  type Job,
+  type JobAssignment,
+  JobType,
+  ParticipantPermission,
+  ServerMessage,
+  WorkerMessage,
+} from '@livekit/protocol';
+import { EventEmitter } from 'events';
+import { AccessToken } from 'livekit-server-sdk';
+import os from 'os';
+import { WebSocket } from 'ws';
+import { HTTPServer } from './http_server.js';
+import { JobProcess } from './ipc/job_process.js';
+import { type AvailRes, JobRequest } from './job_request.js';
+import type { AcceptData } from './job_request.js';
+import { log } from './log.js';
+import { version } from './version.js';
+
+const MAX_RECONNECT_ATTEMPTS = 10;
+const ASSIGNMENT_TIMEOUT = 15 * 1000;
+const LOAD_INTERVAL = 5 * 1000;
+
+const cpuLoad = (): number =>
+  (os
+    .cpus()
+    .reduce(
+      (acc, x) => acc + x.times.idle / Object.values(x.times).reduce((acc, x) => acc + x, 0),
+      0,
+    ) /
+    os.cpus().length) *
+  100;
+
+class WorkerPermissions {
+  canPublish: boolean;
+  canSubscribe: boolean;
+  canPublishData: boolean;
+  canUpdateMetadata: boolean;
+  hidden: boolean;
+
+  constructor(
+    canPublish = true,
+    canSubscribe = true,
+    canPublishData = true,
+    canUpdateMetadata = true,
+    hidden = false,
+  ) {
+    this.canPublish = canPublish;
+    this.canSubscribe = canSubscribe;
+    this.canPublishData = canPublishData;
+    this.canUpdateMetadata = canUpdateMetadata;
+    this.hidden = hidden;
+  }
+}
+
+export class WorkerOptions {
+  requestFunc: (arg: JobRequest) => Promise<void>;
+  cpuLoadFunc: () => number;
+  namespace: string;
+  permissions: WorkerPermissions;
+  workerType: JobType;
+  maxRetry: number;
+  wsURL: string;
+  apiKey?: string;
+  apiSecret?: string;
+  host: string;
+  port: number;
+
+  constructor({
+    requestFunc,
+    cpuLoadFunc = cpuLoad,
+    namespace = 'default',
+    permissions = new WorkerPermissions(),
+    workerType = JobType.JT_PUBLISHER,
+    maxRetry = MAX_RECONNECT_ATTEMPTS,
+    wsURL = 'ws://localhost:7880',
+    apiKey = undefined,
+    apiSecret = undefined,
+    host = 'localhost',
+    port = 8081,
+  }: {
+    requestFunc: (arg: JobRequest) => Promise<void>;
+    cpuLoadFunc?: () => number;
+    namespace?: string;
+    permissions?: WorkerPermissions;
+    workerType?: JobType;
+    maxRetry?: number;
+    wsURL?: string;
+    apiKey?: string;
+    apiSecret?: string;
+    host?: string;
+    port?: number;
+  }) {
+    this.requestFunc = requestFunc;
+    this.cpuLoadFunc = cpuLoadFunc;
+    this.namespace = namespace;
+    this.permissions = permissions;
+    this.workerType = workerType;
+    this.maxRetry = maxRetry;
+    this.wsURL = wsURL;
+    this.apiKey = apiKey;
+    this.apiSecret = apiSecret;
+    this.host = host;
+    this.port = port;
+  }
+}
+
+class ActiveJob {
+  job: Job;
+  acceptData: AcceptData;
+
+  constructor(job: Job, acceptData: AcceptData) {
+    this.job = job;
+    this.acceptData = acceptData;
+  }
+}
+
+type AssignmentPair = {
+  // this string is the JSON string version of the JobAssignment.
+  // we keep it around to unpack it again in the child, because we can't pass Job directly.
+  raw: string;
+  asgn: JobAssignment;
+};
+
+class PendingAssignment {
+  promise = new Promise<AssignmentPair>((resolve) => {
+    this.resolve = resolve; // oh, JavaScript.
+  });
+  resolve(arg: AssignmentPair) {
+    arg;
+  }
+}
+
+export class Worker {
+  opts: WorkerOptions;
+  #id = 'unregistered';
+  session: WebSocket | undefined = undefined;
+  closed = false;
+  httpServer: HTTPServer;
+  logger = log.child({ version });
+  event = new EventEmitter();
+  pending: { [id: string]: { value: PendingAssignment } } = {};
+  processes: { [id: string]: { proc: JobProcess; activeJob: ActiveJob } } = {};
+
+  constructor(opts: WorkerOptions) {
+    opts.wsURL = opts.wsURL || process.env.LIVEKIT_URL || '';
+    opts.apiKey = opts.apiKey || process.env.LIVEKIT_API_KEY || '';
+    opts.apiSecret = opts.apiSecret || process.env.LIVEKIT_API_SECRET || '';
+
+    this.opts = opts;
+    this.httpServer = new HTTPServer(opts.host, opts.port);
+  }
+
+  get id(): string {
+    return this.#id;
+  }
+
+  async run() {
+    this.logger.info('starting worker');
+
+    if (this.opts.wsURL === '') throw new Error('--url is required, or set LIVEKIT_URL env var');
+    if (this.opts.apiKey === '')
+      throw new Error('--api-key is required, or set LIVEKIT_API_KEY env var');
+    if (this.opts.apiSecret === '')
+      throw new Error('--api-secret is required, or set LIVEKIT_API_SECRET env var');
+
+    const workerWS = async () => {
+      let retries = 0;
+      while (!this.closed) {
+        const token = new AccessToken(this.opts.apiKey, this.opts.apiSecret);
+        token.addGrant({ agent: true });
+        const jwt = await token.toJwt();
+
+        const url = new URL(this.opts.wsURL);
+        url.protocol = url.protocol.replace('http', 'ws');
+        this.session = new WebSocket(url + 'agent', {
+          headers: { authorization: 'Bearer ' + jwt },
+        });
+
+        try {
+          await new Promise((resolve, reject) => {
+            this.session!.on('open', resolve);
+            this.session!.on('error', (error) => reject(error));
+            this.session!.on('close', (code) => reject(`WebSocket returned ${code}`));
+          });
+
+          this.runWS(this.session!);
+          return;
+        } catch (e) {
+          if (this.closed) return;
+          if (retries >= this.opts.maxRetry) {
+            throw new Error(`failed to connect to LiveKit server after ${retries} attempts: ${e}`);
+          }
+
+          retries++;
+          const delay = Math.min(retries * 2, 10);
+
+          this.logger.warn(
+            `failed to connect to LiveKit server, retrying in ${delay} seconds: ${e} (${retries}/${this.opts.maxRetry})`,
+          );
+
+          await new Promise((resolve) => setTimeout(resolve, delay * 1000));
+        }
+      }
+    };
+
+    await Promise.all([workerWS(), this.httpServer.run()]);
+  }
+
+  startProcess(job: Job, acceptData: AcceptData, raw: string) {
+    const proc = new JobProcess(job, acceptData, raw, this.opts.wsURL);
+    this.processes[job.id] = { proc, activeJob: new ActiveJob(job, acceptData) };
+    proc
+      .run()
+      .catch((e) => {
+        proc.logger.error(`error running job process ${proc.job.id}: ${e}`);
+      })
+      .finally(() => {
+        proc.clear();
+        delete this.processes[job.id];
+      });
+  }
+
+  runWS(ws: WebSocket) {
+    let closingWS = false;
+
+    const send = (msg: WorkerMessage) => {
+      if (closingWS) {
+        this.event.off('worker_msg', send);
+        return;
+      }
+      ws.send(msg.toBinary());
+    };
+    this.event.on('worker_msg', send);
+
+    ws.addEventListener('close', () => {
+      closingWS = true;
+      this.logger.error('worker connection closed unexpectedly');
+      this.close();
+    });
+
+    ws.addEventListener('message', (event) => {
+      if (event.type !== 'message') {
+        this.logger.warn('unexpected message type: ' + event.type);
+        return;
+      }
+
+      const msg = new ServerMessage();
+      msg.fromBinary(event.data as Uint8Array);
+      switch (msg.message.case) {
+        case 'register': {
+          this.#id = msg.message.value.workerId;
+          log
+            .child({ id: this.id, server_info: msg.message.value.serverInfo })
+            .info('registered worker');
+          break;
+        }
+        case 'availability': {
+          this.availability(msg.message.value);
+          break;
+        }
+        case 'assignment': {
+          const job = msg.message.value.job!;
+          if (job.id in this.pending) {
+            const task = this.pending[job.id];
+            delete this.pending[job.id];
+            task.value.resolve({
+              asgn: msg.message.value,
+              raw: msg.toJsonString(),
+            });
+          } else {
+            log.child({ job }).warn('received assignment for unknown job ' + job.id);
+          }
+          break;
+        }
+      }
+    });
+
+    this.event.emit(
+      'worker_msg',
+      new WorkerMessage({
+        message: {
+          case: 'register',
+          value: {
+            type: this.opts.workerType,
+            namespace: this.opts.namespace,
+            allowedPermissions: new ParticipantPermission({
+              canPublish: this.opts.permissions.canPublish,
+              canSubscribe: this.opts.permissions.canSubscribe,
+              canPublishData: this.opts.permissions.canPublishData,
+              hidden: this.opts.permissions.hidden,
+              agent: true,
+            }),
+            version,
+          },
+        },
+      }),
+    );
+
+    const loadMonitor = setInterval(() => {
+      if (closingWS) clearInterval(loadMonitor);
+      this.event.emit(
+        'worker_msg',
+        new WorkerMessage({
+          message: {
+            case: 'updateWorker',
+            value: {
+              load: cpuLoad(),
+            },
+          },
+        }),
+      );
+    }, LOAD_INTERVAL);
+  }
+
+  async availability(msg: AvailabilityRequest) {
+    const tx = new EventEmitter();
+    const req = new JobRequest(msg.job!, tx);
+
+    tx.on('recv', async (av: AvailRes) => {
+      const msg = new WorkerMessage({
+        message: {
+          case: 'availability',
+          value: {
+            available: av.avail,
+            jobId: req.id,
+            participantIdentity: av.data?.identity,
+            participantName: av.data?.name,
+            participantMetadata: av.data?.metadata,
+          },
+        },
+      });
+
+      this.pending[req.id] = { value: new PendingAssignment() };
+      this.event.emit('worker_msg', msg);
+      if (!av.avail) return;
+
+      const timer = setTimeout(() => {
+        log.child({ req }).warn(`assignment for job ${req.id} timed out`);
+        return;
+      }, ASSIGNMENT_TIMEOUT);
+      this.pending[req.id].value.promise.then(({ asgn, raw }) => {
+        clearTimeout(timer);
+        this.startProcess(asgn!.job!, av.data!, raw);
+      });
+    });
+
+    try {
+      this.opts.requestFunc(req);
+    } catch (e) {
+      log.child({ req }).error(`user request handler for job ${req.id} failed`);
+    } finally {
+      if (!req.answered) {
+        log.child({ req }).error(`no answer for job ${req.id}, automatically rejecting the job`);
+        this.event.emit(
+          'worker_msg',
+          new WorkerMessage({
+            message: {
+              case: 'availability',
+              value: {
+                available: false,
+              },
+            },
+          }),
+        );
+      }
+    }
+  }
+
+  async close() {
+    if (this.closed) return;
+    this.closed = true;
+    this.logger.debug('shutting down worker');
+    await this.httpServer.close();
+    for await (const value of Object.values(this.processes)) {
+      await value.proc.close();
+    }
+    this.session?.close();
+  }
+}

package/tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "extends": "../tsconfig.json",
+  "include": ["./src"],
+  "compilerOptions": {
+    // match output dir to input dir. e.g. dist/index instead of dist/src/index
+    "rootDir": "./src",
+    "declarationDir": "./dist",
+    "outDir": "./dist"
+  }
+}