@livekit/agents 0.0.0-20260120144724

package/src/connection_pool.ts

@@ -0,0 +1,307 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { Mutex } from '@livekit/mutex';
+import { waitForAbort } from './utils.js';
+
+/**
+ * Helper class to manage persistent connections like websockets.
+ */
+export interface ConnectionPoolOptions<T> {
+  /**
+   * Maximum duration in milliseconds before forcing reconnection.
+   * If not set, connections will never expire based on duration.
+   */
+  maxSessionDuration?: number;
+
+  /**
+   * If true, the session will be marked as fresh when get() is called.
+   * Only used when maxSessionDuration is set.
+   */
+  markRefreshedOnGet?: boolean;
+
+  /**
+   * Async callback to create new connections.
+   * @param timeout - Connection timeout in milliseconds
+   * @returns A new connection object
+   */
+  connectCb: (timeout: number) => Promise<T>;
+
+  /**
+   * Optional async callback to close connections.
+   * @param conn - The connection to close
+   */
+  closeCb?: (conn: T) => Promise<void>;
+
+  /**
+   * Default connection timeout in milliseconds.
+   * Defaults to 10000 (10 seconds).
+   */
+  connectTimeout?: number;
+}
+
+/**
+ * Connection pool for managing persistent WebSocket connections.
+ *
+ * Reuses connections efficiently and automatically refreshes them after max duration.
+ * Prevents creating too many connections in a single conversation.
+ */
+export class ConnectionPool<T> {
+  private readonly maxSessionDuration?: number;
+  private readonly markRefreshedOnGet: boolean;
+  private readonly connectCb: (timeout: number) => Promise<T>;
+  private readonly closeCb?: (conn: T) => Promise<void>;
+  private readonly connectTimeout: number;
+
+  // Track connections and their creation timestamps
+  private readonly connections: Map<T, number> = new Map();
+  // Available connections ready for reuse
+  private readonly available: Set<T> = new Set();
+  // Connections queued for closing
+  private readonly toClose: Set<T> = new Set();
+  // Mutex for connection operations
+  private readonly connectLock = new Mutex();
+  // Prewarm task reference
+  private prewarmController?: AbortController;
+
+  constructor(options: ConnectionPoolOptions<T>) {
+    this.maxSessionDuration = options.maxSessionDuration;
+    this.markRefreshedOnGet = options.markRefreshedOnGet ?? false;
+    this.connectCb = options.connectCb;
+    this.closeCb = options.closeCb;
+    this.connectTimeout = options.connectTimeout ?? 10_000;
+  }
+
+  /**
+   * Create a new connection.
+   *
+   * @param timeout - Connection timeout in milliseconds
+   * @returns The new connection object
+   * @throws If connectCb is not provided or connection fails
+   */
+  private async _connect(timeout: number): Promise<T> {
+    const connection = await this.connectCb(timeout);
+    this.connections.set(connection, Date.now());
+    return connection;
+  }
+
+  /**
+   * Drain and close all connections queued for closing.
+   */
+  private async _drainToClose(): Promise<void> {
+    const connectionsToClose = Array.from(this.toClose);
+    this.toClose.clear();
+
+    for (const conn of connectionsToClose) {
+      await this._maybeCloseConnection(conn);
+    }
+  }
+
+  /**
+   * Close a connection if closeCb is provided.
+   *
+   * @param conn - The connection to close
+   */
+  private async _maybeCloseConnection(conn: T): Promise<void> {
+    if (this.closeCb) {
+      await this.closeCb(conn);
+    }
+  }
+
+  private _abortError(): Error {
+    const error = new Error('The operation was aborted.');
+    error.name = 'AbortError';
+    return error;
+  }
+
+  /**
+   * Get an available connection or create a new one if needed.
+   *
+   * @param timeout - Connection timeout in milliseconds
+   * @returns An active connection object
+   */
+  async get(timeout?: number): Promise<T> {
+    const unlock = await this.connectLock.lock();
+    try {
+      await this._drainToClose();
+      const now = Date.now();
+
+      // Try to reuse an available connection that hasn't expired
+      while (this.available.size > 0) {
+        const conn = this.available.values().next().value as T;
+        this.available.delete(conn);
+
+        if (
+          this.maxSessionDuration === undefined ||
+          now - (this.connections.get(conn) ?? 0) <= this.maxSessionDuration
+        ) {
+          if (this.markRefreshedOnGet) {
+            this.connections.set(conn, now);
+          }
+          return conn;
+        }
+
+        // Connection expired; close it now so callers observing get() see it closed promptly.
+        // (Also makes tests deterministic: closeCb should have been called by the time get() resolves.)
+        if (this.connections.has(conn)) {
+          this.connections.delete(conn);
+        }
+        this.toClose.delete(conn);
+        await this._maybeCloseConnection(conn);
+      }
+
+      return await this._connect(timeout ?? this.connectTimeout);
+    } finally {
+      unlock();
+    }
+  }
+
+  /**
+   * Mark a connection as available for reuse.
+   *
+   * If connection has been removed, it will not be added to the pool.
+   *
+   * @param conn - The connection to make available
+   */
+  put(conn: T): void {
+    if (this.connections.has(conn)) {
+      this.available.add(conn);
+      return;
+    }
+  }
+
+  /**
+   * Remove a specific connection from the pool.
+   *
+   * Marks the connection to be closed during the next drain cycle.
+   *
+   * @param conn - The connection to remove
+   */
+  remove(conn: T): void {
+    this.available.delete(conn);
+    if (this.connections.has(conn)) {
+      this.toClose.add(conn);
+      this.connections.delete(conn);
+      // Important for Node websockets: if we just "mark to close later" but remove listeners,
+      // the ws library can buffer incoming frames in memory. Close ASAP in background.
+      void (async () => {
+        const unlock = await this.connectLock.lock();
+        try {
+          if (!this.toClose.has(conn)) return;
+          await this._maybeCloseConnection(conn);
+          this.toClose.delete(conn);
+        } finally {
+          unlock();
+        }
+      })();
+    }
+  }
+
+  /**
+   * Clear all existing connections.
+   *
+   * Marks all current connections to be closed during the next drain cycle.
+   */
+  invalidate(): void {
+    for (const conn of this.connections.keys()) {
+      this.toClose.add(conn);
+    }
+    this.connections.clear();
+    this.available.clear();
+  }
+
+  /**
+   * Initiate prewarming of the connection pool without blocking.
+   *
+   * This method starts a background task that creates a new connection if none exist.
+   * The task automatically cleans itself up when the connection pool is closed.
+   */
+  prewarm(): void {
+    if (this.prewarmController || this.connections.size > 0) {
+      return;
+    }
+
+    const controller = new AbortController();
+    this.prewarmController = controller;
+
+    // Start prewarm in background
+    this._prewarmImpl(controller.signal).catch(() => {
+      // Ignore errors during prewarm
+    });
+  }
+
+  private async _prewarmImpl(signal: AbortSignal): Promise<void> {
+    const unlock = await this.connectLock.lock();
+    try {
+      if (signal.aborted) {
+        return;
+      }
+
+      if (this.connections.size === 0) {
+        const conn = await this._connect(this.connectTimeout);
+        this.available.add(conn);
+      }
+    } finally {
+      unlock();
+    }
+  }
+
+  /**
+   * Get a connection from the pool and automatically return it when done.
+   * Handles abort signals and ensures proper cleanup.
+   *
+   * @param fn - Function to execute with the connection
+   * @param options - Options including timeout and abort signal
+   * @returns The result of the function
+   */
+  async withConnection<R>(
+    fn: (conn: T) => Promise<R>,
+    options?: {
+      timeout?: number;
+      signal?: AbortSignal;
+    },
+  ): Promise<R> {
+    // Check if already aborted before getting connection
+    if (options?.signal?.aborted) {
+      throw this._abortError();
+    }
+
+    const conn = await this.get(options?.timeout);
+
+    const signal = options?.signal;
+
+    try {
+      const fnPromise = fn(conn);
+      const result = signal
+        ? await Promise.race([
+            fnPromise.then((value) => ({ type: 'result' as const, value })),
+            waitForAbort(signal).then(() => ({ type: 'abort' as const })),
+          ]).then((r) => {
+            if (r.type === 'abort') throw this._abortError();
+            return r.value;
+          })
+        : await fnPromise;
+      // Return connection to pool on success
+      this.put(conn);
+      return result;
+    } catch (error) {
+      // Remove connection from pool on error (don't return it)
+      this.remove(conn);
+      throw error;
+    }
+  }
+
+  /**
+   * Close all connections, draining any pending connection closures.
+   */
+  async close(): Promise<void> {
+    // Cancel prewarm task if running
+    if (this.prewarmController) {
+      this.prewarmController.abort();
+      this.prewarmController = undefined;
+    }
+
+    this.invalidate();
+    await this._drainToClose();
+  }
+}

