@genkit-ai/firebase
Version:
Genkit AI framework plugin for Firebase including Firestore trace/state store and deployment helpers for Cloud Functions for Firebase.
103 lines • 2.96 kB
JavaScript
import { getDatabase } from "firebase-admin/database";
import {
GenkitError,
StreamNotFoundError
} from "genkit/beta";
class RtdbActionStream {
db;
streamRef;
metadataRef;
constructor(db, streamRef) {
this.db = db;
this.streamRef = `${streamRef}/stream`;
this.metadataRef = `${streamRef}/metadata`;
}
async update() {
await this.db.ref(this.metadataRef).update({ updatedAt: Date.now() });
}
async write(chunk) {
await this.db.ref(this.streamRef).push({ type: "chunk", chunk });
await this.update();
}
async done(output) {
await this.db.ref(this.streamRef).push({ type: "done", output });
await this.update();
}
async error(err) {
const serializableError = {
message: err.message,
stack: err.stack,
...err
};
await this.db.ref(this.streamRef).push({ type: "error", err: serializableError });
await this.update();
}
}
class RtdbStreamManager {
db;
refRoot;
timeout;
constructor(opts) {
this.refRoot = opts.refPrefix ?? "genkit-streams";
if (this.refRoot && !this.refRoot.endsWith("/")) {
this.refRoot += "/";
}
this.db = opts.db ?? (opts.firebaseApp ? getDatabase(opts.firebaseApp) : getDatabase());
this.timeout = opts.timeout ?? 6e4;
}
async open(streamId) {
const streamRef = this.refRoot + streamId;
await this.db.ref(streamRef).remove();
await this.db.ref(`${streamRef}/metadata`).set({ createdAt: Date.now() });
return new RtdbActionStream(this.db, streamRef);
}
async subscribe(streamId, callbacks) {
const streamRef = this.db.ref(`${this.refRoot}${streamId}`);
const snapshot = await streamRef.get();
if (!snapshot.exists()) {
throw new StreamNotFoundError(`Stream ${streamId} not found.`);
}
const streamDataRef = streamRef.child("stream");
let timeoutId;
const resetTimeout = () => {
clearTimeout(timeoutId);
timeoutId = setTimeout(() => {
callbacks.onError?.(
new GenkitError({
status: "DEADLINE_EXCEEDED",
message: "Stream timed out."
})
);
unsubscribe();
}, this.timeout);
};
const handleEvent = (snapshot2) => {
resetTimeout();
if (!snapshot2.exists()) {
return;
}
const event = snapshot2.val();
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();
}
};
const unsubscribe = () => {
streamDataRef.off("child_added", handleEvent);
};
streamDataRef.on("child_added", handleEvent);
resetTimeout();
return { unsubscribe };
}
}
export {
RtdbStreamManager
};
//# sourceMappingURL=rtdb.mjs.map