@livekit/agents 0.2.0 → 0.3.0

package/src/multimodal/multimodal_agent.ts

@@ -0,0 +1,428 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type {
+  LocalTrackPublication,
+  RemoteAudioTrack,
+  RemoteParticipant,
+  Room,
+} from '@livekit/rtc-node';
+import {
+  AudioSource,
+  AudioStream,
+  LocalAudioTrack,
+  RoomEvent,
+  TrackPublishOptions,
+  TrackSource,
+} from '@livekit/rtc-node';
+import { EventEmitter } from 'events';
+import { AudioByteStream } from '../audio.js';
+import type * as llm from '../llm/index.js';
+import { log } from '../log.js';
+import { BasicTranscriptionForwarder } from '../transcription.js';
+import { findMicroTrackId } from '../utils.js';
+import { AgentPlayout, type PlayoutHandle } from './agent_playout.js';
+
+/**
+ * @internal
+ * @beta
+ */
+export abstract class RealtimeSession extends EventEmitter {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  abstract defaultConversation: any; // openai.realtime.Conversation
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  abstract inputAudioBuffer: any; // openai.realtime.InputAudioBuffer
+  abstract fncCtx: llm.FunctionContext | undefined;
+}
+
+/**
+ * @internal
+ * @beta
+ */
+export abstract class RealtimeModel {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  abstract session(options: any): RealtimeSession; // openai.realtime.ModelOptions
+  abstract close(): Promise<void>;
+  abstract sampleRate: number;
+  abstract numChannels: number;
+  abstract inFrameSize: number;
+  abstract outFrameSize: number;
+}
+
+export type AgentState = 'initializing' | 'thinking' | 'listening' | 'speaking';
+export const AGENT_STATE_ATTRIBUTE = 'lk.agent.state';
+
+/** @beta */
+export class MultimodalAgent {
+  model: RealtimeModel;
+  room: Room | null = null;
+  linkedParticipant: RemoteParticipant | null = null;
+  subscribedTrack: RemoteAudioTrack | null = null;
+  readMicroTask: { promise: Promise<void>; cancel: () => void } | null = null;
+
+  constructor({
+    model,
+    fncCtx,
+  }: {
+    model: RealtimeModel;
+    fncCtx?: llm.FunctionContext | undefined;
+  }) {
+    this.model = model;
+    this.#fncCtx = fncCtx;
+  }
+
+  #participant: RemoteParticipant | string | null = null;
+  #agentPublication: LocalTrackPublication | null = null;
+  #localTrackSid: string | null = null;
+  #localSource: AudioSource | null = null;
+  #agentPlayout: AgentPlayout | null = null;
+  #playingHandle: PlayoutHandle | undefined = undefined;
+  #logger = log();
+  #session: RealtimeSession | null = null;
+  #fncCtx: llm.FunctionContext | undefined = undefined;
+
+  #_started: boolean = false;
+  #_pendingFunctionCalls: Set<string> = new Set();
+  #_speaking: boolean = false;
+
+  get fncCtx(): llm.FunctionContext | undefined {
+    return this.#fncCtx!;
+  }
+
+  set fncCtx(ctx: llm.FunctionContext | undefined) {
+    this.#fncCtx = ctx;
+    if (this.#session) {
+      this.#session.fncCtx = ctx;
+    }
+  }
+
+  get #pendingFunctionCalls(): Set<string> {
+    return this.#_pendingFunctionCalls;
+  }
+
+  set #pendingFunctionCalls(calls: Set<string>) {
+    this.#_pendingFunctionCalls = calls;
+    this.#updateState();
+  }
+
+  get #speaking(): boolean {
+    return this.#_speaking;
+  }
+
+  set #speaking(isSpeaking: boolean) {
+    this.#_speaking = isSpeaking;
+    this.#updateState();
+  }
+
+  get #started(): boolean {
+    return this.#_started;
+  }
+
+  set #started(started: boolean) {
+    this.#_started = started;
+    this.#updateState();
+  }
+
+  start(
+    room: Room,
+    participant: RemoteParticipant | string | null = null,
+  ): Promise<RealtimeSession> {
+    return new Promise(async (resolve, reject) => {
+      if (this.#started) {
+        this.#logger.warn('MultimodalAgent already started');
+        resolve(this.#session!); // TODO: throw error?
+        return;
+      }
+      this.#updateState();
+
+      room.on(RoomEvent.ParticipantConnected, (participant: RemoteParticipant) => {
+        if (!this.linkedParticipant) {
+          return;
+        }
+
+        this.#linkParticipant(participant.identity);
+      });
+
+      this.room = room;
+      this.#participant = participant;
+
+      this.#localSource = new AudioSource(this.model.sampleRate, this.model.numChannels);
+      this.#agentPlayout = new AgentPlayout(
+        this.#localSource,
+        this.model.sampleRate,
+        this.model.numChannels,
+        this.model.inFrameSize,
+        this.model.outFrameSize,
+      );
+      const track = LocalAudioTrack.createAudioTrack('assistant_voice', this.#localSource);
+      const options = new TrackPublishOptions();
+      options.source = TrackSource.SOURCE_MICROPHONE;
+      this.#agentPublication = (await room.localParticipant?.publishTrack(track, options)) || null;
+      if (!this.#agentPublication) {
+        this.#logger.error('Failed to publish track');
+        reject(new Error('Failed to publish track'));
+        return;
+      }
+
+      await this.#agentPublication.waitForSubscription();
+
+      if (participant) {
+        if (typeof participant === 'string') {
+          this.#linkParticipant(participant);
+        } else {
+          this.#linkParticipant(participant.identity);
+        }
+      } else {
+        // No participant specified, try to find the first participant in the room
+        for (const participant of room.remoteParticipants.values()) {
+          this.#linkParticipant(participant.identity);
+          break;
+        }
+      }
+
+      this.#session = this.model.session({ fncCtx: this.#fncCtx });
+      this.#started = true;
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      this.#session.on('response_content_added', (message: any) => {
+        // openai.realtime.RealtimeContent
+        const trFwd = new BasicTranscriptionForwarder(
+          this.room!,
+          this.room!.localParticipant!.identity,
+          this.#getLocalTrackSid()!,
+          message.responseId,
+        );
+
+        const handle = this.#agentPlayout?.play(
+          message.itemId,
+          message.contentIndex,
+          trFwd,
+          message.textStream,
+          message.audioStream,
+        );
+        if (handle) {
+          this.#speaking = true;
+          handle.on('done', () => {
+            if (this.#playingHandle == handle) {
+              this.#speaking = false;
+            }
+          });
+        }
+        this.#playingHandle = handle;
+      });
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      this.#session.on('input_speech_committed', (ev: any) => {
+        // openai.realtime.InputSpeechCommittedEvent
+        const participantIdentity = this.linkedParticipant?.identity;
+        const trackSid = this.subscribedTrack?.sid;
+        if (participantIdentity && trackSid) {
+          this.#publishTranscription(participantIdentity, trackSid, '…', false, ev.itemId);
+        } else {
+          this.#logger.error('Participant or track not set');
+        }
+      });
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      this.#session.on('input_speech_transcription_completed', (ev: any) => {
+        // openai.realtime.InputSpeechTranscriptionCompletedEvent
+        const transcription = ev.transcript;
+        const participantIdentity = this.linkedParticipant?.identity;
+        const trackSid = this.subscribedTrack?.sid;
+        if (participantIdentity && trackSid) {
+          this.#publishTranscription(participantIdentity, trackSid, transcription, true, ev.itemId);
+        } else {
+          this.#logger.error('Participant or track not set');
+        }
+      });
+
+      this.#session.on('input_speech_started', (ev: any) => {
+        if (this.#playingHandle && !this.#playingHandle.done) {
+          this.#playingHandle.interrupt();
+
+          this.#session!.defaultConversation.item.truncate(
+            this.#playingHandle.itemId,
+            this.#playingHandle.contentIndex,
+            Math.floor((this.#playingHandle.audioSamples / 24000) * 1000),
+          );
+
+          this.#playingHandle = undefined;
+        }
+
+        const participantIdentity = this.linkedParticipant?.identity;
+        const trackSid = this.subscribedTrack?.sid;
+        if (participantIdentity && trackSid) {
+          this.#publishTranscription(participantIdentity, trackSid, '…', false, ev.itemId);
+        }
+      });
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      this.#session.on('function_call_started', (ev: any) => {
+        this.#pendingFunctionCalls.add(ev.callId);
+        this.#updateState();
+      });
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      this.#session.on('function_call_completed', (ev: any) => {
+        this.#pendingFunctionCalls.delete(ev.callId);
+        this.#updateState();
+      });
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      this.#session.on('function_call_failed', (ev: any) => {
+        this.#pendingFunctionCalls.delete(ev.callId);
+        this.#updateState();
+      });
+
+      resolve(this.#session!);
+    });
+  }
+
+  // TODO
+  // close() {
+  //   if (!this.connected || !this.ws) return;
+  //   this.logger.debug('stopping assistant');
+  //   this.ws.close();
+  // }
+
+  #linkParticipant(participantIdentity: string): void {
+    if (!this.room) {
+      this.#logger.error('Room is not set');
+      return;
+    }
+
+    this.linkedParticipant = this.room.remoteParticipants.get(participantIdentity) || null;
+    if (!this.linkedParticipant) {
+      this.#logger.error(`Participant with identity ${participantIdentity} not found`);
+      return;
+    }
+
+    if (this.linkedParticipant.trackPublications.size > 0) {
+      this.#subscribeToMicrophone();
+    } else {
+      this.room.on(RoomEvent.TrackPublished, () => {
+        this.#subscribeToMicrophone();
+      });
+    }
+  }
+
+  #subscribeToMicrophone(): void {
+    const readAudioStreamTask = async (audioStream: AudioStream) => {
+      const bstream = new AudioByteStream(
+        this.model.sampleRate,
+        this.model.numChannels,
+        this.model.inFrameSize,
+      );
+
+      for await (const frame of audioStream) {
+        const audioData = frame.data;
+        for (const frame of bstream.write(audioData.buffer)) {
+          this.#session!.inputAudioBuffer.append(frame);
+        }
+      }
+    };
+
+    if (!this.linkedParticipant) {
+      this.#logger.error('Participant is not set');
+      return;
+    }
+
+    for (const publication of this.linkedParticipant.trackPublications.values()) {
+      if (publication.source !== TrackSource.SOURCE_MICROPHONE) {
+        continue;
+      }
+
+      if (!publication.subscribed) {
+        publication.setSubscribed(true);
+      }
+
+      const track = publication.track;
+
+      if (track && track !== this.subscribedTrack) {
+        this.subscribedTrack = track!;
+
+        if (this.readMicroTask) {
+          this.readMicroTask.cancel();
+        }
+
+        let cancel: () => void;
+        this.readMicroTask = {
+          promise: new Promise<void>((resolve, reject) => {
+            cancel = () => {
+              reject(new Error('Task cancelled'));
+            };
+            readAudioStreamTask(
+              new AudioStream(track, this.model.sampleRate, this.model.numChannels),
+            )
+              .then(resolve)
+              .catch(reject);
+          }),
+          cancel: () => cancel(),
+        };
+      }
+    }
+  }
+
+  #getLocalTrackSid(): string | null {
+    if (!this.#localTrackSid && this.room && this.room.localParticipant) {
+      this.#localTrackSid = findMicroTrackId(this.room, this.room.localParticipant?.identity);
+    }
+    return this.#localTrackSid;
+  }
+
+  #publishTranscription(
+    participantIdentity: string,
+    trackSid: string,
+    text: string,
+    isFinal: boolean,
+    id: string,
+  ): void {
+    this.#logger.debug(
+      `Publishing transcription ${participantIdentity} ${trackSid} ${text} ${isFinal} ${id}`,
+    );
+    if (!this.room?.localParticipant) {
+      this.#logger.error('Room or local participant not set');
+      return;
+    }
+
+    this.room.localParticipant.publishTranscription({
+      participantIdentity,
+      trackSid,
+      segments: [
+        {
+          text,
+          final: isFinal,
+          id,
+          startTime: BigInt(0),
+          endTime: BigInt(0),
+          language: '',
+        },
+      ],
+    });
+  }
+
+  #updateState() {
+    let newState: AgentState = 'initializing';
+    if (this.#pendingFunctionCalls.size > 0) {
+      newState = 'thinking';
+    } else if (this.#speaking) {
+      newState = 'speaking';
+    } else if (this.#started) {
+      newState = 'listening';
+    }
+
+    this.#setState(newState);
+  }
+
+  #setState(state: AgentState) {
+    if (this.room?.isConnected && this.room.localParticipant) {
+      const currentState = this.room.localParticipant.attributes[AGENT_STATE_ATTRIBUTE];
+      if (currentState !== state) {
+        this.room.localParticipant!.setAttributes({
+          [AGENT_STATE_ATTRIBUTE]: state,
+        });
+        this.#logger.debug(`${AGENT_STATE_ATTRIBUTE}: ${currentState} ->${state}`);
+      }
+    }
+  }
+}