package/src/constants.ts

@@ -0,0 +1,9 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+export const ATTRIBUTE_TRANSCRIPTION_TRACK_ID = 'lk.transcribed_track_id';
+export const ATTRIBUTE_TRANSCRIPTION_FINAL = 'lk.transcription_final';
+export const TOPIC_TRANSCRIPTION = 'lk.transcription';
+export const ATTRIBUTE_TRANSCRIPTION_SEGMENT_ID = 'lk.segment_id';
+export const ATTRIBUTE_PUBLISH_ON_BEHALF = 'lk.publish_on_behalf';
+export const TOPIC_CHAT = 'lk.chat';

package/src/generator.ts

@@ -0,0 +1,38 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import type { JobContext, JobProcess } from './job.js';
+
+/** @see {@link defineAgent} */
+export interface Agent {
+  entry: (ctx: JobContext) => Promise<void>;
+  prewarm?: (proc: JobProcess) => unknown;
+}
+
+/** Helper to check if an object is an agent before running it.
+ *
+ * @internal
+ */
+export function isAgent(obj: unknown): obj is Agent {
+  return (
+    typeof obj === 'object' &&
+    obj !== null &&
+    'entry' in obj &&
+    typeof (obj as Agent).entry === 'function' &&
+    (('prewarm' in obj && typeof (obj as Agent).prewarm === 'function') || !('prewarm' in obj))
+  );
+}
+
+/**
+ * Helper to define an agent according to the required interface.
+ * @example A basic agent with entry and prewarm functions
+ * ```
+ * export default defineAgent({
+ *   entry: async (ctx: JobContext) => { ... },
+ *   prewarm: (proc: JobProcess) => { ... },
+ * })
+ * ```
+ */
+export function defineAgent(agent: Agent): Agent {
+  return agent;
+}

