Merge pull request #3742 from bmen25124/custom_request_stream Added stream support to "custom-request"

50334890a64b14949690e40ff215c56b5bd2e7da

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed

Parents: 4f52c369faa80564df154c964593716c5d09ef84 0b937237c3e01100a8a325a82a4c450eae891646

public/scripts/custom-request.js+148 -44 Viewed

 import { getTextGenServer } from './textgen-settings.js';
 import { extractReasoningFromData } from './reasoning.js';
 import { formatInstructModeChat, formatInstructModePrompt, names_behavior_types } from './instruct-mode.js';
 +import { getStreamingReply, tryParseStreamingError } from './openai.js';
 +import EventSourceStream from './sse-stream.js';
 // #region Type Definitions
 /**
  * @typedef {Object} TextCompletionRequestBase
 + * @property {boolean?} [stream=false] - Whether to stream the response
  * @property {number} max_tokens - Maximum number of tokens to generate
  * @property {string} [model] - Optional model name
  * @property {string} api_type - Type of API to use

 /**
  * @typedef {Object} TextCompletionPayloadBase
 + * @property {boolean?} [stream=false] - Whether to stream the response
  * @property {string} prompt - The text prompt for completion
  * @property {number} max_tokens - Maximum number of tokens to generate
  * @property {number} max_new_tokens - Alias for max_tokens

 /**
  * @typedef {Object} ChatCompletionPayloadBase
 + * @property {boolean?} [stream=false] - Whether to stream the response
  * @property {ChatCompletionMessage[]} messages - Array of chat messages
  * @property {string} [model] - Optional model name to use for completion
  * @property {string} chat_completion_source - Source provider for chat completion

  * @property {string} reasoning - Extracted reasoning.
  */
 +/**
 + * @typedef {Object} StreamResponse
 + * @property {string} text - Generated text.
 + * @property {string[]} swipes - Generated swipes
 + * @property {Object} state - Generated state
 + * @property {string?} [state.reasoning] - Generated reasoning
 + * @property {string?} [state.image] - Generated image
 + */
