mongodb-rag-core
Version:
Common elements used by MongoDB Chatbot Framework components.
98 lines • 3.33 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.makeDataStreamer = exports.escapeNewlines = void 0;
const logger_1 = require("./logger");
function escapeNewlines(str) {
return str.replaceAll(`\n`, `\\n`);
}
exports.escapeNewlines = escapeNewlines;
function makeServerSentEventDispatcher(res) {
return {
connect() {
// Define SSE headers and respond to the client to establish a connection
res.writeHead(200, {
"Cache-Control": "no-cache",
"Content-Type": "text/event-stream",
"Access-Control-Allow-Credentials": "true",
Connection: "keep-alive",
});
},
disconnect() {
res.end();
},
sendData(data) {
res.write(`data: ${JSON.stringify(data)}\n\n`);
},
sendEvent(eventType, data) {
res.write(`event: ${eventType}\n`);
res.write(`data: ${JSON.stringify(data)}\n\n`);
},
};
}
/**
Create a {@link DataStreamer} that streams data to the client.
*/
function makeDataStreamer() {
let connected = false;
let sse;
return {
get connected() {
return connected;
},
connect(res) {
if (this.connected) {
throw new Error("Tried to connect SSE, but it was already connected.");
}
sse = makeServerSentEventDispatcher(res);
res.on("close", () => {
logger_1.logger.info("SSE connection was closed.");
});
sse.connect();
connected = true;
},
disconnect() {
if (!this.connected) {
throw new Error("Tried to disconnect SSE, but it was already disconnected.");
}
sse?.disconnect();
sse = undefined;
connected = false;
},
/**
Streams single item of data in an event stream.
*/
streamData(data) {
if (!this.connected) {
throw new Error(`Tried to stream data, but there's no SSE connection. Call DataStreamer.connect() first.`);
}
sse?.sendData(data);
},
/**
Streams all message events in an event stream.
*/
async stream({ stream }) {
if (!this.connected) {
throw new Error(`Tried to stream data, but there's no SSE connection. Call DataStreamer.connect() first.`);
}
let streamedData = "";
for await (const event of stream) {
if (event.choices.length === 0) {
continue;
}
// The event could contain many choices, but we only want the first one
const choice = event.choices[0];
if (choice.delta) {
const content = escapeNewlines(choice.delta.content ?? "");
this.streamData({
type: "delta",
data: content,
});
streamedData += content;
}
}
return streamedData;
},
};
}
exports.makeDataStreamer = makeDataStreamer;
//# sourceMappingURL=DataStreamer.js.map