@genkit-ai/core/lib/streaming.js

Genkit AI framework core libraries.

"use strict";
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
  for (var name in all)
    __defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
  if (from && typeof from === "object" || typeof from === "function") {
    for (let key of __getOwnPropNames(from))
      if (!__hasOwnProp.call(to, key) && key !== except)
        __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
  }
  return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var streaming_exports = {};
__export(streaming_exports, {
  InMemoryStreamManager: () => InMemoryStreamManager,
  StreamNotFoundError: () => StreamNotFoundError
});
module.exports = __toCommonJS(streaming_exports);
var import_error = require("./error");
class StreamNotFoundError extends import_error.GenkitError {
  constructor(message) {
    super({ status: "NOT_FOUND", message });
    this.name = "StreamNotFoundError";
  }
}
class InMemoryStreamManager {
  /**
   * @param options Configuration options.
   * @param options.ttlSeconds Time-to-live for streams in seconds. Defaults to 5 minutes.
   */
  constructor(options = {}) {
    this.options = options;
  }
  streams = /* @__PURE__ */ new Map();
  _cleanup() {
    const ttl = (this.options.ttlSeconds ?? 5 * 60) * 1e3;
    const now = Date.now();
    for (const [streamId, stream] of this.streams.entries()) {
      if (stream.status !== "open" && now - stream.lastTouched > ttl) {
        this.streams.delete(streamId);
      }
    }
  }
  async open(streamId) {
    this._cleanup();
    if (this.streams.has(streamId)) {
      throw new Error(`Stream with id ${streamId} already exists.`);
    }
    this.streams.set(streamId, {
      status: "open",
      chunks: [],
      subscribers: [],
      lastTouched: Date.now()
    });
    return {
      write: async (chunk) => {
        const stream = this.streams.get(streamId);
        if (stream?.status === "open") {
          stream.chunks.push(chunk);
          stream.subscribers.forEach((s) => s.onChunk(chunk));
          stream.lastTouched = Date.now();
        }
      },
      done: async (output) => {
        const stream = this.streams.get(streamId);
        if (stream?.status === "open") {
          this.streams.set(streamId, {
            status: "done",
            chunks: stream.chunks,
            output,
            lastTouched: Date.now()
          });
          stream.subscribers.forEach((s) => s.onDone(output));
        }
      },
      error: async (err) => {
        const stream = this.streams.get(streamId);
        if (stream?.status === "open") {
          stream.subscribers.forEach((s) => s.onError(err));
          this.streams.set(streamId, {
            status: "error",
            chunks: stream.chunks,
            error: err,
            lastTouched: Date.now()
          });
        }
      }
    };
  }
  async subscribe(streamId, subscriber) {
    const stream = this.streams.get(streamId);
    if (!stream) {
      throw new StreamNotFoundError(`Stream with id ${streamId} not found.`);
    }
    if (stream.status === "done") {
      for (const chunk of stream.chunks) {
        subscriber.onChunk(chunk);
      }
      subscriber.onDone(stream.output);
    } else if (stream.status === "error") {
      for (const chunk of stream.chunks) {
        subscriber.onChunk(chunk);
      }
      subscriber.onError(stream.error);
    } else {
      stream.chunks.forEach((chunk) => subscriber.onChunk(chunk));
      stream.subscribers.push(subscriber);
    }
    return {
      unsubscribe: () => {
        const currentStream = this.streams.get(streamId);
        if (currentStream?.status === "open") {
          const index = currentStream.subscribers.indexOf(subscriber);
          if (index > -1) {
            currentStream.subscribers.splice(index, 1);
          }
        }
      }
    };
  }
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
  InMemoryStreamManager,
  StreamNotFoundError
});
//# sourceMappingURL=streaming.js.map