1
import { getPresetManager } from './preset-manager.js';
2
import { extractJsonFromData, extractMessageFromData, getGenerateUrl, getRequestHeaders, name1, name2 } from '../script.js';
3
import { getTextGenServer, createTextGenGenerationData, setting_names, textgenerationwebui_settings } from './textgen-settings.js';
4
import { extractReasoningFromData } from './reasoning.js';
5
import { formatInstructModeChat, formatInstructModePrompt, getInstructStoppingSequences } from './instruct-mode.js';
6
import { getStreamingReply, tryParseStreamingError, createGenerationParameters, settingsToUpdate, oai_settings } from './openai.js';
7
import EventSourceStream from './sse-stream.js';
8

9
// #region Type Definitions
10
/**
11
 * @typedef {Object} TextCompletionRequestBase
12
 * @property {boolean?} [stream=false] - Whether to stream the response
13
 * @property {number} max_tokens - Maximum number of tokens to generate
14
 * @property {string} [model] - Optional model name
15
 * @property {string} api_type - Type of API to use
16
 * @property {string} [api_server] - Optional API server URL
17
 * @property {number} [temperature] - Optional temperature parameter
18
 * @property {number} [min_p] - Optional min_p parameter
19
 */
20

21
/**
22
 * @typedef {Object} TextCompletionPayloadBase
23
 * @property {boolean?} [stream=false] - Whether to stream the response
24
 * @property {string} prompt - The text prompt for completion
25
 * @property {number} max_tokens - Maximum number of tokens to generate
26
 * @property {number} max_new_tokens - Alias for max_tokens
27
 * @property {string} [model] - Optional model name
28
 * @property {string} api_type - Type of API to use
29
 * @property {string} api_server - API server URL
30
 * @property {number} [temperature] - Optional temperature parameter
31
 */
32

33
/** @typedef {Record<string, any> & TextCompletionPayloadBase} TextCompletionPayload */
34

35
/**
36
 * @typedef {Object} ChatCompletionMessage
37
 * @property {string} [name] - The name of the message author (optional)
38
 * @property {string} role - The role of the message author (e.g., "user", "assistant", "system")
39
 * @property {string} content - The content of the message
40
 */
41

42
/**
43
 * @typedef {Object} ChatCompletionPayloadBase
44
 * @property {boolean?} [stream=false] - Whether to stream the response
45
 * @property {ChatCompletionMessage[]} messages - Array of chat messages
46
 * @property {string} [model] - Optional model name to use for completion
47
 * @property {string} chat_completion_source - Source provider
48
 * @property {number} max_tokens - Maximum number of tokens to generate
49
 * @property {number} [temperature] - Optional temperature parameter for response randomness
50
 * @property {string} [custom_url] - Optional custom URL
51
 * @property {string} [reverse_proxy] - Optional reverse proxy URL
52
 * @property {string} [proxy_password] - Optional proxy password
53
 * @property {string} [custom_prompt_post_processing] - Optional custom prompt post-processing
54
 * @property {import('../script.js').JsonSchema} [json_schema] - Optional JSON schema for structured generation
55
 */
56

57
/** @typedef {Record<string, any> & ChatCompletionPayloadBase} ChatCompletionPayload */
58

59
/**
60
 * @typedef {Object} ExtractedData
61
 * @property {string} content - Extracted content.
62
 * @property {string} reasoning - Extracted reasoning.
63
 */
64

65
/**
66
 * @typedef {Object} StreamResponse
67
 * @property {string} text - Generated text.
68
 * @property {string[]} swipes - Generated swipes
69
 * @property {Object} state - Generated state
70
 * @property {string?} [state.reasoning] - Generated reasoning
71
 * @property {string?} [state.image] - Generated image
72
 */
73

74
// #endregion
75

76
/**
77
 * Creates & sends a text completion request.
78
 */