package/src/http_server.ts

@@ -0,0 +1,64 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { type IncomingMessage, type Server, type ServerResponse, createServer } from 'node:http';
+import { log } from './log.js';
+
+const healthCheck = async (res: ServerResponse) => {
+  res.writeHead(200);
+  res.end('OK');
+};
+
+interface WorkerResponse {
+  agent_name: string;
+  worker_type: string;
+  active_jobs: number;
+  sdk_version: string;
+  project_type: string;
+}
+
+export class HTTPServer {
+  host: string;
+  port: number;
+  app: Server;
+  #logger = log();
+
+  constructor(host: string, port: number, workerListener: () => WorkerResponse) {
+    this.host = host;
+    this.port = port;
+
+    this.app = createServer((req: IncomingMessage, res: ServerResponse) => {
+      if (req.url === '/') {
+        healthCheck(res);
+      } else if (req.url === '/worker') {
+        res.writeHead(200, { 'Content-Type': 'application/json' });
+        res.end(JSON.stringify(workerListener()));
+      } else {
+        res.writeHead(404);
+        res.end('not found');
+      }
+    });
+  }
+
+  async run(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      this.app.listen(this.port, this.host, (err?: Error) => {
+        if (err) reject(err);
+        const address = this.app.address();
+        if (typeof address! !== 'string') {
+          this.#logger.info(`Server is listening on port ${address!.port}`);
+        }
+        resolve();
+      });
+    });
+  }
+
+  async close(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      this.app.close((err?: Error) => {
+        if (err) reject(err);
+        resolve();
+      });
+    });
+  }
+}

package/src/index.ts