package/src/transcription.ts

@@ -0,0 +1,128 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { AudioFrame, Room } from '@livekit/rtc-node';
+import { log } from './log.js';
+
+export interface TranscriptionForwarder {
+  start(): void;
+  pushAudio(frame: AudioFrame): void;
+  pushText(text: string): void;
+  markTextComplete(): void;
+  markAudioComplete(): void;
+  close(interrupt: boolean): Promise<void>;
+  currentCharacterIndex: number;
+}
+
+export class BasicTranscriptionForwarder implements TranscriptionForwarder {
+  #room: Room;
+  #participantIdentity: string;
+  #trackSid: string;
+  #currentText: string = '';
+  #totalAudioDuration: number = 0;
+  #currentPlayoutTime: number = 0;
+  #DEFAULT_CHARS_PER_SECOND = 16;
+  #charsPerSecond: number = this.#DEFAULT_CHARS_PER_SECOND;
+  #messageId: string;
+  #isRunning: boolean = false;
+  currentCharacterIndex: number = 0;
+
+  constructor(room: Room, participantIdentity: string, trackSid: string, messageId: string) {
+    this.#room = room;
+    this.#participantIdentity = participantIdentity;
+    this.#trackSid = trackSid;
+    this.#messageId = messageId;
+  }
+
+  start(): void {
+    if (!this.#isRunning) {
+      this.#isRunning = true;
+      this.#startPublishingLoop().catch((error) => {
+        log().error('Error in publishing loop:', error);
+        this.#isRunning = false;
+      });
+    }
+  }
+
+  pushAudio(frame: AudioFrame): void {
+    this.#totalAudioDuration += frame.samplesPerChannel / frame.sampleRate;
+  }
+
+  pushText(text: string): void {
+    this.#currentText += text;
+  }
+
+  #textIsComplete: boolean = false;
+  #audioIsComplete: boolean = false;
+
+  markTextComplete(): void {
+    this.#textIsComplete = true;
+    this.#adjustTimingIfBothFinished();
+  }
+
+  markAudioComplete(): void {
+    this.#audioIsComplete = true;
+    this.#adjustTimingIfBothFinished();
+  }
+
+  #adjustTimingIfBothFinished(): void {
+    if (this.#textIsComplete && this.#audioIsComplete) {
+      const actualDuration = this.#totalAudioDuration;
+      if (actualDuration > 0 && this.#currentText.length > 0) {
+        this.#charsPerSecond = this.#currentText.length / actualDuration;
+      }
+    }
+  }
+
+  #computeSleepInterval(): number {
+    return Math.min(Math.max(1 / this.#charsPerSecond, 0.0625), 0.5);
+  }
+
+  async #startPublishingLoop(): Promise<void> {
+    this.#isRunning = true;
+    let sleepInterval = this.#computeSleepInterval();
+    let isComplete = false;
+    while (this.#isRunning && !isComplete) {
+      this.#currentPlayoutTime += sleepInterval;
+      this.currentCharacterIndex = Math.floor(this.#currentPlayoutTime * this.#charsPerSecond);
+      isComplete = this.#textIsComplete && this.currentCharacterIndex >= this.#currentText.length;
+      await this.#publishTranscription(false);
+      if (this.#isRunning && !isComplete) {
+        sleepInterval = this.#computeSleepInterval();
+        await new Promise((resolve) => setTimeout(resolve, sleepInterval * 1000));
+      }
+    }
+
+    if (this.#isRunning) {
+      this.close(false);
+    }
+  }
+
+  async #publishTranscription(final: boolean): Promise<void> {
+    const textToPublish = this.#currentText.slice(0, this.currentCharacterIndex);
+    await this.#room.localParticipant?.publishTranscription({
+      participantIdentity: this.#participantIdentity,
+      trackSid: this.#trackSid,
+      segments: [
+        {
+          text: textToPublish,
+          final: final,
+          id: this.#messageId,
+          startTime: BigInt(0),
+          endTime: BigInt(0),
+          language: '',
+        },
+      ],
+    });
+  }
+
+  async close(interrupt: boolean): Promise<void> {
+    this.#isRunning = false;
+
+    // Publish whatever we had as final
+    if (!interrupt) {
+      this.currentCharacterIndex = this.#currentText.length;
+    }
+    await this.#publishTranscription(true);
+  }
+}