@livekit/agents-plugin-baseten 0.0.0-next-20260624041820

package/src/llm.ts

@@ -0,0 +1,195 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Baseten LLM plugin for LiveKit Agents
+ * Configures the OpenAI plugin to work with Baseten's OpenAI-compatible API
+ */
+import type { APIConnectOptions } from '@livekit/agents';
+import { DEFAULT_API_CONNECT_OPTIONS, inference, llm } from '@livekit/agents';
+import { OpenAI } from 'openai';
+import type { BasetenLLMOptions } from './types.js';
+
+export interface LLMOptions {
+  model: string;
+  apiKey?: string;
+  baseURL?: string;
+  user?: string;
+  temperature?: number;
+  topP?: number;
+  presencePenalty?: number;
+  frequencyPenalty?: number;
+  client?: OpenAI;
+  toolChoice?: llm.ToolChoice;
+  parallelToolCalls?: boolean;
+  metadata?: Record<string, string>;
+  maxCompletionTokens?: number;
+  serviceTier?: string;
+  store?: boolean;
+  strictToolSchema?: boolean;
+}
+
+const defaultLLMOptions: LLMOptions = {
+  model: 'openai/gpt-4o-mini',
+  apiKey: process.env.OPENAI_API_KEY,
+  parallelToolCalls: true,
+  strictToolSchema: false,
+};
+
+export class OpenAILLM extends llm.LLM {
+  #opts: LLMOptions;
+  #client: OpenAI;
+  #providerFmt: llm.ProviderFormat;
+
+  constructor(
+    opts: Partial<LLMOptions> = defaultLLMOptions,
+    providerFmt: llm.ProviderFormat = 'openai',
+  ) {
+    super();
+
+    this.#opts = { ...defaultLLMOptions, ...opts };
+    this.#providerFmt = providerFmt;
+    if (this.#opts.apiKey === undefined) {
+      throw new Error('OpenAI API key is required, whether as an argument or as $OPENAI_API_KEY');
+    }
+
+    this.#client =
+      this.#opts.client ||
+      new OpenAI({
+        baseURL: opts.baseURL,
+        apiKey: opts.apiKey,
+      });
+  }
+
+  label(): string {
+    return 'openai.LLM';
+  }
+
+  get model(): string {
+    return this.#opts.model;
+  }
+
+  chat({
+    chatCtx,
+    toolCtx: toolCtxInput,
+    connOptions = DEFAULT_API_CONNECT_OPTIONS,
+    parallelToolCalls,
+    toolChoice,
+    extraKwargs,
+  }: {
+    chatCtx: llm.ChatContext;
+    toolCtx?: llm.ToolContextLike;
+    connOptions?: APIConnectOptions;
+    parallelToolCalls?: boolean;
+    toolChoice?: llm.ToolChoice;
+    extraKwargs?: Record<string, unknown>;
+  }): inference.LLMStream {
+    const toolCtx = llm.toToolContext(toolCtxInput);
+    const extras: Record<string, unknown> = { ...extraKwargs };
+
+    if (this.#opts.metadata) {
+      extras.metadata = this.#opts.metadata;
+    }
+
+    if (this.#opts.user) {
+      extras.user = this.#opts.user;
+    }
+
+    if (this.#opts.maxCompletionTokens) {
+      extras.max_completion_tokens = this.#opts.maxCompletionTokens;
+    }
+
+    if (this.#opts.temperature !== undefined) {
+      extras.temperature = this.#opts.temperature;
+    }
+
+    if (this.#opts.topP !== undefined) {
+      extras.top_p = this.#opts.topP;
+    }
+
+    if (this.#opts.presencePenalty !== undefined) {
+      extras.presence_penalty = this.#opts.presencePenalty;
+    }
+
+    if (this.#opts.frequencyPenalty !== undefined) {
+      extras.frequency_penalty = this.#opts.frequencyPenalty;
+    }
+
+    if (this.#opts.serviceTier) {
+      extras.service_tier = this.#opts.serviceTier;
+    }
+
+    if (this.#opts.store !== undefined) {
+      extras.store = this.#opts.store;
+    }
+
+    parallelToolCalls =
+      parallelToolCalls !== undefined ? parallelToolCalls : this.#opts.parallelToolCalls;
+    if (
+      toolCtx &&
+      Object.keys(toolCtx.functionTools).length > 0 &&
+      parallelToolCalls !== undefined
+    ) {
+      extras.parallel_tool_calls = parallelToolCalls;
+    }
+
+    toolChoice = toolChoice !== undefined ? toolChoice : this.#opts.toolChoice;
+    if (toolChoice) {
+      extras.tool_choice = toolChoice;
+    }
+
+    return new LLMStream(this as unknown as inference.LLM, {
+      model: this.#opts.model,
+      providerFmt: this.#providerFmt,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      client: this.#client as any,
+      chatCtx,
+      toolCtx,
+      connOptions,
+      modelOptions: extras,
+      strictToolSchema: this.#opts.strictToolSchema || false,
+      gatewayOptions: undefined, // OpenAI plugin doesn't use gateway authentication
+    });
+  }
+}
+
+export class LLMStream extends inference.LLMStream {}
+
+export class LLM extends OpenAILLM {
+  constructor(opts: BasetenLLMOptions) {
+    const apiKey = opts.apiKey ?? process.env.BASETEN_API_KEY;
+    if (!apiKey) {
+      throw new Error(
+        'Baseten API key is required. Set BASETEN_API_KEY environment variable or pass apiKey in options.',
+      );
+    }
+
+    if (!opts.model) {
+      throw new Error(
+        'Model is required. Please specify a model name (e.g., "openai/gpt-4o-mini").',
+      );
+    }
+
+    const model = opts.model;
+
+    // Configure the OpenAI plugin with Baseten's endpoint
+    super({
+      model,
+      apiKey,
+      baseURL: 'https://inference.baseten.co/v1',
+      temperature: opts.temperature,
+      topP: opts.topP,
+      presencePenalty: opts.presencePenalty,
+      frequencyPenalty: opts.frequencyPenalty,
+      user: opts.user,
+      maxCompletionTokens: opts.maxTokens,
+      toolChoice: opts.toolChoice,
+      parallelToolCalls: opts.parallelToolCalls,
+    });
+  }
+
+  label(): string {
+    return 'baseten.LLM';
+  }
+}