++
 // #endregion
 /**
  * Creates & sends a text completion request. Streaming is not supported.
  */
 export class TextCompletionService {
     static TYPE = 'textgenerationwebui';

      * @param {Record<string, any> & TextCompletionRequestBase & {prompt: string}} custom
      * @returns {TextCompletionPayload}
      */
     static createRequestData({ stream = false, prompt, max_tokens, model, api_type, api_server, temperature, min_p, ...props }) {
         const payload = {
             ...props,
 +            stream,
             prompt,
             max_tokens,
             max_new_tokens: max_tokens,

             api_server: api_server ?? getTextGenServer(api_type),
             temperature,
             min_p,
 -            stream: false,
         };
         // Remove undefined values to avoid API errors

      * Sends a text completion request to the specified server
      * @param {TextCompletionPayload} data Request data
      * @param {boolean?} extractData Extract message from the response. Default true
 -     * @returns {Promise<ExtractedData | any>} Extracted data or the raw response
 +     * @param {AbortSignal?} signal
 +     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
      * @throws {Error}
      */
     static async sendRequest(data, extractData = true, signal = null) {
 -        const response = await fetch(getGenerateUrl(this.TYPE), {
 +        if (!data.stream) {
 +            const response = await fetch(getGenerateUrl(this.TYPE), {
 +                method: 'POST',
 +                headers: getRequestHeaders(),
 +                cache: 'no-cache',
 +                body: JSON.stringify(data),
 +                signal: signal ?? new AbortController().signal,
 +            });
++
 +            const json = await response.json();
 +            if (!response.ok || json.error) {
 +                throw json;
 +            }
++
 +            if (!extractData) {
 +                return json;
 +            }
++
 +            return {
 +                content: extractMessageFromData(json, this.TYPE),
 +                reasoning: extractReasoningFromData(json, {
 +                    mainApi: this.TYPE,
 +                    textGenType: data.api_type,
 +                    ignoreShowThoughts: true,
 +                }),
 +            };
 +        }
++
 +        const response = await fetch('/api/backends/text-completions/generate', {
             method: 'POST',
             headers: getRequestHeaders(),
             cache: 'no-cache',
             body: JSON.stringify(data),
             signal: signal ?? new AbortController().signal,
         });
 -        const json = await response.json();
 +        if (!response.ok) {
 -        if (!response.ok || json.error) {
 +            const text = await response.text();
 -            throw json;
 +            tryParseStreamingError(response, text, { quiet: true });
 -        }
 -        if (!extractData) {
 +            throw new Error(`Got response status ${response.status}`);
 -            return json;
+        }
 -        return {
 +        const eventStream = new EventSourceStream();
 -            content: extractMessageFromData(json, this.TYPE),
 +        response.body.pipeThrough(eventStream);
 -            reasoning: extractReasoningFromData(json, {
 +        const reader = eventStream.readable.getReader();
 -                mainApi: this.TYPE,
 +        return async function* streamData() {
 -                textGenType: data.api_type,
 +            let text = '';
 -                ignoreShowThoughts: true,
 +            const swipes = [];
 -            }),
 +            const state = { reasoning: '' };
 +            while (true) {
 +                const { done, value } = await reader.read();
 +                if (done) return;
 +                if (value.data === '[DONE]') return;
++
 +                tryParseStreamingError(response, value.data, { quiet: true });
++
 +                let data = JSON.parse(value.data);
++
 +                if (data?.choices?.[0]?.index > 0) {
 +                    const swipeIndex = data.choices[0].index - 1;
 +                    swipes[swipeIndex] = (swipes[swipeIndex] || '') + data.choices[0].text;
 +                } else {
 +                    const newText = data?.choices?.[0]?.text || data?.content || '';
 +                    text += newText;
 +                    state.reasoning += data?.choices?.[0]?.reasoning ?? '';
 +                }
++
 +                yield { text, swipes, state };
 +            }
         };
+    }

      * @param {string?} [options.presetName] - Name of the preset to use for generation settings
      * @param {string?} [options.instructName] - Name of instruct preset for message formatting
      * @param {boolean} extractData - Whether to extract structured data from response
 -     * @returns {Promise<ExtractedData | any>} Extracted data or the raw response
 +     * @param {AbortSignal?} [signal]
 +     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
      * @throws {Error}
      */
     static async processRequest(
         custom,
         options = {},
         extractData = true,
 +        signal = null,
     ) {
         const { presetName, instructName } = options;
         let requestData = { ...custom };

         // @ts-ignore
         const data = this.createRequestData(requestData);
         return await this.sendRequest(data, extractData, signal);
+    }
     /**

+}
 /**
  * Creates & sends a chat completion request. Streaming is not supported.
  */
 export class ChatCompletionService {
     static TYPE = 'openai';

      * @param {ChatCompletionPayload} custom
      * @returns {ChatCompletionPayload}
      */
     static createRequestData({ stream = false, messages, model, chat_completion_source, max_tokens, temperature, custom_url, ...props }) {
         const payload = {
             ...props,
 +            stream,
             messages,
             model,
             chat_completion_source,
             max_tokens,
             temperature,
             custom_url,
 -            stream: false,
         };
         // Remove undefined values to avoid API errors

      * Sends a chat completion request
      * @param {ChatCompletionPayload} data Request data
      * @param {boolean?} extractData Extract message from the response. Default true
 -     * @returns {Promise<ExtractedData | any>} Extracted data or the raw response
 +     * @param {AbortSignal?} signal Abort signal
 +     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
      * @throws {Error}
      */
     static async sendRequest(data, extractData = true, signal = null) {
         const response = await fetch('/api/backends/chat-completions/generate', {
             method: 'POST',
             headers: getRequestHeaders(),
             cache: 'no-cache',
             body: JSON.stringify(data),
             signal: signal ?? new AbortController().signal,
         });
 -        const json = await response.json();
 +        if (!data.stream) {
 -        if (!response.ok || json.error) {
 +            const json = await response.json();
 -            throw json;
 +            if (!response.ok || json.error) {
 +                throw json;
 +            }
++
 +            if (!extractData) {
 +                return json;
 +            }
++
 +            return {
 +                content: extractMessageFromData(json, this.TYPE),
 +                reasoning: extractReasoningFromData(json, {
 +                    mainApi: this.TYPE,
 +                    textGenType: data.chat_completion_source,
 +                    ignoreShowThoughts: true,
 +                }),
 +            };
+        }
         if (!extractDataresponse.ok) {
 -            return json;
 +            const text = await response.text();
 +            tryParseStreamingError(response, text, { quiet: true });
++
 +            throw new Error(`Got response status ${response.status}`);
+        }
 -        return {
 +        const eventStream = new EventSourceStream();
 -            content: extractMessageFromData(json, this.TYPE),
 +        response.body.pipeThrough(eventStream);
 -            reasoning: extractReasoningFromData(json, {
 +        const reader = eventStream.readable.getReader();
 -                mainApi: this.TYPE,
 +        return async function* streamData() {
 -                textGenType: data.chat_completion_source,
 +            let text = '';
 -                ignoreShowThoughts: true,
 +            const swipes = [];
 -            }),
 +            const state = { reasoning: '', image: '' };
 +            while (true) {
 +                const { done, value } = await reader.read();
 +                if (done) return;
 +                const rawData = value.data;
 +                if (rawData === '[DONE]') return;
 +                tryParseStreamingError(response, rawData, { quiet: true });
 +                const parsed = JSON.parse(rawData);
++
 +                const reply = getStreamingReply(parsed, state, {
 +                    chatCompletionSource: data.chat_completion_source,
 +                    overrideShowThoughts: true,
 +                });
 +                if (Array.isArray(parsed?.choices) && parsed?.choices?.[0]?.index > 0) {
 +                    const swipeIndex = parsed.choices[0].index - 1;
 +                    swipes[swipeIndex] = (swipes[swipeIndex] || '') + reply;
 +                } else {
 +                    text += reply;
 +                }
++
 +                yield { text, swipes: swipes, state };
 +            }
         };
+    }

      * @param {ChatCompletionPayload} custom
      * @param {Object} options - Configuration options
      * @param {string?} [options.presetName] - Name of the preset to use for generation settings
      * @param {boolean} [extractData=true] - Whether to extract structured data from response
 -     * @returns {Promise<ExtractedData | any>} Extracted data or the raw response
 +     * @param {AbortSignal?} [signal] - Abort signal
 +     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
      * @throws {Error}
      */
     static async processRequest(custom, options, extractData = true, signal = null) {
         const { presetName } = options;
         let requestData = { ...custom };

@@ -354,7 +458,7 @@ export class ChatCompletionService {
354	458
355	459	const data = this.createRequestData(requestData);
356	460
357	461	return await this.sendRequest(data, extractData, signal);
358	462	}
359	463
360	464	/**

public/scripts/extensions/shared.js+10 -6 Viewed

+}
 /**
  * It uses the profiles to send a generate request to the API. Doesn't support streaming.
  */
 export class ConnectionManagerRequestService {
     static defaultSendRequestParams = {
 +        stream: false,
 +        signal: null,
         extractData: true,
         includePreset: true,
         includeInstruct: true,

      * @param {string} profileId
      * @param {string | (import('../custom-request.js').ChatCompletionMessage & {ignoreInstruct?: boolean})[]} prompt
      * @param {number} maxTokens
      * @param {{stream?: boolean, signal?: AbortSignal, extractData?: boolean, includePreset?: boolean, includeInstruct?: boolean}} custom - default values are true
      * @returns {Promise<import('../custom-request.js').ExtractedData | any(() => AsyncGenerator<import('../custom-request.js').StreamResponse>)>} ExtractedIf not streaming, returns extracted data; orif thestreaming, rawreturns responsea function that creates an AsyncGenerator
      */
     static async sendRequest(profileId, prompt, maxTokens, custom = this.defaultSendRequestParams) {
         const { stream, signal, extractData, includePreset, includeInstruct } = { ...this.defaultSendRequestParams, ...custom };
         const context = SillyTavern.getContext();
         if (context.extensionSettings.disabledExtensions.includes('connection-manager')) {

                     const messages = Array.isArray(prompt) ? prompt : [{ role: 'user', content: prompt }];
                     return await context.ChatCompletionService.processRequest({
 +                        stream,
                         messages,
                         max_tokens: maxTokens,
                         model: profile.model,

                         custom_url: profile['api-url'],
                     }, {
                         presetName: includePreset ? profile.preset : undefined,
                     }, extractData, signal);
+                }
                 case 'textgenerationwebui': {
                     if (!selectedApiMap.type) {

+                    }
                     return await context.TextCompletionService.processRequest({
 +                        stream,
                         prompt,
                         max_tokens: maxTokens,
                         model: profile.model,

                     }, {
                         instructName: includeInstruct ? profile.instruct : undefined,
                         presetName: includePreset ? profile.preset : undefined,
                     }, extractData, signal);
+                }
                 default: {
                     throw new Error(`Unknown API type ${selectedApiMap.selected}`);

public/scripts/openai.js+36 -21 Viewed

  * Handles errors during streaming requests.
  * @param {Response} response
  * @param {string} decoded - response text or decoded stream data
 + * @param {object} [options]
 + * @param {boolean?} [options.quiet=false] Suppress toast messages
  */
 export function tryParseStreamingError(response, decoded, { quiet = false } = {}) {
     try {
         const data = JSON.parse(decoded);

@@ -1453,19 +1455,19 @@ function tryParseStreamingError(response, decoded) {
1453	1455	return;
1454	1456	}
1455	1457
1456	1458	checkQuotaError(data, { quiet });
1457	1459	checkModerationError(data, { quiet });
1458	1460
1459	1461	// these do not throw correctly (equiv to Error("[object Object]"))
1460	1462	// if trying to fix "[object Object]" displayed to users, start here
1461	1463
1462	1464	if (data.error) {
1463	1465	!quiet && toastr.error(data.error.message \|\| response.statusText, 'Chat Completion API');
1464	1466	throw new Error(data);
1465	1467	}
1466	1468
1467	1469	if (data.message) {
1468	1470	!quiet && toastr.error(data.message, 'Chat Completion API');
1469	1471	throw new Error(data);
1470	1472	}
1471	1473	}

 /**
  * Checks if the response contains a quota error and displays a popup if it does.
  * @param data
 + * @param {object} [options]
 + * @param {boolean?} [options.quiet=false] Suppress toast messages
  * @returns {void}
  * @throws {object} - response JSON
  */
 function checkQuotaError(data, { quiet = false } = {}) {
     if (!data) {
         return;
+    }
     if (data.quota_error) {
         !quiet && renderTemplateAsync('quotaError').then((html) => Popup.show.text('Quota Error', html));
         // this does not throw correctly (equiv to Error("[object Object]"))
         // if trying to fix "[object Object]" displayed to users, start here

+    }
+}
 -function checkModerationError(data) {
 +/**
 + * @param {any} data
 + * @param {object} [options]
 + * @param {boolean?} [options.quiet=false] Suppress toast messages
 + */
 +function checkModerationError(data, { quiet = false } = {}) {
     const moderationError = data?.error?.message?.includes('requires moderation');
     if (moderationError && !quiet) {
         const moderationReason = `Reasons: ${data?.error?.metadata?.reasons?.join(', ') ?? '(N/A)'}`;
         const flaggedText = data?.error?.metadata?.flagged_input ?? '(N/A)';
         toastr.info(flaggedText, moderationReason, { timeOut: 10000 });

  * Extracts the reply from the response data from a chat completions-like source
  * @param {object} data Response data from the chat completions-like source
  * @param {object} state Additional state to keep track of
 + * @param {object} [options] Additional options
 + * @param {string?} [options.chatCompletionSource] Chat completion source
 + * @param {boolean?} [options.overrideShowThoughts] Override show thoughts
  * @returns {string} The reply extracted from the response data
  */
 -function getStreamingReply(data, state) {
 +export function getStreamingReply(data, state, { chatCompletionSource = null, overrideShowThoughts = null } = {}) {
 -    if (oai_settings.chat_completion_source === chat_completion_sources.CLAUDE) {
 +    const chat_completion_source = chatCompletionSource ?? oai_settings.chat_completion_source;
 -        if (oai_settings.show_thoughts) {
 +    const show_thoughts = overrideShowThoughts ?? oai_settings.show_thoughts;
++
 +    if (chat_completion_source === chat_completion_sources.CLAUDE) {
 +        if (show_thoughts) {
             state.reasoning += data?.delta?.thinking || '';
+        }
         return data?.delta?.text || '';
     } else if (oai_settings.chat_completion_source === chat_completion_sources.MAKERSUITE) {
         const inlineData = data?.candidates?.[0]?.content?.parts?.find(x => x.inlineData)?.inlineData;
         if (inlineData) {
             state.image = `data:${inlineData.mimeType};base64,${inlineData.data}`;
+        }
         if (oai_settings.show_thoughts) {
             state.reasoning += (data?.candidates?.[0]?.content?.parts?.filter(x => x.thought)?.map(x => x.text)?.[0] || '');
+        }
         return data?.candidates?.[0]?.content?.parts?.filter(x => !x.thought)?.map(x => x.text)?.[0] || '';
     } else if (oai_settings.chat_completion_source === chat_completion_sources.COHERE) {
         return data?.delta?.message?.content?.text || data?.delta?.message?.tool_plan || '';
     } else if (oai_settings.chat_completion_source === chat_completion_sources.DEEPSEEK) {
         if (oai_settings.show_thoughts) {
             state.reasoning += (data.choices?.filter(x => x?.delta?.reasoning_content)?.[0]?.delta?.reasoning_content || '');
+        }
         return data.choices?.[0]?.delta?.content || '';
     } else if (oai_settings.chat_completion_source === chat_completion_sources.OPENROUTER) {
         if (oai_settings.show_thoughts) {
             state.reasoning += (data.choices?.filter(x => x?.delta?.reasoning)?.[0]?.delta?.reasoning || '');
+        }
         return data.choices?.[0]?.delta?.content ?? data.choices?.[0]?.message?.content ?? data.choices?.[0]?.text ?? '';
     } else if (oai_settings.chat_completion_source === chat_completion_sources.CUSTOM) {
         if (oai_settings.show_thoughts) {
             state.reasoning +=
                 data.choices?.filter(x => x?.delta?.reasoning_content)?.[0]?.delta?.reasoning_content ??
                 data.choices?.filter(x => x?.delta?.reasoning)?.[0]?.delta?.reasoning ??