import { eventSource, event_types } from '../script.js'; import { power_user } from './power-user.js'; import { delay } from './utils.js'; /** * A stream which handles Server-Sent Events from a binary ReadableStream like you get from the fetch API. */ class EventSourceStream { constructor() { const decoder = new TextDecoderStream('utf-8'); let streamBuffer = ''; let lastEventId = ''; function processChunk(controller) { // Events are separated by two newlines const events = streamBuffer.split(/\r\n\r\n|\r\r|\n\n/g); if (events.length === 0) return; // The leftover text to remain in the buffer is whatever doesn't have two newlines after it. If the buffer ended // with two newlines, this will be an empty string. streamBuffer = events.pop(); for (const eventChunk of events) { let eventType = ''; // Split up by single newlines. const lines = eventChunk.split(/\n|\r|\r\n/g); let eventData = ''; for (const line of lines) { const lineMatch = /([^:]+)(?:: ?(.*))?/.exec(line); if (lineMatch) { const field = lineMatch[1]; const value = lineMatch[2] || ''; switch (field) { case 'event': eventType = value; break; case 'data': eventData += value; eventData += '\n'; break; case 'id': // The ID field cannot contain null, per the spec if (!value.includes('\0')) lastEventId = value; break; // We do nothing for the `delay` type, and other types are explicitly ignored } } } // https://html.spec.whatwg.org/multipage/server-sent-events.html#dispatchMessage // Skip the event if the data buffer is the empty string. if (eventData === '') continue; if (eventData[eventData.length - 1] === '\n') { eventData = eventData.slice(0, -1); } // Trim the *last* trailing newline only. const event = new MessageEvent(eventType || 'message', { data: eventData, lastEventId }); controller.enqueue(event); } } const sseStream = new TransformStream({ transform(chunk, controller) { streamBuffer += chunk; processChunk(controller); }, }); decoder.readable.pipeThrough(sseStream); this.readable = sseStream.readable; this.writable = decoder.writable; } } /** * Gets a delay based on the character. * @param {string} s The character. * @returns {number} The delay in milliseconds. */ function getDelay(s) { if (!s) { return 0; } const speedFactor = Math.max(100 - power_user.smooth_streaming_speed, 1); const defaultDelayMs = speedFactor * 0.4; const punctuationDelayMs = defaultDelayMs * 25; if ([',', '\n'].includes(s)) { return punctuationDelayMs / 2; } if (['.', '!', '?'].includes(s)) { return punctuationDelayMs; } return defaultDelayMs; } /** * Parses the stream data and returns the parsed data and the chunk to be sent. * @param {object} json The JSON data. * @returns {AsyncGenerator<{data: object, chunk: string}>} The parsed data and the chunk to be sent. */ async function* parseStreamData(json) { // Claude if (typeof json.delta === 'object') { if (typeof json.delta.text === 'string' && json.delta.text.length > 0) { for (let i = 0; i < json.delta.text.length; i++) { const str = json.delta.text[i]; yield { data: { ...json, delta: { text: str } }, chunk: str, }; } } return; } // MakerSuite else if (Array.isArray(json.candidates)) { for (let i = 0; i < json.candidates.length; i++) { const isNotPrimary = json.candidates?.[0]?.index > 0; if (isNotPrimary || json.candidates.length === 0) { return null; } if (typeof json.candidates[0].content === 'object' && Array.isArray(json.candidates[i].content.parts)) { for (let j = 0; j < json.candidates[i].content.parts.length; j++) { if (typeof json.candidates[i].content.parts[j].text === 'string') { for (let k = 0; k < json.candidates[i].content.parts[j].text.length; k++) { const str = json.candidates[i].content.parts[j].text[k]; const candidateClone = structuredClone(json.candidates[0]); candidateClone.content.parts[j].text = str; const candidates = [candidateClone]; yield { data: { ...json, candidates }, chunk: str, }; } } } } } return; } // NovelAI / KoboldCpp Classic else if (typeof json.token === 'string' && json.token.length > 0) { for (let i = 0; i < json.token.length; i++) { const str = json.token[i]; yield { data: { ...json, token: str }, chunk: str, }; } return; } // llama.cpp? else if (typeof json.content === 'string' && json.content.length > 0) { for (let i = 0; i < json.content.length; i++) { const str = json.content[i]; yield { data: { ...json, content: str }, chunk: str, }; } return; } // OpenAI-likes else if (Array.isArray(json.choices)) { const isNotPrimary = json?.choices?.[0]?.index > 0; if (isNotPrimary || json.choices.length === 0) { throw new Error('Not a primary swipe'); } if (typeof json.choices[0].text === 'string' && json.choices[0].text.length > 0) { for (let j = 0; j < json.choices[0].text.length; j++) { const str = json.choices[0].text[j]; const choiceClone = structuredClone(json.choices[0]); choiceClone.text = str; const choices = [choiceClone]; yield { data: { ...json, choices }, chunk: str, }; } return; } else if (typeof json.choices[0].delta === 'object') { if (typeof json.choices[0].delta.text === 'string' && json.choices[0].delta.text.length > 0) { for (let j = 0; j < json.choices[0].delta.text.length; j++) { const str = json.choices[0].delta.text[j]; const choiceClone = structuredClone(json.choices[0]); choiceClone.delta.text = str; const choices = [choiceClone]; yield { data: { ...json, choices }, chunk: str, }; } return; } else if (typeof json.choices[0].delta.content === 'string' && json.choices[0].delta.content.length > 0) { for (let j = 0; j < json.choices[0].delta.content.length; j++) { const str = json.choices[0].delta.content[j]; const choiceClone = structuredClone(json.choices[0]); choiceClone.delta.content = str; const choices = [choiceClone]; yield { data: { ...json, choices }, chunk: str, }; } return; } } else if (typeof json.choices[0].message === 'object') { if (typeof json.choices[0].message.content === 'string' && json.choices[0].message.content.length > 0) { for (let j = 0; j < json.choices[0].message.content.length; j++) { const str = json.choices[0].message.content[j]; const choiceClone = structuredClone(json.choices[0]); choiceClone.message.content = str; const choices = [choiceClone]; yield { data: { ...json, choices }, chunk: str, }; } return; } } } throw new Error('Unknown event data format'); } /** * Like the default one, but multiplies the events by the number of letters in the event data. */ export class SmoothEventSourceStream extends EventSourceStream { constructor() { super(); let lastStr = ''; const transformStream = new TransformStream({ async transform(chunk, controller) { const event = chunk; const data = event.data; try { const hasFocus = document.hasFocus(); if (data === '[DONE]') { lastStr = ''; return controller.enqueue(event); } const json = JSON.parse(data); if (!json) { lastStr = ''; return controller.enqueue(event); } for await (const parsed of parseStreamData(json)) { hasFocus && await delay(getDelay(lastStr)); controller.enqueue(new MessageEvent(event.type, { data: JSON.stringify(parsed.data) })); lastStr = parsed.chunk; hasFocus && await eventSource.emit(event_types.SMOOTH_STREAM_TOKEN_RECEIVED, parsed.chunk); } } catch (error) { console.debug('Smooth Streaming parsing error', error); controller.enqueue(event); } }, }); this.readable = this.readable.pipeThrough(transformStream); } } export function getEventSourceStream() { if (power_user.smooth_streaming) { return new SmoothEventSourceStream(); } return new EventSourceStream(); } export default EventSourceStream;