package/src/stt.test.ts

@@ -0,0 +1,22 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { VAD } from '@livekit/agents-plugin-silero';
+import { stt } from '@livekit/agents-plugins-test';
+import { describe, it } from 'vitest';
+import { STT } from './stt.js';
+
+const hasBasetenStreamingConfig = Boolean(
+  process.env.BASETEN_API_KEY &&
+    (process.env.BASETEN_MODEL_ENDPOINT || process.env.BASETEN_STT_MODEL_ID),
+);
+
+if (hasBasetenStreamingConfig) {
+  describe('Baseten', async () => {
+    await stt(new STT(), await VAD.load(), { streaming: true });
+  });
+} else {
+  describe('Baseten', () => {
+    it.skip('requires Baseten streaming credentials/config', () => {});
+  });
+}

package/src/stt.ts

@@ -0,0 +1,335 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import {
+  type AudioBuffer,
+  AudioByteStream,
+  Task,
+  log,
+  normalizeLanguage,
+  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!,
+      alignedTranscript: 'word',
+    });
+
+    const apiKey = opts.apiKey ?? process.env.BASETEN_API_KEY;
+    const modelEndpoint = opts.modelEndpoint ?? process.env.BASETEN_MODEL_ENDPOINT;
+    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 (!modelEndpoint && !modelId) {
+      throw new Error(
+        'Baseten model endpoint is required, either pass it as `modelEndpoint` or set $BASETEN_MODEL_ENDPOINT',
+      );
+    }
+
+    this.#opts = {
+      ...defaultSTTOptions,
+      ...opts,
+      apiKey,
+      modelEndpoint,
+      modelId,
+      audioLanguage: normalizeLanguage((opts.audioLanguage ?? defaultSTTOptions.audioLanguage)!),
+    } 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,
+      audioLanguage:
+        opts.audioLanguage !== undefined
+          ? normalizeLanguage(opts.audioLanguage)
+          : this.#opts.audioLanguage,
+    };
+  }
+
+  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 {
+    if (this.#opts.modelEndpoint) {
+      return this.#opts.modelEndpoint;
+    }
+    // Fallback to constructing URL from modelId (deprecated)
+    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
+    // Note: Baseten server expects 'vad_params' and 'streaming_whisper_params' field names
+    // (not 'streaming_vad_config', 'streaming_params', 'whisper_params' as in older versions)
+    const metadata = {
+      vad_params: {
+        threshold: this.#opts.vadThreshold,
+        min_silence_duration_ms: this.#opts.vadMinSilenceDurationMs,
+        speech_pad_ms: this.#opts.vadSpeechPadMs,
+      },
+      streaming_whisper_params: {
+        encoding: this.#opts.encoding ?? 'pcm_s16le',
+        sample_rate: this.#opts.sampleRate ?? 16000,
+        enable_partial_transcripts: false,
+        audio_language: this.#opts.audioLanguage ?? 'en',
+        show_word_timestamps: true,
+      },
+    };
+
+    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);
+            const isFinal = msg.is_final ?? true;
+            const segments = msg.segments ?? [];
+            const transcript = msg.transcript ?? '';
+            const confidence = msg.confidence ?? 0.0;
+            const languageCode = normalizeLanguage(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 });
+            }
+
+            // Note: Baseten uses 'start_time' and 'end_time' field names (with underscores)
+            const startTime =
+              segments.length > 0
+                ? (segments[0].start_time ?? 0.0) + this.startTimeOffset
+                : this.startTimeOffset;
+            const endTime =
+              segments.length > 0
+                ? (segments[segments.length - 1].end_time ?? 0.0) + this.startTimeOffset
+                : this.startTimeOffset;
+
+            // Note: Baseten returns segments (chunks) which we treat as words for aligned transcripts
+            const words = segments.map(
+              (segment: { text?: string; start_time?: number; end_time?: number }) => ({
+                text: segment.text ?? '',
+                startTime: (segment.start_time ?? 0.0) + this.startTimeOffset,
+                endTime: (segment.end_time ?? 0.0) + this.startTimeOffset,
+                startTimeOffset: this.startTimeOffset,
+                confidence: confidence,
+              }),
+            );
+
+            const speechData: stt.SpeechData = {
+              language: languageCode,
+              text: transcript,
+              startTime,
+              endTime,
+              confidence,
+              words: words.length > 0 ? words : undefined,
+            };
+
+            // 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,21 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { tts } from '@livekit/agents-plugins-test';
+import { describe, it } from 'vitest';
+import { STT } from './stt.js';
+import { TTS } from './tts.js';
+
+const hasBasetenStreamingConfig = Boolean(
+  process.env.BASETEN_API_KEY && process.env.BASETEN_MODEL_ENDPOINT,
+);
+
+if (hasBasetenStreamingConfig) {
+  describe('Baseten', async () => {
+    await tts(new TTS(), new STT(), { streaming: false });
+  });
+} else {
+  describe('Baseten', () => {
+    it.skip('requires Baseten streaming credentials/config', () => {});
+  });
+}