@livekit/agents-plugin-baseten 1.0.31

package/src/stt.ts

@@ -0,0 +1,298 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { type AudioBuffer, AudioByteStream, Task, log, stt, waitForAbort } from '@livekit/agents';
+import type { AudioFrame } from '@livekit/rtc-node';
+import { WebSocket } from 'ws';
+import type { BasetenSttOptions } from './types.js';
+
+const defaultSTTOptions: Partial<BasetenSttOptions> = {
+  environment: 'production',
+  encoding: 'pcm_s16le',
+  sampleRate: 16000,
+  bufferSizeSeconds: 0.032,
+  enablePartialTranscripts: true,
+  partialTranscriptIntervalS: 0.5,
+  finalTranscriptMaxDurationS: 5,
+  audioLanguage: 'en',
+  languageDetectionOnly: false,
+  vadThreshold: 0.5,
+  vadMinSilenceDurationMs: 300,
+  vadSpeechPadMs: 30,
+};
+
+export class STT extends stt.STT {
+  #opts: BasetenSttOptions;
+  #logger = log();
+  label = 'baseten.STT';
+
+  constructor(opts: Partial<BasetenSttOptions> = {}) {
+    super({
+      streaming: true,
+      interimResults: opts.enablePartialTranscripts ?? defaultSTTOptions.enablePartialTranscripts!,
+    });
+
+    const apiKey = opts.apiKey ?? process.env.BASETEN_API_KEY;
+    const modelId = opts.modelId ?? process.env.BASETEN_STT_MODEL_ID;
+
+    if (!apiKey) {
+      throw new Error(
+        'Baseten API key is required, either pass it as `apiKey` or set $BASETEN_API_KEY',
+      );
+    }
+    if (!modelId) {
+      throw new Error(
+        'Baseten model ID is required, either pass it as `modelId` or set $BASETEN_STT_MODEL_ID',
+      );
+    }
+
+    this.#opts = {
+      ...defaultSTTOptions,
+      ...opts,
+      apiKey,
+      modelId,
+    } as BasetenSttOptions;
+  }
+
+  // eslint-disable-next-line
+  async _recognize(_: AudioBuffer): Promise<stt.SpeechEvent> {
+    throw new Error('Recognize is not supported on Baseten STT');
+  }
+
+  updateOptions(opts: Partial<BasetenSttOptions>) {
+    this.#opts = { ...this.#opts, ...opts };
+  }
+
+  stream(): SpeechStream {
+    return new SpeechStream(this, this.#opts);
+  }
+}
+
+export class SpeechStream extends stt.SpeechStream {
+  #opts: BasetenSttOptions;
+  #logger = log();
+  #speaking = false;
+  #requestId = '';
+  label = 'baseten.SpeechStream';
+
+  constructor(stt: STT, opts: BasetenSttOptions) {
+    super(stt, opts.sampleRate);
+    this.#opts = opts;
+    this.closed = false;
+  }
+
+  private getWsUrl(): string {
+    return `wss://model-${this.#opts.modelId}.api.baseten.co/environments/${this.#opts.environment}/websocket`;
+  }
+
+  protected async run() {
+    const maxRetry = 32;
+    let retries = 0;
+
+    while (!this.input.closed && !this.closed) {
+      const url = this.getWsUrl();
+      const headers = {
+        Authorization: `Api-Key ${this.#opts.apiKey}`,
+      };
+
+      const ws = new WebSocket(url, { headers });
+
+      try {
+        await new Promise((resolve, reject) => {
+          ws.on('open', resolve);
+          ws.on('error', (error) => reject(error));
+          ws.on('close', (code) => reject(`WebSocket returned ${code}`));
+        });
+
+        await this.#runWS(ws);
+      } catch (e) {
+        if (!this.closed && !this.input.closed) {
+          if (retries >= maxRetry) {
+            throw new Error(`failed to connect to Baseten after ${retries} attempts: ${e}`);
+          }
+
+          const delay = Math.min(retries * 5, 10);
+          retries++;
+
+          this.#logger.warn(
+            `failed to connect to Baseten, retrying in ${delay} seconds: ${e} (${retries}/${maxRetry})`,
+          );
+          await new Promise((resolve) => setTimeout(resolve, delay * 1000));
+        } else {
+          this.#logger.warn(
+            `Baseten disconnected, connection is closed: ${e} (inputClosed: ${this.input.closed}, isClosed: ${this.closed})`,
+          );
+        }
+      }
+    }
+
+    this.closed = true;
+  }
+
+  async #runWS(ws: WebSocket) {
+    let closing = false;
+
+    // Send initial metadata
+    const metadata = {
+      streaming_vad_config: {
+        threshold: this.#opts.vadThreshold,
+        min_silence_duration_ms: this.#opts.vadMinSilenceDurationMs,
+        speech_pad_ms: this.#opts.vadSpeechPadMs,
+      },
+      streaming_params: {
+        encoding: this.#opts.encoding ?? 'pcm_s16le',
+        sample_rate: this.#opts.sampleRate ?? 16000,
+        enable_partial_transcripts: this.#opts.enablePartialTranscripts,
+        partial_transcript_interval_s: this.#opts.partialTranscriptIntervalS,
+        final_transcript_max_duration_s: this.#opts.finalTranscriptMaxDurationS,
+      },
+      whisper_params: {
+        prompt: this.#opts.prompt,
+        audio_language: this.#opts.audioLanguage ?? 'en',
+        language_detection_only: this.#opts.languageDetectionOnly ?? false,
+      },
+    };
+    ws.send(JSON.stringify(metadata));
+
+    const sendTask = async () => {
+      const sampleRate = this.#opts.sampleRate ?? 16000;
+      const samplesPerChunk = sampleRate === 16000 ? 512 : 256;
+      const audioByteStream = new AudioByteStream(sampleRate, 1, samplesPerChunk);
+
+      try {
+        while (!this.closed) {
+          const result = await this.input.next();
+          if (result.done) {
+            break;
+          }
+
+          const data = result.value;
+
+          let frames: AudioFrame[];
+          if (data === SpeechStream.FLUSH_SENTINEL) {
+            // Flush any remaining buffered audio
+            frames = audioByteStream.flush();
+          } else {
+            if (data.sampleRate !== sampleRate || data.channels !== 1) {
+              throw new Error(
+                `sample rate or channel count mismatch: expected ${sampleRate}Hz/1ch, got ${data.sampleRate}Hz/${data.channels}ch`,
+              );
+            }
+            frames = audioByteStream.write(data.data.buffer as ArrayBuffer);
+          }
+
+          for (const frame of frames) {
+            const buffer = Buffer.from(
+              frame.data.buffer,
+              frame.data.byteOffset,
+              frame.data.byteLength,
+            );
+            ws.send(buffer);
+          }
+        }
+      } finally {
+        closing = true;
+        ws.close();
+      }
+    };
+
+    const listenTask = Task.from(async (controller) => {
+      const listenMessage = new Promise<void>((resolve, reject) => {
+        ws.on('message', (data) => {
+          try {
+            let jsonString: string;
+
+            if (typeof data === 'string') {
+              jsonString = data;
+            } else if (data instanceof Buffer) {
+              jsonString = data.toString('utf-8');
+            } else if (Array.isArray(data)) {
+              jsonString = Buffer.concat(data).toString('utf-8');
+            } else {
+              return;
+            }
+
+            const msg = JSON.parse(jsonString);
+
+            // Parse response format matching Python implementation
+            const isFinal = msg.is_final ?? true;
+            const segments = msg.segments ?? [];
+            const transcript = msg.transcript ?? '';
+            const confidence = msg.confidence ?? 0.0;
+            const languageCode = msg.language_code ?? this.#opts.audioLanguage;
+
+            // Skip if no transcript text
+            if (!transcript) {
+              this.#logger.debug('Received non-transcript message:', msg);
+              return;
+            }
+
+            // Emit START_OF_SPEECH if not already speaking (only for interim or first final)
+            if (!this.#speaking && !isFinal) {
+              this.#speaking = true;
+              this.queue.put({ type: stt.SpeechEventType.START_OF_SPEECH });
+            }
+
+            // Extract timing from segments
+            const startTime = segments.length > 0 ? segments[0].start ?? 0.0 : 0.0;
+            const endTime = segments.length > 0 ? segments[segments.length - 1].end ?? 0.0 : 0.0;
+
+            const speechData: stt.SpeechData = {
+              language: languageCode!,
+              text: transcript,
+              startTime,
+              endTime,
+              confidence,
+            };
+
+            // Handle interim vs final transcripts (matching Python implementation)
+            if (!isFinal) {
+              // Interim transcript
+              this.queue.put({
+                type: stt.SpeechEventType.INTERIM_TRANSCRIPT,
+                alternatives: [speechData],
+              });
+            } else {
+              // Final transcript
+              this.queue.put({
+                type: stt.SpeechEventType.FINAL_TRANSCRIPT,
+                alternatives: [speechData],
+              });
+
+              // Emit END_OF_SPEECH after final transcript
+              if (this.#speaking) {
+                this.#speaking = false;
+                this.queue.put({ type: stt.SpeechEventType.END_OF_SPEECH });
+              }
+            }
+
+            if (this.closed || closing) {
+              resolve();
+            }
+          } catch (err) {
+            this.#logger.error(`STT: Error processing message: ${data}`);
+            reject(err);
+          }
+        });
+
+        ws.on('error', (err) => {
+          if (!closing) {
+            reject(err);
+          }
+        });
+
+        ws.on('close', () => {
+          if (!closing) {
+            resolve();
+          }
+        });
+      });
+
+      await Promise.race([listenMessage, waitForAbort(controller.signal)]);
+    }, this.abortController);
+
+    await Promise.all([sendTask(), listenTask.result]);
+    closing = true;
+    ws.close();
+  }
+}

