Switch branch
feat/multi-feature-pack maindefault
View all branches
Blame Raw
Cohee · e3f41666 · · 389 lines (16.4 KB)
1 contributor 
1import { power_user } from './power-user.js';


2import { delay } from './utils.js';


3


4// Symbol for not primary swipe error


5const NOT_PRIMARY = Symbol('not_primary_swipe');


6


7/**


8 * A stream which handles Server-Sent Events from a binary ReadableStream like you get from the fetch API.


9 */


10class EventSourceStream {


11    constructor() {


12        const decoder = new TextDecoderStream('utf-8');


13


14        let streamBuffer = '';


15        let lastEventId = '';


16


17        function processChunk(controller) {


18            // Events are separated by two newlines


19            const events = streamBuffer.split(/\r\n\r\n|\r\r|\n\n/g);


20            if (events.length === 0) return;


21


22            // The leftover text to remain in the buffer is whatever doesn't have two newlines after it. If the buffer ended


23            // with two newlines, this will be an empty string.


24            streamBuffer = events.pop();


25


26            for (const eventChunk of events) {


27                let eventType = '';


28                // Split up by single newlines.


29                const lines = eventChunk.split(/\n|\r|\r\n/g);


30                let eventData = '';


31                for (const line of lines) {


32                    const lineMatch = /([^:]+)(?:: ?(.*))?/.exec(line);


33                    if (lineMatch) {


34                        const field = lineMatch[1];


35                        const value = lineMatch[2] || '';


36


37                        switch (field) {


38                            case 'event':


39                                eventType = value;


40                                break;


41                            case 'data':


42                                eventData += value;


43                                eventData += '\n';


44                                break;


45                            case 'id':


46                                // The ID field cannot contain null, per the spec


47                                if (!value.includes('\0')) lastEventId = value;


48                                break;


49                            // We do nothing for the `delay` type, and other types are explicitly ignored


50                        }


51                    }


52                }


53


54


55                // https://html.spec.whatwg.org/multipage/server-sent-events.html#dispatchMessage


56                // Skip the event if the data buffer is the empty string.


57                if (eventData === '') continue;


58


59                if (eventData[eventData.length - 1] === '\n') {


60                    eventData = eventData.slice(0, -1);


61                }


62


63                // Trim the *last* trailing newline only.


64                const event = new MessageEvent(eventType || 'message', { data: eventData, lastEventId });


65                controller.enqueue(event);


66            }


67        }


68


69        const sseStream = new TransformStream({


70            transform(chunk, controller) {


71                streamBuffer += chunk;


72                processChunk(controller);


73            },


74        });


75


76        decoder.readable.pipeThrough(sseStream);


77


78        this.readable = sseStream.readable;


79        this.writable = decoder.writable;


80    }


81}


82


83/**


84 * Gets a delay based on the character.


85 * @param {string} s The character.


86 * @returns {number} The delay in milliseconds.


87 */


88function getDelay(s) {


89    if (!s) {


90        return 0;


91    }


92


93    const speedFactor = Math.max(100 - power_user.smooth_streaming_speed, 1);


94    const defaultDelayMs = speedFactor * 0.4;


95    const punctuationDelayMs = defaultDelayMs * 25;


96


97    if ([',', '\n'].includes(s)) {


98        return punctuationDelayMs / 2;


99    }


100


101    if (['.', '!', '?'].includes(s)) {


102        return punctuationDelayMs;


103    }


104


105    return defaultDelayMs;


106}


107


108/**


109 * Parses the stream data and returns the parsed data and the chunk to be sent.


110 * @param {object} json The JSON data.


111 * @returns {AsyncGenerator<{data: object, chunk: string, reasoning?: boolean}>} The parsed data and the chunk to be sent.


112 */


