chore: refactor progressCB to one place, fix sydney, and sanitize html

This commit is contained in:
Danny Avila 2023-03-14 15:42:59 -04:00
parent 9a17e94f8f
commit 2e20b28c4d
12 changed files with 351 additions and 69 deletions

View file

@ -1,3 +1,7 @@
const { citeText } = require('../../app/');
const _ = require('lodash');
const sanitizeHtml = require('sanitize-html');
const handleError = (res, message) => {
res.write(`event: error\ndata: ${JSON.stringify(message)}\n\n`);
res.end();
@ -10,4 +14,32 @@ const sendMessage = (res, message) => {
res.write(`event: message\ndata: ${JSON.stringify(message)}\n\n`);
};
module.exports = { handleError, sendMessage };
const createOnProgress = () => {
let i = 0;
let tokens = '';
const progressCallback = async (partial, { res, text, bing = false, ...rest }) => {
tokens += partial === text ? '' : partial;
tokens = tokens.trim();
tokens = tokens.replaceAll('[DONE]', '');
if (tokens.includes('```')) {
tokens = sanitizeHtml(tokens);
}
if (bing) {
tokens = citeText(tokens, true);
}
sendMessage(res, { text: tokens, message: true, initial: i === 0, ...rest });
i++;
};
const onProgress = (model, opts) => {
const bingModels = new Set(['bingai', 'sydney']);
return _.partialRight(progressCallback, { ...opts, bing: bingModels.has(model) });
};
return onProgress;
};
module.exports = { handleError, sendMessage, createOnProgress };