package/src/tts.test.ts

@@ -0,0 +1,11 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { tts } from '@livekit/agents-plugins-test';
+import { describe } from 'vitest';
+import { STT } from './stt.js';
+import { TTS } from './tts.js';
+
+describe('Baseten', async () => {
+  await tts(new TTS(), new STT(), { streaming: false });
+});

package/src/tts.ts

@@ -0,0 +1,202 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import {
+  type APIConnectOptions,
+  AudioByteStream,
+  shortuuid,
+  tts,
+  waitForAbort,
+} from '@livekit/agents';
+import type { AudioFrame } from '@livekit/rtc-node';
+import type { BasetenTTSOptions } from './types.js';
+
+const defaultTTSOptions: Partial<BasetenTTSOptions> = {
+  voice: 'tara',
+  language: 'en',
+  temperature: 0.6,
+};
+
+/**
+ * Baseten TTS implementation (streaming, 24kHz mono)
+ */
+export class TTS extends tts.TTS {
+  private opts: BasetenTTSOptions;
+  label = 'baseten.TTS';
+  private abortController = new AbortController();
+  constructor(opts: Partial<BasetenTTSOptions> = {}) {
+    /**
+     * Baseten audio is 24kHz mono.
+     * The Orpheus model generates audio chunks that are processed as they arrive,
+     * which reduces latency and improves agent responsiveness.
+     */
+    super(24000, 1, { streaming: false });
+
+    // Apply defaults and environment fallbacks.
+    const apiKey = opts.apiKey ?? process.env.BASETEN_API_KEY;
+    const modelEndpoint = opts.modelEndpoint ?? process.env.BASETEN_MODEL_ENDPOINT;
+
+    if (!apiKey) {
+      throw new Error(
+        'Baseten API key is required, either pass it as `apiKey` or set $BASETEN_API_KEY',
+      );
+    }
+    if (!modelEndpoint) {
+      throw new Error(
+        'Baseten model endpoint is required, either pass it as `modelEndpoint` or set $BASETEN_MODEL_ENDPOINT',
+      );
+    }
+
+    this.opts = {
+      ...defaultTTSOptions,
+      ...opts,
+      apiKey,
+      modelEndpoint,
+    } as BasetenTTSOptions;
+  }
+
+  updateOptions(opts: Partial<Omit<BasetenTTSOptions, 'apiKey' | 'modelEndpoint'>>) {
+    this.opts = {
+      ...this.opts,
+      ...opts,
+    } as BasetenTTSOptions;
+  }
+
+  /**
+   * Synthesize speech for a given piece of text.  Returns a `ChunkedStream`
+   * which will asynchronously fetch audio from Baseten and push frames into
+   * LiveKit's playback pipeline.  If you need to cancel synthesis you can
+   * call {@link ChunkedStream.stop} on the returned object.
+   */
+  synthesize(
+    text: string,
+    connOptions?: APIConnectOptions,
+    abortSignal?: AbortSignal,
+  ): ChunkedStream {
+    return new ChunkedStream(this, text, this.opts, connOptions, abortSignal);
+  }
+
+  stream(): tts.SynthesizeStream {
+    throw new Error('Streaming is not supported on Baseten TTS');
+  }
+
+  async close(): Promise<void> {
+    this.abortController.abort();
+  }
+}
+
+/**
+ * Internal helper that performs the actual HTTP request and converts the
+ * response into audio frames.  It inherits from `tts.ChunkedStream` to
+ * integrate with LiveKit's event and cancellation framework.
+ *
+ * This implementation streams audio chunks as they arrive from the Baseten
+ * model endpoint, processing them incrementally instead of waiting for the
+ * complete response.
+ */
+export class ChunkedStream extends tts.ChunkedStream {
+  label = 'baseten.ChunkedStream';
+  private readonly opts: BasetenTTSOptions;
+
+  constructor(
+    tts: TTS,
+    text: string,
+    opts: BasetenTTSOptions,
+    connOptions?: APIConnectOptions,
+    abortSignal?: AbortSignal,
+  ) {
+    super(text, tts, connOptions, abortSignal);
+    this.opts = opts;
+  }
+
+  /**
+   * Execute the synthesis request.  This method is automatically invoked
+   * by the base class when the stream starts.  It performs a POST request
+   * to the configured `modelEndpoint` with the input text and optional
+   * parameters.  Audio chunks are streamed as they arrive and transformed
+   * into a sequence of `AudioFrame` objects that are enqueued immediately
+   * for playback.
+   */
+  protected async run() {
+    const { apiKey, modelEndpoint, voice, language, temperature, maxTokens } = this.opts;
+    const payload: Record<string, unknown> = {
+      prompt: this.inputText,
+    };
+    if (voice) payload.voice = voice;
+    if (language) payload.language = language;
+    if (temperature !== undefined) payload.temperature = temperature;
+    if (maxTokens !== undefined) payload.max_tokens = maxTokens;
+
+    const headers: Record<string, string> = {
+      Authorization: `Api-Key ${apiKey}`,
+      'Content-Type': 'application/json',
+    };
+
+    const response = await fetch(modelEndpoint, {
+      method: 'POST',
+      headers,
+      body: JSON.stringify(payload),
+      signal: this.abortSignal,
+    });
+
+    if (!response.ok) {
+      let errText: string;
+      try {
+        errText = await response.text();
+      } catch {
+        errText = response.statusText;
+      }
+      throw new Error(`Baseten TTS request failed: ${response.status} ${errText}`);
+    }
+
+    // Stream the response body as chunks arrive
+    if (!response.body) {
+      throw new Error('Response body is not available for streaming');
+    }
+
+    const requestId = shortuuid();
+    const audioByteStream = new AudioByteStream(24000, 1);
+    const reader = response.body.getReader();
+
+    try {
+      let lastFrame: AudioFrame | undefined;
+      const sendLastFrame = (segmentId: string, final: boolean) => {
+        if (lastFrame) {
+          this.queue.put({ requestId, segmentId, frame: lastFrame, final });
+          lastFrame = undefined;
+        }
+      };
+
+      // waitForAbort internally sets up an abort listener on the abort signal
+      // we need to put it outside loop to avoid constant re-registration of the listener
+      const abortPromise = waitForAbort(this.abortSignal);
+
+      while (!this.abortSignal.aborted) {
+        const result = await Promise.race([reader.read(), abortPromise]);
+
+        if (result === undefined) break; // aborted
+
+        const { done, value } = result;
+
+        if (done) {
+          break;
+        }
+
+        // Process the chunk and convert to audio frames
+        // Convert Uint8Array to ArrayBuffer for AudioByteStream
+        const frames = audioByteStream.write(value.buffer);
+
+        for (const frame of frames) {
+          sendLastFrame(requestId, false);
+          lastFrame = frame;
+        }
+      }
+
+      // Send the final frame
+      sendLastFrame(requestId, true);
+    } finally {
+      reader.releaseLock();
+      this.queue.close();
+    }
+  }
+}