113async function* parseStreamData(json) {


114    if (typeof json.delta === 'object' && typeof json.delta.message === 'object' && ['tool-plan-delta', 'content-delta'].includes(json.type)) {


115        // Cohere


116        const text = json?.delta?.message?.content?.text ?? '';


117        for (let i = 0; i < text.length; i++) {


118            const str = json.delta.message.content.text[i];


119            yield {


120                data: { ...json, delta: { message: { content: { text: str } } } },


121                chunk: str,


122            };


123        }


124        return;


125    } else if (typeof json.delta === 'object' && typeof json.delta.text === 'string') {


126        // Claude


127        if (json.delta.text.length > 0) {


128            for (let i = 0; i < json.delta.text.length; i++) {


129                const str = json.delta.text[i];


130                yield {


131                    data: { ...json, delta: { text: str } },


132                    chunk: str,


133                };


134            }


135        }


136        return;


137    } else if (typeof json.delta === 'object' && typeof json.delta.thinking === 'string') {


138        // Claude (reasoning content)


139        if (json.delta.thinking.length > 0) {


140            for (let i = 0; i < json.delta.thinking.length; i++) {


141                const str = json.delta.thinking[i];


142                yield {


143                    data: { ...json, delta: { thinking: str } },


144                    chunk: str,


145                    reasoning: true,


146                };


147            }


148        }


149        return;


150    } else if (Array.isArray(json.candidates)) {


151        // Google VertexAI / AI Studio


152        for (let i = 0; i < json.candidates.length; i++) {


153            const isNotPrimary = json.candidates?.[0]?.index > 0;


154            const hasToolCalls = json?.candidates?.[0]?.content?.parts?.some(p => p?.functionCall);


155            const hasInlineData = json?.candidates?.[0]?.content?.parts?.some(p => p?.inlineData);


156            if (isNotPrimary || json.candidates.length === 0) {


157                return null;


158            }


159            if (hasToolCalls || hasInlineData) {


160                yield { data: json, chunk: '' };


161                return;


162            }


163            if (typeof json.candidates[0].content === 'object' && Array.isArray(json.candidates[i].content.parts)) {


164                for (let j = 0; j < json.candidates[i].content.parts.length; j++) {


165                    if (typeof json.candidates[i].content.parts[j].text === 'string') {


166                        for (let k = 0; k < json.candidates[i].content.parts[j].text.length; k++) {


167                            const moreThanOnePart = json.candidates[i].content.parts.length > 1;


168                            const isNotLastPart = j !== json.candidates[i].content.parts.length - 1;


169                            const isLastSymbol = k === json.candidates[i].content.parts[j].text.length - 1;


170                            const addNewline = moreThanOnePart && isNotLastPart && isLastSymbol;


171                            const str = json.candidates[i].content.parts[j].text[k] + (addNewline ? '\n\n' : '');


172                            const candidateClone = structuredClone(json.candidates[0]);


173                            candidateClone.content.parts[j].text = str;


174                            candidateClone.content.parts = [candidateClone.content.parts[j]];


175                            const candidates = [candidateClone];


176                            const reasoning = json.candidates[i].content.parts[j].thought ?? false;


177                            yield {


178                                data: { ...json, candidates },


179                                chunk: str,


180                                reasoning,


181                            };


182                        }


183                    }


184                }


185            }


186        }


187        return;


188    } else if (typeof json.token === 'string' && json.token.length > 0) {


189        // NovelAI / KoboldCpp Classic


190        for (let i = 0; i < json.token.length; i++) {


191            const str = json.token[i];


192            yield {


193                data: { ...json, token: str },


194                chunk: str,


195            };


196        }


197        return;


198    } else if (typeof json.content === 'string' && json.content.length > 0 && json.object !== 'chat.completion.chunk') {


199        // llama.cpp?


200        const isNotPrimary = json?.index > 0;


201        if (isNotPrimary) {


202            throw new Error('Not a primary swipe', { cause: NOT_PRIMARY });


203        }


204        for (let i = 0; i < json.content.length; i++) {


205            const str = json.content[i];


206            yield {


207                data: { ...json, content: str },


208                chunk: str,


209            };


210        }


211        return;


212    } else if (Array.isArray(json.choices)) {


213        // OpenAI-likes and friends


214        const isNotPrimary = json?.choices?.[0]?.index > 0;


215        if (isNotPrimary || json.choices.length === 0) {


216            throw new Error('Not a primary swipe', { cause: NOT_PRIMARY });


217        }


218


219        if (typeof json.choices[0].text === 'string' && json.choices[0].text.length > 0) {


220            for (let j = 0; j < json.choices[0].text.length; j++) {


221                const str = json.choices[0].text[j];


222                const choiceClone = structuredClone(json.choices[0]);


223                choiceClone.text = str;


224                const choices = [choiceClone];


225                yield {


226                    data: { ...json, choices },


227                    chunk: str,


228                };


229            }


230            return;


231        } else if (typeof json.choices[0].thinking === 'string' && json.choices[0].thinking.length > 0) {


232            for (let j = 0; j < json.choices[0].thinking.length; j++) {


233                const str = json.choices[0].thinking[j];


234                const choiceClone = structuredClone(json.choices[0]);


235                choiceClone.thinking = str;


236                const choices = [choiceClone];


237                yield {


238                    data: { ...json, choices },


239                    chunk: str,


240                    reasoning: true,


241                };


242            }


243            return;


244        } else if (typeof json.choices[0].delta === 'object') {


245            if (typeof json.choices[0].delta.text === 'string' && json.choices[0].delta.text.length > 0) {


246                for (let j = 0; j < json.choices[0].delta.text.length; j++) {


247                    const str = json.choices[0].delta.text[j];


248                    const choiceClone = structuredClone(json.choices[0]);


249                    choiceClone.delta.text = str;


250                    const choices = [choiceClone];


251                    yield {


252                        data: { ...json, choices },


253                        chunk: str,


254                    };


255                }


256                return;


257            } else if (typeof json.choices[0].delta.reasoning_content === 'string' && json.choices[0].delta.reasoning_content.length > 0) {


258                for (let j = 0; j < json.choices[0].delta.reasoning_content.length; j++) {


259                    const str = json.choices[0].delta.reasoning_content[j];


260                    const isLastSymbol = j === json.choices[0].delta.reasoning_content.length - 1;


261                    const choiceClone = structuredClone(json.choices[0]);


262                    choiceClone.delta.reasoning_content = str;


263                    choiceClone.delta.content = isLastSymbol ? choiceClone.delta.content : '';


264                    const choices = [choiceClone];


265                    yield {


266                        data: { ...json, choices },


267                        chunk: str,


268                        reasoning: true,


269                    };


270                }


271                return;


272            } else if (typeof json.choices[0].delta.reasoning === 'string' && json.choices[0].delta.reasoning.length > 0) {


273                for (let j = 0; j < json.choices[0].delta.reasoning.length; j++) {


274                    const str = json.choices[0].delta.reasoning[j];


275                    const isLastSymbol = j === json.choices[0].delta.reasoning.length - 1;


276                    const choiceClone = structuredClone(json.choices[0]);


277                    choiceClone.delta.reasoning = str;


278                    choiceClone.delta.content = isLastSymbol ? choiceClone.delta.content : '';


279                    const choices = [choiceClone];


280                    yield {


281                        data: { ...json, choices },


282                        chunk: str,


283                        reasoning: true,


284                    };


285                }


286                return;


287            } else if (typeof json.choices[0].delta.content === 'string' && json.choices[0].delta.content.length > 0) {


288                for (let j = 0; j < json.choices[0].delta.content.length; j++) {


289                    const str = json.choices[0].delta.content[j];


290                    const choiceClone = structuredClone(json.choices[0]);


291                    choiceClone.delta.content = str;


292                    const choices = [choiceClone];


293                    yield {


294                        data: { ...json, choices },


295                        chunk: str,


296                    };


297                }


298                return;


299            } else if (Array.isArray(json.choices[0].delta.content) && json.choices[0].delta.content.length > 0) {


300                if (Array.isArray(json.choices[0].delta.content[0].thinking) && json.choices[0].delta.content[0].thinking.length > 0) {


301                    if (typeof json.choices[0].delta.content[0].thinking[0].text === 'string' && json.choices[0].delta.content[0].thinking[0].text.length > 0) {


302                        for (let j = 0; j < json.choices[0].delta.content[0].thinking[0].text.length; j++) {


303                            const str = json.choices[0].delta.content[0].thinking[0].text[j];


304                            const choiceClone = structuredClone(json.choices[0]);


305                            choiceClone.delta.content[0].thinking[0].text = str;


306                            const choices = [choiceClone];


307                            yield {


308                                data: { ...json, choices },


309                                chunk: str,


310                                reasoning: true,


311                            };


312                        }


313                        return;


314                    }


315                }


316            }


317        } else if (typeof json.choices[0].message === 'object') {


318            if (typeof json.choices[0].message.content === 'string' && json.choices[0].message.content.length > 0) {


319                for (let j = 0; j < json.choices[0].message.content.length; j++) {


320                    const str = json.choices[0].message.content[j];


321                    const choiceClone = structuredClone(json.choices[0]);


322                    choiceClone.message.content = str;


323                    const choices = [choiceClone];


324                    yield {


325                        data: { ...json, choices },


326                        chunk: str,


327                    };


328                }


329                return;


330            }


331        }


332    }


333


334    throw new Error('Unknown event data format');


335}