79
export class TextCompletionService {
80
    static TYPE = 'textgenerationwebui';
81

82
    /**
83
     * @param {Record<string, any> & TextCompletionRequestBase & {prompt: string}} custom
84
     * @returns {TextCompletionPayload}
85
     */
86
    static createRequestData({ stream = false, prompt, max_tokens, model, api_type, api_server, temperature, min_p, ...props }) {
87
        const payload = {
88
            stream,
89
            prompt,
90
            max_tokens,
91
            max_new_tokens: max_tokens,
92
            model,
93
            api_type,
94
            api_server: api_server ?? getTextGenServer(api_type),
95
            temperature,
96
            min_p,
97
            ...props,
98
        };
99

100
        // Remove undefined values to avoid API errors
101
        Object.keys(payload).forEach(key => {
102
            if (payload[key] === undefined) {
103
                delete payload[key];
104
            }
105
        });
106

107
        return payload;
108
    }
109

110
    /**
111
     * Sends a text completion request to the specified server
112
     * @param {TextCompletionPayload} data Request data
113
     * @param {boolean?} extractData Extract message from the response. Default true
114
     * @param {AbortSignal?} signal
115
     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
116
     * @throws {Error}
117
     */
118
    static async sendRequest(data, extractData = true, signal = null) {
119
        if (!data.stream) {
120
            const response = await fetch(getGenerateUrl(this.TYPE), {
121
                method: 'POST',
122
                headers: getRequestHeaders(),
123
                cache: 'no-cache',
124
                body: JSON.stringify(data),
125
                signal: signal ?? new AbortController().signal,
126
            });
127

128
            const json = await response.json();
129
            if (!response.ok || json.error) {
130
                throw new Error(String(json.error?.message || 'Response not OK'));
131
            }
132

133
            if (!extractData) {
134
                return json;
135
            }
136

137
            return {
138
                content: extractMessageFromData(json, this.TYPE),
139
                reasoning: extractReasoningFromData(json, {
140
                    mainApi: this.TYPE,
141
                    textGenType: data.api_type,
142
                    ignoreShowThoughts: true,
143
                }),
144
            };
145
        }
146

147
        const response = await fetch('/api/backends/text-completions/generate', {
148
            method: 'POST',
149
            headers: getRequestHeaders(),
150
            cache: 'no-cache',
151
            body: JSON.stringify(data),
152
            signal: signal ?? new AbortController().signal,
153
        });
154

155
        if (!response.ok) {
156
            const text = await response.text();
157
            tryParseStreamingError(response, text, { quiet: true });
158

159
            throw new Error(`Got response status ${response.status}`);
160
        }
161

162
        const eventStream = new EventSourceStream();
163
        response.body.pipeThrough(eventStream);
164
        const reader = eventStream.readable.getReader();
165
        return async function* streamData() {
166
            let text = '';
167
            const swipes = [];
168
            const state = { reasoning: '' };
169
            while (true) {
170
                const { done, value } = await reader.read();
171
                if (done) return;
172
                if (value.data === '[DONE]') return;
173

174
                tryParseStreamingError(response, value.data, { quiet: true });
175

176
                let data = JSON.parse(value.data);
177

178
                if (data?.choices?.[0]?.index > 0) {
179
                    const swipeIndex = data.choices[0].index - 1;
180
                    swipes[swipeIndex] = (swipes[swipeIndex] || '') + data.choices[0].text;
181
                } else {
182
                    const newText = data?.choices?.[0]?.text || data?.content || '';
183
                    text += newText;
184
                    state.reasoning += data?.choices?.[0]?.reasoning ?? '';
185
                }
186

187
                yield { text, swipes, state };
188
            }
189
        };
190
    }
191

192
    /**
193
    * Return a formatted prompt string given an array of messages, a chosen instruct preset, and instruct settings.
194
    * @param {(ChatCompletionMessage & {ignoreInstruct?: boolean})[]} prompt An array of messages
195
    * @param {InstructSettings|string} instructPreset Either the name of an instruct preset or the instruct preset object itself.
196
    * @param {Partial<InstructSettings>} instructSettings Optional instruct settings
197
    */
198
    static constructPrompt(prompt, instructPreset, instructSettings) {
199
        // InstructPreset may either be a name or itself a preset
200
        if (typeof instructPreset === 'string') {
201
            const instructPresetManager = getPresetManager('instruct');
202
            instructPreset = instructPresetManager?.getCompletionPresetByName(instructPreset);
203
        }
204

205
        // Clone the preset to avoid modifying the original
206
        instructPreset = structuredClone(instructPreset);
207
        if (instructSettings) {  // apply any additional settings
208
            Object.assign(instructPreset, instructSettings);
209
        }
210

211
        // Make the type check shut up. We 100% don't have a string here.
212
        if (typeof instructPreset === 'string') {
213
            return;
214
        }
215

216
        // Format messages using instruct formatting
217
        const formattedMessages = [];
218
        const prefillActive = prompt.length > 0 ? prompt[prompt.length - 1].role === 'assistant' : false;
219
        for (const message of prompt) {
220
            let messageContent = message.content;
221
            if (!message.ignoreInstruct) {
222
                const isLastMessage = message === prompt[prompt.length - 1];
223

224
                // This complicated logic means:
225
                // 1. If prefill is not active, format all messages
226
                // 2. If prefill is active, format all messages except the last one
227
                if (!isLastMessage || !prefillActive) {
228
                    messageContent = formatInstructModeChat(
229
                        message.name ?? message.role,
230
                        message.content,
231
                        message.role === 'user',
232
                        message.role === 'system',
233
                        undefined,
234
                        name1,  // for macros
235
                        name2,  // for macros
236
                        undefined,
237
                        instructPreset,
238
                    );
239
                }
240

241
                // Add prompt formatting for the last message.
242
                // e.g. "<|im_start|>assistant"
243
                if (isLastMessage) {
244
                    let last_line = formatInstructModePrompt(
245
                        'assistant',  // for sequences using {{name}}
246
                        false,  // not an impersonation
247
                        prefillActive ? message.content : undefined,  // if using prefill, last message is the prefill
248
                        name1,  // for macros
249
                        name2,  // for macros
250
                        true,   // quiet
251
                        false,
252
                        instructPreset,
253
                    );
254

255
                    if (prefillActive) {  // content is the prefilled message
256
                        if (last_line.endsWith('\n') && !message.content.endsWith('\n')) {
257
                            last_line = last_line.slice(0, -1);  // remove newline after prefill if it's not in the prefill itself
258
                        }
259
                        messageContent = last_line;
260
                    } else {  // append last line to content (e.g. "<|im_start|>assistant:")
261
                        messageContent += last_line;
262
                    }
263
                }
264
            }
265
            formattedMessages.push(messageContent);
266
        }
267
        return formattedMessages.join('');
268
    }
269

270

271
    /**
272
     * Process and send a text completion request with optional preset & instruct
273
     * @param {TextCompletionPayload} requestData
274
     * @param {Object} options - Configuration options
275
     * @param {string?} [options.presetName] - Name of the preset to use for generation settings
276
     * @param {string?} [options.instructName] - Name of instruct preset for message formatting
277
     * @param {Partial<InstructSettings>?} [options.instructSettings] - Override instruct settings
278
     * @param {boolean} extractData - Whether to extract structured data from response
279
     * @param {AbortSignal?} [signal]
280
     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
281
     * @throws {Error}
282
     */
283
    static async processRequest(requestData, options = {}, extractData = true, signal = null) {
284
        const { presetName, instructName } = options;
285

286
        // remove any undefined params in given request data
287
        requestData = this.createRequestData(requestData);
288

289
        /** @type {InstructSettings | undefined} */
290
        let instructPreset;
291
        const prompt = requestData.prompt;
292
        // Handle instruct formatting if requested
293
        if (Array.isArray(prompt)) {
294
            if (instructName) {
295
                const instructPresetManager = getPresetManager('instruct');
296
                instructPreset = instructPresetManager?.getCompletionPresetByName(instructName);
297
                if (instructPreset) {
298
                    requestData.prompt = this.constructPrompt(prompt, instructPreset, options.instructSettings);
299
                    const stoppingStrings = getInstructStoppingSequences({ customInstruct: instructPreset, useStopStrings: false });
300
                    requestData.stop = stoppingStrings;
301
                    requestData.stopping_strings = stoppingStrings;
302
                } else {
303
                    console.warn(`Instruct preset "${instructName}" not found, using basic formatting`);
304
                    requestData.prompt = prompt.map(x => x.content).join('\n\n');
305
                }
306
            } else {
307
                requestData.prompt = prompt.map(x => x.content).join('\n\n');
308
            }
309
        } else if (typeof prompt === 'string') {
310
            requestData.prompt = prompt;
311
        }
312

313
        // Apply generation preset if specified
314
        if (presetName) {
315
            const presetManager = getPresetManager(this.TYPE);
316
            if (presetManager) {
317
                const preset = presetManager.getCompletionPresetByName(presetName);
318
                if (preset) {
319
                    // Convert preset to payload and merge with custom data
320
                    requestData = this.presetToGeneratePayload(preset, {}, requestData);
321
                } else {
322
                    console.warn(`Preset "${presetName}" not found, continuing with default settings`);
323
                }
324
            } else {
325
                console.warn('Preset manager not found, continuing with default settings');
326
            }
327
        }
328

329
        const response = await this.sendRequest(requestData, extractData, signal);
330

331
        // Remove stopping strings from the end
332
        if (!requestData.stream && extractData) {
333
            /** @type {ExtractedData} */
334
            // @ts-ignore
335
            const extractedData = response;
336

337
            let message = extractedData.content;
338

339
            message = message.replace(/[^\S\r\n]+$/gm, '');
340

341
            if (requestData.stopping_strings) {
342
                for (const stoppingString of requestData.stopping_strings) {
343
                    if (stoppingString.length) {
344
                        for (let j = stoppingString.length; j > 0; j--) {
345
                            if (message.slice(-j) === stoppingString.slice(0, j)) {
346
                                message = message.slice(0, -j);
347
                                break;
348
                            }
349
                        }
350
                    }
351
                }
352
            }
353

354
            if (instructPreset) {
355
                [
356
                    instructPreset.stop_sequence,
357
                    instructPreset.input_sequence,
358
                ].forEach(sequence => {
359
                    if (sequence?.trim()) {
360
                        const index = message.indexOf(sequence);
361
                        if (index !== -1) {
362
                            message = message.substring(0, index);
363
                        }
364
                    }
365
                });
366

367
                [
368
                    instructPreset.output_sequence,
369
                    instructPreset.last_output_sequence,
370
                ].forEach(sequences => {
371
                    if (sequences) {
372
                        sequences.split('\n')
373
                            .filter(line => line.trim() !== '')
374
                            .forEach(line => {
375
                                message = message.replaceAll(line, '');
376
                            });
377
                    }
378
                });
379
            }
380

381
            extractedData.content = message;
382
        }
383

384
        return response;
385
    }
386

387
    /**
388
     * Converts a preset to a valid text completion payload.
389
     * Only supports temperature.
390
     * @param {Object} preset - The preset configuration
391
     * @param {Object} overridePreset - Additional parameters to override preset values
392
     * @param {Object} overridePayload - Additional parameters to override payload values
393
     * @returns {Object} - Formatted payload for text completion API
394
     */
395
    static presetToGeneratePayload(preset, overridePreset = {}, overridePayload = {}) {
396
        if (!preset || typeof preset !== 'object') {
397
            throw new Error('Invalid preset: must be an object');
398
        }
399

400
        // apply preset overrides
401
        preset = { ...preset, ...overridePreset };
402

403
        // Only take fields from the preset specified in setting_names to use as TextCompletionSettings
404
        const settings = structuredClone(textgenerationwebui_settings);
405
        for (const [key, value] of Object.entries(preset)) {
406
            if (!setting_names.includes(key)) continue;
407
            settings[key] = value;
408
        }
409

410
        // convert to a generation payload
411
        const payload = createTextGenGenerationData(settings, overridePayload.model, overridePayload.prompt, preset.genamt);
412

413
        // apply overrides
414
        return this.createRequestData({ ...payload, ...overridePayload });
415
    }
416
}
417