@@ -0,0 +1,41 @@
+// SPDX-FileCopyrightText: 2024 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * LiveKit Agents is a framework for building realtime programmable participants that run on
+ * servers.
+ *
+ * @see {@link https://docs.livekit.io/agents/overview | LiveKit Agents documentation}
+ * @packageDocumentation
+ */
+import * as cli from './cli.js';
+import * as inference from './inference/index.js';
+import * as ipc from './ipc/index.js';
+import * as llm from './llm/index.js';
+import * as metrics from './metrics/index.js';
+import * as stream from './stream/index.js';
+import * as stt from './stt/index.js';
+import * as telemetry from './telemetry/index.js';
+import * as tokenize from './tokenize/index.js';
+import * as tts from './tts/index.js';
+import * as voice from './voice/index.js';
+
+export * from './_exceptions.js';
+export * from './audio.js';
+export * from './connection_pool.js';
+export * from './generator.js';
+export * from './inference_runner.js';
+export * from './job.js';
+export * from './log.js';
+export * from './plugin.js';
+export * from './transcription.js';
+export * from './types.js';
+export * from './utils.js';
+export * from './vad.js';
+export * from './version.js';
+export * from './worker.js';
+
+export * from './inference/interruption/index.js';
+
+export { cli, inference, ipc, llm, metrics, stream, stt, telemetry, tokenize, tts, voice };

package/src/inference/api_protos.ts

@@ -0,0 +1,82 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import { z } from 'zod';
+
+export const ttsSessionCreateEventSchema = z.object({
+  type: z.literal('session.create'),
+  sample_rate: z.string(),
+  encoding: z.string(),
+  model: z.string().optional(),
+  voice: z.string().optional(),
+  language: z.string().optional(),
+  extra: z.record(z.string(), z.unknown()),
+  transcript: z.string().optional(),
+});
+
+export const ttsInputTranscriptEventSchema = z.object({
+  type: z.literal('input_transcript'),
+  transcript: z.string(),
+});
+
+export const ttsSessionFlushEventSchema = z.object({
+  type: z.literal('session.flush'),
+});
+
+export const ttsSessionCloseEventSchema = z.object({
+  type: z.literal('session.close'),
+});
+
+export const ttsSessionCreatedEventSchema = z.object({
+  type: z.literal('session.created'),
+  session_id: z.string(),
+});
+
+export const ttsOutputAudioEventSchema = z.object({
+  type: z.literal('output_audio'),
+  audio: z.string(),
+  session_id: z.string(),
+});
+
+export const ttsDoneEventSchema = z.object({
+  type: z.literal('done'),
+  session_id: z.string(),
+});
+
+export const ttsSessionClosedEventSchema = z.object({
+  type: z.literal('session.closed'),
+  session_id: z.string(),
+});
+
+export const ttsErrorEventSchema = z.object({
+  type: z.literal('error'),
+  message: z.string().optional(),
+  session_id: z.string().optional(),
+});
+
+export const ttsClientEventSchema = z.discriminatedUnion('type', [
+  ttsSessionCreateEventSchema,
+  ttsInputTranscriptEventSchema,
+  ttsSessionFlushEventSchema,
+  ttsSessionCloseEventSchema,
+]);
+
+export const ttsServerEventSchema = z.discriminatedUnion('type', [
+  ttsSessionCreatedEventSchema,
+  ttsOutputAudioEventSchema,
+  ttsDoneEventSchema,
+  ttsSessionClosedEventSchema,
+  ttsErrorEventSchema,
+]);
+
+export type TtsSessionCreateEvent = z.infer<typeof ttsSessionCreateEventSchema>;
+export type TtsInputTranscriptEvent = z.infer<typeof ttsInputTranscriptEventSchema>;
+export type TtsSessionFlushEvent = z.infer<typeof ttsSessionFlushEventSchema>;
+export type TtsSessionCloseEvent = z.infer<typeof ttsSessionCloseEventSchema>;
+export type TtsSessionCreatedEvent = z.infer<typeof ttsSessionCreatedEventSchema>;
+export type TtsOutputAudioEvent = z.infer<typeof ttsOutputAudioEventSchema>;
+export type TtsDoneEvent = z.infer<typeof ttsDoneEventSchema>;
+export type TtsSessionClosedEvent = z.infer<typeof ttsSessionClosedEventSchema>;
+export type TtsErrorEvent = z.infer<typeof ttsErrorEventSchema>;
+export type TtsClientEvent = z.infer<typeof ttsClientEventSchema>;
+export type TtsServerEvent = z.infer<typeof ttsServerEventSchema>;

package/src/inference/index.ts

@@ -0,0 +1,32 @@
+// SPDX-FileCopyrightText: 2025 LiveKit, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+import * as llm from './llm.js';
+import * as stt from './stt.js';
+import * as tts from './tts.js';
+
+export {
+  LLM,
+  LLMStream,
+  type ChatCompletionOptions,
+  type GatewayOptions,
+  type InferenceLLMOptions,
+  type LLMModels,
+} from './llm.js';
+
+export {
+  STT,
+  type STTLanguages,
+  type STTModels,
+  type ModelWithLanguage as STTModelString,
+  type STTOptions,
+} from './stt.js';
+
+export {
+  TTS,
+  type TTSModels,
+  type ModelWithVoice as TTSModelString,
+  type TTSOptions,
+} from './tts.js';
+
+export { llm, stt, tts };