336


337/**


338 * Like the default one, but multiplies the events by the number of letters in the event data.


339 */


340export class SmoothEventSourceStream extends EventSourceStream {


341    constructor() {


342        super();


343        let lastStr = '';


344        const transformStream = new TransformStream({


345            async transform(chunk, controller) {


346                const event = chunk;


347                const data = event.data;


348                try {


349                    const hasFocus = document.hasFocus();


350


351                    if (data === '[DONE]') {


352                        lastStr = '';


353                        return controller.enqueue(event);


354                    }


355


356                    const json = JSON.parse(data);


357


358                    if (!json) {


359                        lastStr = '';


360                        return controller.enqueue(event);


361                    }


362


363                    for await (const parsed of parseStreamData(json)) {


364                        !(power_user.smooth_streaming_no_think && parsed.reasoning) && hasFocus && await delay(getDelay(lastStr));


365                        controller.enqueue(new MessageEvent(event.type, { data: JSON.stringify(parsed.data) }));


366                        lastStr = parsed.chunk;


367                    }


368                } catch (error) {


369                    if (error instanceof Error && error.cause !== NOT_PRIMARY) {


370                        console.debug('Smooth Streaming parsing error', error);


371                    }


372                    controller.enqueue(event);


373                }


374            },


375        });


376


377        this.readable = this.readable.pipeThrough(transformStream);


378    }


379}


380


381export function getEventSourceStream() {


382    if (power_user.smooth_streaming) {


383        return new SmoothEventSourceStream();


384    }


385


386    return new EventSourceStream();


387}


388


389export default EventSourceStream;