418
/**
419
 * Creates & sends a chat completion request.
420
 */
421
export class ChatCompletionService {
422
    static TYPE = 'openai';
423

424
    /**
425
     * @param {ChatCompletionPayload} custom
426
     * @returns {ChatCompletionPayload}
427
     */
428
    static createRequestData({ stream = false, messages, model, chat_completion_source, max_tokens, temperature, custom_url, reverse_proxy, proxy_password, custom_prompt_post_processing, ...props }) {
429
        const payload = {
430
            stream,
431
            messages,
432
            model,
433
            chat_completion_source,
434
            max_tokens,
435
            temperature,
436
            custom_url,
437
            reverse_proxy,
438
            proxy_password,
439
            custom_prompt_post_processing,
440
            use_sysprompt: true,
441
            ...props,
442
        };
443

444
        // Remove undefined values to avoid API errors
445
        Object.keys(payload).forEach(key => {
446
            if (payload[key] === undefined) {
447
                delete payload[key];
448
            }
449
        });
450

451
        return payload;
452
    }
453

454
    /**
455
     * Sends a chat completion request
456
     * @param {ChatCompletionPayload} data Request data
457
     * @param {boolean?} extractData Extract message from the response. Default true
458
     * @param {AbortSignal?} signal Abort signal
459
     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
460
     * @throws {Error}
461
     */
462
    static async sendRequest(data, extractData = true, signal = null) {
463
        const response = await fetch('/api/backends/chat-completions/generate', {
464
            method: 'POST',
465
            headers: getRequestHeaders(),
466
            cache: 'no-cache',
467
            body: JSON.stringify(data),
468
            signal: signal ?? new AbortController().signal,
469
        });
470

471
        if (!data.stream) {
472
            const json = await response.json();
473
            if (!response.ok || json.error) {
474
                throw new Error(String(json.error?.message || 'Response not OK'));
475
            }
476

477
            if (!extractData) {
478
                return json;
479
            }
480

481
            const result = {
482
                content: extractMessageFromData(json, this.TYPE),
483
                reasoning: extractReasoningFromData(json, {
484
                    mainApi: this.TYPE,
485
                    textGenType: data.chat_completion_source,
486
                    ignoreShowThoughts: true,
487
                }),
488
            };
489
            // Try parse JSON
490
            if (data.json_schema) {
491
                result.content = JSON.parse(extractJsonFromData(json, { mainApi: this.TYPE, chatCompletionSource: data.chat_completion_source }));
492
            }
493
            return result;
494
        }
495

496
        if (!response.ok) {
497
            const text = await response.text();
498
            tryParseStreamingError(response, text, { quiet: true });
499

500
            throw new Error(`Got response status ${response.status}`);
501
        }
502

503
        const eventStream = new EventSourceStream();
504
        response.body.pipeThrough(eventStream);
505
        const reader = eventStream.readable.getReader();
506
        return async function* streamData() {
507
            let text = '';
508
            const swipes = [];
509
            const state = { reasoning: '', images: [], signature: '', toolSignatures: {} };
510
            while (true) {
511
                const { done, value } = await reader.read();
512
                if (done) return;
513
                const rawData = value.data;
514
                if (rawData === '[DONE]') return;
515
                tryParseStreamingError(response, rawData, { quiet: true });
516
                const parsed = JSON.parse(rawData);
517

518
                const reply = getStreamingReply(parsed, state, {
519
                    chatCompletionSource: data.chat_completion_source,
520
                    overrideShowThoughts: true,
521
                });
522
                if (Array.isArray(parsed?.choices) && parsed?.choices?.[0]?.index > 0) {
523
                    const swipeIndex = parsed.choices[0].index - 1;
524
                    swipes[swipeIndex] = (swipes[swipeIndex] || '') + reply;
525
                } else {
526
                    text += reply;
527
                }
528

529
                yield { text, swipes: swipes, state };
530
            }
531
        };
532
    }
533

534
    /**
535
     * Process and send a chat completion request with optional preset
536
     * @param {ChatCompletionPayload} requestData - payload data, overriding preset if given
537
     * @param {Object} options - Configuration options
538
     * @param {string?} [options.presetName] - Name of the preset to use for generation settings
539
     * @param {boolean} [extractData=true] - Whether to extract structured data from response
540
     * @param {AbortSignal?} [signal] - Abort signal
541
     * @returns {Promise<ExtractedData | (() => AsyncGenerator<StreamResponse>)>} If not streaming, returns extracted data; if streaming, returns a function that creates an AsyncGenerator
542
     * @throws {Error}
543
     */
544
    static async processRequest(requestData, options, extractData = true, signal = null) {
545
        const { presetName } = options;
546
        requestData = this.createRequestData(requestData);
547

548
        // Apply generation preset if specified
549
        if (presetName) {
550
            const presetManager = getPresetManager(this.TYPE);
551
            if (presetManager) {
552
                const preset = presetManager.getCompletionPresetByName(presetName);
553
                if (preset) {
554
                    // Convert preset to payload and merge with custom parameters
555
                    requestData = await this.presetToGeneratePayload(preset, {}, requestData);
556
                } else {
557
                    console.warn(`Preset "${presetName}" not found, continuing with default settings`);
558
                }
559
            } else {
560
                console.warn('Preset manager not found, continuing with default settings');
561
            }
562
        }
563

564
        return await this.sendRequest(requestData, extractData, signal);
565
    }
566

567
    /**
568
     * Converts a preset to a valid chat completion payload
569
     * Only supports temperature.
570
     * @param {Object} preset - The preset configuration
571
     * @param {Object} overridePreset - Additional parameters to override preset values
572
     * @param {Object} overridePayload - Additional parameters to override payload values
573
     * @returns {Promise<any>} - Formatted payload for chat completion API
574
     */
575
    static async presetToGeneratePayload(preset, overridePreset = {}, overridePayload = {}) {
576
        if (!preset || typeof preset !== 'object') {
577
            throw new Error('Invalid preset: must be an object');
578
        }
579

580
        // apply preset overrides
581
        preset = { ...preset, ...overridePreset };
582

583
        // Fix any fields before converting to settings
584
        preset.bias_preset_selected = preset.bias_presets !== undefined ? preset.bias_preset_selected : undefined;  // presets might have bias_preset_selected but not bias_presets, but settings need both or neither.
585

586
        // Convert from preset to ChatCompletionSettings
587
        const settings = structuredClone(oai_settings);
588
        for (const [key, value] of Object.entries(preset)) {
589
            const settingToUpdate = settingsToUpdate[key];
590
            if (!settingToUpdate) continue;
591
            settings[settingToUpdate[1]] = value;
592
        }
593

594
        // Ensure api-url is properly applied for all sources that accept it
595
        ['custom_url', 'vertexai_region', 'zai_endpoint', 'siliconflow_endpoint', 'minimax_endpoint'].forEach(field => {
596
            // The order is: connection profile => CC preset => CC settings
597
            overridePayload[field] = overridePayload[field] || settings[field] || oai_settings[field];
598
        });
599

600
        // Convert from settings to generation payload
601
        const data = await createGenerationParameters(settings, overridePayload.model, 'quiet', overridePayload.messages);
602
        const payload = data.generate_data;
603

604
        // apply overrides
605
        return this.createRequestData({ ...payload, ...overridePayload });
606
    }
607
}