@genkit-ai/firebase/lib/stream-manager/firestore.mjs

Genkit AI framework plugin for Firebase including Firestore trace/state store and deployment helpers for Cloud Functions for Firebase.

import { randomUUID } from "crypto";
import {
  FieldValue,
  getFirestore
} from "firebase-admin/firestore";
import {
  GenkitError,
  StreamNotFoundError
} from "genkit/beta";
class FirestoreActionStream {
  streamDoc;
  constructor(streamDoc) {
    this.streamDoc = streamDoc;
  }
  async update(data) {
    await this.streamDoc.update({
      ...data,
      updatedAt: FieldValue.serverTimestamp()
    });
  }
  async write(chunk) {
    await this.update({
      // We add a random ID to the chunk to prevent Firestore from deduplicating chunks
      // that have the same content.
      stream: FieldValue.arrayUnion({
        type: "chunk",
        chunk,
        uuid: randomUUID()
      })
    });
  }
  async done(output) {
    await this.update({
      stream: FieldValue.arrayUnion({ type: "done", output })
    });
  }
  async error(err) {
    const serializableError = {
      message: err.message,
      stack: err.stack,
      ...err
    };
    await this.update({
      stream: FieldValue.arrayUnion({ type: "error", err: serializableError })
    });
  }
}
class FirestoreStreamManager {
  db;
  collection;
  timeout;
  constructor(opts) {
    this.collection = opts.collection;
    this.db = opts.db ?? (opts.firebaseApp ? getFirestore(opts.firebaseApp) : getFirestore());
    this.timeout = opts.timeout ?? 6e4;
  }
  async open(streamId) {
    const streamDoc = this.db.collection(this.collection).doc(streamId);
    await streamDoc.set({
      createdAt: FieldValue.serverTimestamp(),
      stream: []
    });
    return new FirestoreActionStream(streamDoc);
  }
  async subscribe(streamId, callbacks) {
    const streamDoc = this.db.collection(this.collection).doc(streamId);
    const snapshot = await streamDoc.get();
    if (!snapshot.exists) {
      throw new StreamNotFoundError(`Stream ${streamId} not found.`);
    }
    let lastIndex = -1;
    let timeoutId;
    const resetTimeout = () => {
      clearTimeout(timeoutId);
      timeoutId = setTimeout(() => {
        callbacks.onError?.(
          new GenkitError({
            status: "DEADLINE_EXCEEDED",
            message: "Stream timed out."
          })
        );
        unsubscribe();
      }, this.timeout);
    };
    const unsubscribe = streamDoc.onSnapshot((snapshot2) => {
      resetTimeout();
      const data = snapshot2.data();
      if (!data) {
        return;
      }
      const stream = data.stream || [];
      for (let i = lastIndex + 1; i < stream.length; i++) {
        const event = stream[i];
        if (event.type === "chunk") {
          callbacks.onChunk?.(event.chunk);
        } else if (event.type === "done") {
          clearTimeout(timeoutId);
          callbacks.onDone?.(event.output);
          unsubscribe();
        } else if (event.type === "error") {
          clearTimeout(timeoutId);
          callbacks.onError?.(event.err);
          unsubscribe();
        }
      }
      lastIndex = stream.length - 1;
    });
    resetTimeout();
    return { unsubscribe };
  }
}
export {
  FirestoreStreamManager
};
//# sourceMappingURL=firestore.mjs.map