@livekit/agents 1.1.0 → 1.2.0

package/dist/inference/stt.js

@@ -1,8 +1,9 @@
 import {} from "@livekit/rtc-node";
-import { WebSocket } from "ws";
 import { APIError, APIStatusError } from "../_exceptions.js";
 import { AudioByteStream } from "../audio.js";
+import { areLanguagesEquivalent, normalizeLanguage } from "../language.js";
 import { log } from "../log.js";
+import { createStreamChannel } from "../stream/stream_channel.js";
 import {
   STT as BaseSTT,
   SpeechStream as BaseSpeechStream,
@@ -10,16 +11,40 @@ import {
 } from "../stt/index.js";
 import { DEFAULT_API_CONNECT_OPTIONS } from "../types.js";
 import { Event, Task, cancelAndWait, shortuuid, waitForAbort } from "../utils.js";
-import { connectWs, createAccessToken } from "./utils.js";
+import { createTimedString } from "../voice/io.js";
+import {
+  sttServerEventSchema
+} from "./api_protos.js";
+import { connectWs, createAccessToken, getDefaultInferenceUrl } from "./utils.js";
+function parseSTTModelString(model) {
+  const idx = model.lastIndexOf(":");
+  if (idx !== -1) {
+    return [model.slice(0, idx), normalizeLanguage(model.slice(idx + 1))];
+  }
+  return [model, void 0];
+}
+function normalizeSTTFallback(fallback) {
+  const makeFallback = (model) => {
+    if (typeof model === "string") {
+      const [name] = parseSTTModelString(model);
+      return { model: name };
+    }
+    return model;
+  };
+  if (Array.isArray(fallback)) {
+    return fallback.map(makeFallback);
+  }
+  return [makeFallback(fallback)];
+}
 const DEFAULT_ENCODING = "pcm_s16le";
 const DEFAULT_SAMPLE_RATE = 16e3;
-const DEFAULT_BASE_URL = "wss://agent-gateway.livekit.cloud/v1";
 const DEFAULT_CANCEL_TIMEOUT = 5e3;
 class STT extends BaseSTT {
   opts;
   streams = /* @__PURE__ */ new Set();
+  #logger = log();
   constructor(opts) {
-    super({ streaming: true, interimResults: true });
+    super({ streaming: true, interimResults: true, alignedTranscript: "word" });
     const {
       model,
       language,
@@ -28,9 +53,11 @@ class STT extends BaseSTT {
       sampleRate = DEFAULT_SAMPLE_RATE,
       apiKey,
       apiSecret,
-      extraKwargs = {}
+      modelOptions = {},
+      fallback,
+      connOptions
     } = opts || {};
-    const lkBaseURL = baseURL || process.env.LIVEKIT_INFERENCE_URL || DEFAULT_BASE_URL;
+    const lkBaseURL = baseURL || getDefaultInferenceUrl();
     const lkApiKey = apiKey || process.env.LIVEKIT_INFERENCE_API_KEY || process.env.LIVEKIT_API_KEY;
     if (!lkApiKey) {
       throw new Error("apiKey is required: pass apiKey or set LIVEKIT_API_KEY");
@@ -39,39 +66,113 @@ class STT extends BaseSTT {
     if (!lkApiSecret) {
       throw new Error("apiSecret is required: pass apiSecret or set LIVEKIT_API_SECRET");
     }
+    let nextModel = model;
+    let nextLanguage = language;
+    if (typeof nextModel === "string") {
+      const [parsedModel, parsedLanguage] = parseSTTModelString(nextModel);
+      if (parsedLanguage !== void 0) {
+        if (nextLanguage && !areLanguagesEquivalent(nextLanguage, parsedLanguage)) {
+          this.#logger.warn(
+            "`language` is provided via both argument and model, using the one from the argument",
+            { language: nextLanguage, model: nextModel }
+          );
+        } else {
+          nextLanguage = parsedLanguage;
+        }
+        nextModel = parsedModel;
+      }
+    }
+    const normalizedFallback = fallback ? normalizeSTTFallback(fallback) : void 0;
     this.opts = {
-      model,
-      language,
+      model: nextModel,
+      language: nextLanguage ? normalizeLanguage(nextLanguage) : void 0,
       encoding,
       sampleRate,
       baseURL: lkBaseURL,
       apiKey: lkApiKey,
       apiSecret: lkApiSecret,
-      extraKwargs
+      modelOptions,
+      fallback: normalizedFallback,
+      connOptions: connOptions ?? DEFAULT_API_CONNECT_OPTIONS
     };
   }
   get label() {
     return "inference.STT";
   }
+  get model() {
+    return this.opts.model ?? "auto";
+  }
+  get provider() {
+    return "livekit";
+  }
+  static fromModelString(modelString) {
+    const [model, language] = parseSTTModelString(modelString);
+    return new STT({ model, language });
+  }
   async _recognize(_) {
     throw new Error("LiveKit STT does not support batch recognition, use stream() instead");
   }
   updateOptions(opts) {
-    this.opts = { ...this.opts, ...opts };
+    this.opts = {
+      ...this.opts,
+      ...opts,
+      language: opts.language !== void 0 ? normalizeLanguage(opts.language) : this.opts.language
+    };
     for (const stream of this.streams) {
       stream.updateOptions(opts);
     }
   }
   stream(options) {
-    const { language, connOptions = DEFAULT_API_CONNECT_OPTIONS } = options || {};
+    const { language, connOptions = this.opts.connOptions ?? DEFAULT_API_CONNECT_OPTIONS } = options || {};
     const streamOpts = {
       ...this.opts,
-      language: language ?? this.opts.language
+      language: language !== void 0 ? normalizeLanguage(language) : this.opts.language
     };
     const stream = new SpeechStream(this, streamOpts, connOptions);
     this.streams.add(stream);
     return stream;
   }
+  async connectWs(timeout) {
+    var _a;
+    const params = {
+      settings: {
+        sample_rate: String(this.opts.sampleRate),
+        encoding: this.opts.encoding,
+        extra: this.opts.modelOptions
+      }
+    };
+    if (this.opts.model && this.opts.model !== "auto") {
+      params.model = this.opts.model;
+    }
+    if (this.opts.language) {
+      params.settings.language = this.opts.language;
+    }
+    if ((_a = this.opts.fallback) == null ? void 0 : _a.length) {
+      params.fallback = {
+        models: this.opts.fallback.map((m) => ({
+          model: m.model,
+          extra: m.extraKwargs ?? {}
+        }))
+      };
+    }
+    if (this.opts.connOptions) {
+      params.connection = {
+        timeout: this.opts.connOptions.timeoutMs / 1e3,
+        retries: this.opts.connOptions.maxRetry
+      };
+    }
+    let baseURL = this.opts.baseURL;
+    if (baseURL.startsWith("http://") || baseURL.startsWith("https://")) {
+      baseURL = baseURL.replace("http", "ws");
+    }
+    const token = await createAccessToken(this.opts.apiKey, this.opts.apiSecret);
+    const url = `${baseURL}/stt`;
+    const headers = { Authorization: `Bearer ${token}` };
+    const socket = await connectWs(url, headers, timeout);
+    const msg = { ...params, type: "session.create" };
+    socket.send(JSON.stringify(msg));
+    return socket;
+  }
 }
 class SpeechStream extends BaseSpeechStream {
   opts;
@@ -79,206 +180,254 @@ class SpeechStream extends BaseSpeechStream {
   speaking = false;
   speechDuration = 0;
   reconnectEvent = new Event();
+  stt;
+  connOptions;
   #logger = log();
   constructor(sttImpl, opts, connOptions) {
     super(sttImpl, opts.sampleRate, connOptions);
     this.opts = opts;
+    this.stt = sttImpl;
+    this.connOptions = connOptions;
   }
   get label() {
     return "inference.SpeechStream";
   }
   updateOptions(opts) {
-    this.opts = { ...this.opts, ...opts };
+    this.opts = {
+      ...this.opts,
+      ...opts,
+      language: opts.language !== void 0 ? normalizeLanguage(opts.language) : this.opts.language
+    };
+    this.reconnectEvent.set();
   }
   async run() {
-    let ws = null;
-    let closingWs = false;
-    this.reconnectEvent.set();
-    const connect = async () => {
-      const params = {
-        settings: {
-          sample_rate: String(this.opts.sampleRate),
-          encoding: this.opts.encoding,
-          extra: this.opts.extraKwargs
-        }
+    while (true) {
+      let ws = null;
+      let closing = false;
+      let finalReceived = false;
+      const eventChannel = createStreamChannel();
+      const resourceCleanup = () => {
+        if (closing) return;
+        closing = true;
+        eventChannel.close();
+        ws == null ? void 0 : ws.removeAllListeners();
+        ws == null ? void 0 : ws.close();
       };
-      if (this.opts.model) {
-        params.model = this.opts.model;
-      }
-      if (this.opts.language) {
-        params.settings.language = this.opts.language;
-      }
-      let baseURL = this.opts.baseURL;
-      if (baseURL.startsWith("http://") || baseURL.startsWith("https://")) {
-        baseURL = baseURL.replace("http", "ws");
-      }
-      const token = await createAccessToken(this.opts.apiKey, this.opts.apiSecret);
-      const url = `${baseURL}/stt`;
-      const headers = { Authorization: `Bearer ${token}` };
-      const socket = await connectWs(url, headers, 1e4);
-      const msg = { ...params, type: "session.create" };
-      socket.send(JSON.stringify(msg));
-      return socket;
-    };
-    const send = async (socket, signal) => {
-      const audioStream = new AudioByteStream(
-        this.opts.sampleRate,
-        1,
-        Math.floor(this.opts.sampleRate / 20)
-        // 50ms
-      );
-      for await (const ev of this.input) {
-        if (signal.aborted) break;
-        let frames;
-        if (ev === SpeechStream.FLUSH_SENTINEL) {
-          frames = audioStream.flush();
-        } else {
-          const frame = ev;
-          frames = audioStream.write(new Int16Array(frame.data).buffer);
-        }
-        for (const frame of frames) {
-          this.speechDuration += frame.samplesPerChannel / frame.sampleRate;
-          const base64 = Buffer.from(frame.data.buffer).toString("base64");
-          const msg = { type: "input_audio", audio: base64 };
-          socket.send(JSON.stringify(msg));
-        }
-      }
-      closingWs = true;
-      socket.send(JSON.stringify({ type: "session.finalize" }));
-    };
-    const recv = async (socket, signal) => {
-      while (!this.closed && !signal.aborted) {
-        const dataPromise = new Promise((resolve, reject) => {
-          const messageHandler = (d) => {
-            resolve(d.toString());
-            removeListeners();
+      const createWsListener = async (ws2, signal) => {
+        return new Promise((resolve, reject) => {
+          const onAbort = () => {
+            resourceCleanup();
+            reject(new Error("WebSocket connection aborted"));
           };
-          const errorHandler = (e) => {
+          signal.addEventListener("abort", onAbort, { once: true });
+          ws2.on("message", (data) => {
+            const json = JSON.parse(data.toString());
+            eventChannel.write(json);
+          });
+          ws2.on("error", (e) => {
+            this.#logger.error({ error: e }, "WebSocket error");
+            resourceCleanup();
             reject(e);
-            removeListeners();
-          };
-          const closeHandler = (code) => {
-            if (closingWs) {
-              resolve("");
+          });
+          ws2.on("close", (code) => {
+            resourceCleanup();
+            if (!closing) return this.#logger.error("WebSocket closed unexpectedly");
+            if (finalReceived) return resolve();
+            reject(
+              new APIStatusError({
+                message: "LiveKit STT connection closed unexpectedly",
+                options: { statusCode: code }
+              })
+            );
+          });
+        });
+      };
+      const send = async (socket, signal) => {
+        const audioStream = new AudioByteStream(
+          this.opts.sampleRate,
+          1,
+          Math.floor(this.opts.sampleRate / 20)
+          // 50ms
+        );
+        const abortPromise = new Promise((_, reject) => {
+          if (signal.aborted) {
+            return reject(new Error("Send aborted"));
+          }
+          const onAbort = () => reject(new Error("Send aborted"));
+          signal.addEventListener("abort", onAbort, { once: true });
+        });
+        const iterator = this.input[Symbol.asyncIterator]();
+        try {
+          while (true) {
+            const result = await Promise.race([iterator.next(), abortPromise]);
+            if (result.done) break;
+            const ev = result.value;
+            let frames;
+            if (ev === SpeechStream.FLUSH_SENTINEL) {
+              frames = audioStream.flush();
             } else {
-              reject(
-                new APIStatusError({
-                  message: "LiveKit STT connection closed unexpectedly",
-                  options: { statusCode: code }
-                })
+              const frame = ev;
+              frames = audioStream.write(new Int16Array(frame.data).buffer);
+            }
+            for (const frame of frames) {
+              this.speechDuration += frame.samplesPerChannel / frame.sampleRate;
+              const base64 = Buffer.from(frame.data.buffer).toString("base64");
+              const msg = { type: "input_audio", audio: base64 };
+              socket.send(JSON.stringify(msg));
+            }
+          }
+          closing = true;
+          socket.send(JSON.stringify({ type: "session.finalize" }));
+        } catch (e) {
+          if (e.message === "Send aborted") {
+            return;
+          }
+          throw e;
+        }
+      };
+      const recv = async (signal) => {
+        const serverEventStream = eventChannel.stream();
+        const reader = serverEventStream.getReader();
+        try {
+          while (!this.closed && !signal.aborted) {
+            const result = await reader.read();
+            if (signal.aborted) return;
+            if (result.done) return;
+            const parseResult = await sttServerEventSchema.safeParseAsync(result.value);
+            if (!parseResult.success) {
+              this.#logger.warn(
+                { error: parseResult.error, rawData: result.value },
+                "Failed to parse STT server event"
               );
+              continue;
             }
-            removeListeners();
-          };
-          const removeListeners = () => {
-            socket.removeListener("message", messageHandler);
-            socket.removeListener("error", errorHandler);
-            socket.removeListener("close", closeHandler);
-          };
-          socket.once("message", messageHandler);
-          socket.once("error", errorHandler);
-          socket.once("close", closeHandler);
-        });
-        const data = await Promise.race([dataPromise, waitForAbort(signal)]);
-        if (!data || signal.aborted) return;
-        const json = JSON.parse(data);
-        const type = json.type;
-        switch (type) {
-          case "session.created":
-          case "session.finalized":
-          case "session.closed":
-            break;
-          case "interim_transcript":
-            this.processTranscript(json, false);
-            break;
-          case "final_transcript":
-            this.processTranscript(json, true);
-            break;
-          case "error":
-            this.#logger.error("received error from LiveKit STT: %o", json);
-            throw new APIError(`LiveKit STT returned error: ${JSON.stringify(json)}`);
-          default:
-            this.#logger.warn("received unexpected message from LiveKit STT: %o", json);
-            break;
+            const event = parseResult.data;
+            switch (event.type) {
+              case "session.created":
+              case "session.finalized":
+                break;
+              case "session.closed":
+                finalReceived = true;
+                resourceCleanup();
+                break;
+              case "interim_transcript":
+                this.processTranscript(event, false);
+                break;
+              case "final_transcript":
+                this.processTranscript(event, true);
+                break;
+              case "error":
+                this.#logger.error({ error: event }, "Received error from LiveKit STT");
+                resourceCleanup();
+                throw new APIError(`LiveKit STT returned error: ${JSON.stringify(event)}`);
+            }
+          }
+        } finally {
+          reader.releaseLock();
+          try {
+            await serverEventStream.cancel();
+          } catch (e) {
+            this.#logger.debug("Error cancelling serverEventStream (may already be cancelled):", e);
+          }
         }
-      }
-    };
-    while (true) {
+      };
       try {
-        ws = await connect();
-        const sendTask = Task.from(async ({ signal }) => {
-          await send(ws, signal);
-        });
-        const recvTask = Task.from(async ({ signal }) => {
-          await recv(ws, signal);
-        });
-        const tasks = [sendTask, recvTask];
-        const waitReconnectTask = Task.from(async ({ signal }) => {
-          await Promise.race([this.reconnectEvent.wait(), waitForAbort(signal)]);
-        });
+        ws = await this.stt.connectWs(this.connOptions.timeoutMs);
+        const controller = this.abortController;
+        const sendTask = Task.from(({ signal }) => send(ws, signal), controller);
+        const wsListenerTask = Task.from(({ signal }) => createWsListener(ws, signal), controller);
+        const recvTask = Task.from(({ signal }) => recv(signal), controller);
+        const waitReconnectTask = Task.from(
+          ({ signal }) => Promise.race([this.reconnectEvent.wait(), waitForAbort(signal)]),
+          controller
+        );
         try {
           await Promise.race([
-            Promise.all(tasks.map((task) => task.result)),
+            Promise.all([sendTask.result, wsListenerTask.result, recvTask.result]),
             waitReconnectTask.result
           ]);
           if (!waitReconnectTask.done) break;
           this.reconnectEvent.clear();
         } finally {
-          await cancelAndWait([sendTask, recvTask, waitReconnectTask], DEFAULT_CANCEL_TIMEOUT);
+          await cancelAndWait(
+            [sendTask, wsListenerTask, recvTask, waitReconnectTask],
+            DEFAULT_CANCEL_TIMEOUT
+          );
+          resourceCleanup();
         }
       } finally {
-        try {
-          if (ws) ws.close();
-        } catch {
-        }
+        resourceCleanup();
       }
     }
   }
   processTranscript(data, isFinal) {
-    const requestId = data.request_id ?? this.requestId;
-    const text = data.transcript ?? "";
-    const language = data.language ?? this.opts.language ?? "en";
+    if (this.queue.closed) return;
+    const requestId = data.session_id || this.requestId;
+    const text = data.transcript;
+    const language = normalizeLanguage(data.language || this.opts.language || "en");
     if (!text && !isFinal) return;
-    if (!this.speaking) {
-      this.speaking = true;
-      this.queue.put({ type: SpeechEventType.START_OF_SPEECH });
-    }
-    const speechData = {
-      language,
-      startTime: data.start ?? 0,
-      endTime: data.duration ?? 0,
-      confidence: data.confidence ?? 1,
-      text
-    };
-    if (isFinal) {
-      if (this.speechDuration > 0) {
+    try {
+      if (!this.speaking) {
+        this.speaking = true;
+        this.queue.put({ type: SpeechEventType.START_OF_SPEECH });
+      }
+      const speechData = {
+        language,
+        startTime: this.startTimeOffset + data.start,
+        endTime: this.startTimeOffset + data.start + data.duration,
+        confidence: data.confidence,
+        text,
+        words: data.words.map(
+          (word) => createTimedString({
+            text: word.word,
+            startTime: word.start + this.startTimeOffset,
+            endTime: word.end + this.startTimeOffset,
+            startTimeOffset: this.startTimeOffset,
+            confidence: word.confidence
+          })
+        )
+      };
+      if (isFinal) {
+        if (this.speechDuration > 0) {
+          this.queue.put({
+            type: SpeechEventType.RECOGNITION_USAGE,
+            requestId,
+            recognitionUsage: { audioDuration: this.speechDuration }
+          });
+          this.speechDuration = 0;
+        }
+        this.queue.put({
+          type: SpeechEventType.FINAL_TRANSCRIPT,
+          requestId,
+          alternatives: [speechData]
+        });
+        if (this.speaking) {
+          this.speaking = false;
+          this.queue.put({ type: SpeechEventType.END_OF_SPEECH });
+        }
+      } else {
         this.queue.put({
-          type: SpeechEventType.RECOGNITION_USAGE,
+          type: SpeechEventType.INTERIM_TRANSCRIPT,
           requestId,
-          recognitionUsage: { audioDuration: this.speechDuration }
+          alternatives: [speechData]
         });
-        this.speechDuration = 0;
       }
-      this.queue.put({
-        type: SpeechEventType.FINAL_TRANSCRIPT,
-        requestId,
-        alternatives: [speechData]
-      });
-      if (this.speaking) {
-        this.speaking = false;
-        this.queue.put({ type: SpeechEventType.END_OF_SPEECH });
+    } catch (e) {
+      if (e instanceof Error && e.message.includes("Queue is closed")) {
+        this.#logger.warn(
+          { err: e },
+          "Queue closed during transcript processing (expected during disconnect)"
+        );
+      } else {
+        this.#logger.error({ err: e }, "Error putting transcript to queue");
       }
-    } else {
-      this.queue.put({
-        type: SpeechEventType.INTERIM_TRANSCRIPT,
-        requestId,
-        alternatives: [speechData]
-      });
     }
   }
 }
 export {
   STT,
-  SpeechStream
+  SpeechStream,
+  normalizeSTTFallback,
+  parseSTTModelString
 };
 //# sourceMappingURL=stt.js.map