package/src/types.ts

@@ -0,0 +1,55 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Baseten plugin types and interfaces
+ */
+
+/**
+ * Options for configuring the Baseten LLM
+ * Since Baseten provides an OpenAI-compatible API, these options
+ * map to standard OpenAI parameters.
+ */
+export interface BasetenLLMOptions {
+  apiKey?: string;
+  model: string;
+  temperature?: number;
+  maxTokens?: number;
+  user?: string;
+  toolChoice?: 'none' | 'auto' | 'required' | { type: 'function'; function: { name: string } };
+  parallelToolCalls?: boolean;
+}
+
+/**
+ * Options for configuring the Baseten STT service
+ */
+export interface BasetenSttOptions {
+  apiKey: string;
+  modelId: string;
+  environment?: string;
+  encoding?: string;
+  sampleRate?: number;
+  bufferSizeSeconds?: number;
+  vadThreshold?: number;
+  vadMinSilenceDurationMs?: number;
+  vadSpeechPadMs?: number;
+  enablePartialTranscripts?: boolean;
+  partialTranscriptIntervalS?: number;
+  finalTranscriptMaxDurationS?: number;
+  audioLanguage?: string;
+  prompt?: string;
+  languageDetectionOnly?: boolean;
+}
+
+/**
+ * Options for configuring the Baseten TTS service
+ */
+export interface BasetenTTSOptions {
+  apiKey: string;
+  modelEndpoint: string;
+  voice?: string;
+  language?: string;
+  temperature?: number;
+  maxTokens?: number;
+}