@llamaindex/ui
Version:
A comprehensive UI component library built with React, TypeScript, and Tailwind CSS for LlamaIndex applications
219 lines (215 loc) • 6.99 kB
JavaScript
;
var chunkCODV6TUI_js = require('./chunk-CODV6TUI.js');
var chunk4E3IDRQJ_js = require('./chunk-4E3IDRQJ.js');
// src/lib/shared-streaming.ts
var SharedStreamingManager = class {
constructor() {
chunk4E3IDRQJ_js.__publicField(this, "activeStreams", /* @__PURE__ */ new Map());
}
/**
* Subscribe to a shared stream identified by key.
* If stream already exists, reuses it and sends historical events.
* If stream doesn't exist, creates a new one using the executor.
*
* @param streamKey - Unique identifier for the stream
* @param subscriber - Event handlers for the stream
* @param executor - Function that performs the actual streaming
* @param externalSignal - Optional abort signal from caller
* @returns Promise that resolves with all events and unsubscribe function
*/
subscribe(streamKey, subscriber, executor, canceler) {
const existingStream = this.activeStreams.get(streamKey);
if (existingStream) {
return this.subscribeToExistingStream(
streamKey,
existingStream,
subscriber
);
}
return this.createNewStream(streamKey, subscriber, executor, canceler);
}
/**
* Get current events for a stream without subscribing
*/
getStreamEvents(streamKey) {
const stream = this.activeStreams.get(streamKey);
return stream ? [...stream.events] : [];
}
/**
* Check if a stream is currently active
*/
isStreamActive(streamKey) {
return this.activeStreams.has(streamKey);
}
/**
* Get number of subscribers for a stream
*/
getSubscriberCount(streamKey) {
const stream = this.activeStreams.get(streamKey);
return stream ? stream.subscribers.size : 0;
}
/**
* Force close a stream and all its subscribers
*/
closeStream(streamKey) {
const stream = this.activeStreams.get(streamKey);
if (stream) {
stream.controller.abort();
this.cleanupStream(streamKey);
}
}
/**
* Close all active streams
*/
closeAllStreams() {
for (const streamKey of this.activeStreams.keys()) {
this.closeStream(streamKey);
}
}
subscribeToExistingStream(streamKey, stream, subscriber) {
var _a, _b, _c, _d, _e;
stream.subscribers.add(subscriber);
try {
(_a = subscriber.onStart) == null ? void 0 : _a.call(subscriber);
for (const event of stream.events) {
(_b = subscriber.onData) == null ? void 0 : _b.call(subscriber, event);
}
if (stream.isCompleted) {
if (stream.error) {
(_c = subscriber.onError) == null ? void 0 : _c.call(subscriber, stream.error);
} else {
(_d = subscriber.onSuccess) == null ? void 0 : _d.call(subscriber, stream.events);
}
(_e = subscriber.onComplete) == null ? void 0 : _e.call(subscriber);
}
} catch (error) {
console.error("Error sending historical events to subscriber:", error);
}
return {
promise: stream.promise,
unsubscribe: () => this.unsubscribe(streamKey, subscriber),
disconnect: () => this.disconnect(streamKey),
cancel: () => this.cancel(streamKey)
};
}
createNewStream(streamKey, subscriber, executor, canceler) {
const controller = new AbortController();
const subscribers = /* @__PURE__ */ new Set([subscriber]);
const events = [];
const streamState = {
controller,
promise: Promise.resolve([]),
// Will be replaced below
subscribers,
events,
isCompleted: false,
error: null,
canceler
};
this.activeStreams.set(streamKey, streamState);
const compositeSubscriber = {
onStart: () => {
streamState.subscribers.forEach((sub) => {
var _a;
try {
(_a = sub.onStart) == null ? void 0 : _a.call(sub);
} catch (error) {
console.error("Error in subscriber onStart:", error);
}
});
},
onData: (event) => {
events.push(event);
streamState.subscribers.forEach((sub) => {
var _a;
try {
(_a = sub.onData) == null ? void 0 : _a.call(sub, event);
} catch (error) {
console.error("Error in subscriber onData:", error);
}
});
},
onError: (error) => {
streamState.error = error;
streamState.isCompleted = true;
streamState.subscribers.forEach((sub) => {
var _a, _b;
try {
(_a = sub.onError) == null ? void 0 : _a.call(sub, error);
(_b = sub.onComplete) == null ? void 0 : _b.call(sub);
} catch (err) {
console.error("Error in subscriber onError:", err);
}
});
this.cleanupStream(streamKey);
},
onSuccess: (allEvents) => {
streamState.isCompleted = true;
streamState.subscribers.forEach((sub) => {
var _a, _b;
try {
(_a = sub.onSuccess) == null ? void 0 : _a.call(sub, allEvents);
(_b = sub.onComplete) == null ? void 0 : _b.call(sub);
} catch (error) {
console.error("Error in subscriber onFinish:", error);
}
});
this.cleanupStream(streamKey);
}
};
const streamPromise = this.executeStream(
executor,
compositeSubscriber,
controller.signal
);
streamState.promise = streamPromise;
return {
promise: streamPromise,
unsubscribe: () => this.unsubscribe(streamKey, subscriber),
cancel: () => this.cancel(streamKey),
disconnect: () => this.disconnect(streamKey)
};
}
async executeStream(executor, subscriber, signal) {
var _a;
try {
return await executor(subscriber, signal);
} catch (error) {
const err = error instanceof Error ? error : new Error(String(error));
(_a = subscriber.onError) == null ? void 0 : _a.call(subscriber, err);
throw err;
}
}
async cancel(streamKey) {
const stream = this.activeStreams.get(streamKey);
this.disconnect(streamKey);
await (stream == null ? void 0 : stream.canceler());
}
disconnect(streamKey) {
const stream = this.activeStreams.get(streamKey);
if (stream) {
for (const subscriber of stream.subscribers) {
this.unsubscribe(streamKey, subscriber);
}
this.cleanupStream(streamKey);
}
}
unsubscribe(streamKey, subscriber) {
const stream = this.activeStreams.get(streamKey);
if (!stream) return;
stream.subscribers.delete(subscriber);
if (stream.subscribers.size === 0) {
try {
stream.controller.abort();
} catch (error) {
chunkCODV6TUI_js.logger.debug("Error aborting stream in unsubscribe", error);
}
this.cleanupStream(streamKey);
}
}
cleanupStream(streamKey) {
this.activeStreams.delete(streamKey);
}
};
var workflowStreamingManager = new SharedStreamingManager();
exports.workflowStreamingManager = workflowStreamingManager;