1
import { Readable } from 'node:stream';
2
import fetch from 'node-fetch';
3
import express from 'express';
4
import _ from 'lodash';
5

6
import {
7
    TEXTGEN_TYPES,
8
    TOGETHERAI_KEYS,
9
    OLLAMA_KEYS,
10
    INFERMATICAI_KEYS,
11
    OPENROUTER_KEYS,
12
    VLLM_KEYS,
13
    FEATHERLESS_KEYS,
14
    OPENAI_KEYS,
15
} from '../../constants.js';
16
import { forwardFetchResponse, trimV1, getConfigValue } from '../../util.js';
17
import { setAdditionalHeaders } from '../../additional-headers.js';
18
import { createHash } from 'node:crypto';
19

20
export const router = express.Router();
21

22
/**
23
 * Special boy's steaming routine. Wrap this abomination into proper SSE stream.
24
 * @param {import('node-fetch').Response} jsonStream JSON stream
25
 * @param {import('express').Request} request Express request
26
 * @param {import('express').Response} response Express response
27
 * @returns {Promise<any>} Nothing valuable
28
 */
29
async function parseOllamaStream(jsonStream, request, response) {
30
    try {
31
        if (!jsonStream.body) {
32
            throw new Error('No body in the response');
33
        }
34

35
        let partialData = '';
36
        jsonStream.body.on('data', (data) => {
37
            const chunk = data.toString();
38
            partialData += chunk;
39
            while (true) {
40
                let json;
41
                try {
42
                    json = JSON.parse(partialData);
43
                } catch (e) {
44
                    break;
45
                }
46
                const text = json.response || '';
47
                const thinking = json.thinking || '';
48
                const chunk = { choices: [{ text, thinking }] };
49
                response.write(`data: ${JSON.stringify(chunk)}\n\n`);
50
                partialData = '';
51
            }
52
        });
53

54
        request.socket.on('close', function () {
55
            if (jsonStream.body instanceof Readable) jsonStream.body.destroy();
56
            response.end();
57
        });
58

59
        jsonStream.body.on('end', () => {
60
            console.info('Streaming request finished');
61
            response.write('data: [DONE]\n\n');
62
            response.end();
63
        });
64
    } catch (error) {
65
        console.error('Error forwarding streaming response:', error);
66
        if (!response.headersSent) {
67
            return response.status(500).send({ error: true });
68
        } else {
69
            return response.end();
70
        }
71
    }
72
}
73

74
/**
75
 * Abort KoboldCpp generation request.
76
 * @param {import('express').Request} request the generation request
77
 * @param {string} url Server base URL
78
 * @returns {Promise<void>} Promise resolving when we are done
79
 */
80
async function abortKoboldCppRequest(request, url) {
81
    try {
82
        console.info('Aborting Kobold generation...');
83
        const args = {
84
            method: 'POST',
85
            headers: {},
86
        };
87

88
        setAdditionalHeaders(request, args, url);
89
        const abortResponse = await fetch(`${url}/api/extra/abort`, args);
90

91
        if (!abortResponse.ok) {
92
            console.error('Error sending abort request to Kobold:', abortResponse.status, abortResponse.statusText);
93
        }
94
    } catch (error) {
95
        console.error(error);
96
    }
97
}
98

99
//************** Ooba/OpenAI text completions API
100
router.post('/status', async function (request, response) {
101
    if (!request.body) return response.sendStatus(400);
102

103
    try {
104
        if (request.body.api_server.indexOf('localhost') !== -1) {
105
            request.body.api_server = request.body.api_server.replace('localhost', '127.0.0.1');
106
        }
107

108
        console.debug('Trying to connect to API', request.body);
109
        const baseUrl = trimV1(request.body.api_server);
110

111
        const args = {
112
            headers: { 'Content-Type': 'application/json' },
113
        };
114

115
        setAdditionalHeaders(request, args, baseUrl);
116

117
        const apiType = request.body.api_type;
118
        let url = baseUrl;
119
        let result = '';
120

121
        switch (apiType) {
122
            case TEXTGEN_TYPES.GENERIC:
123
            case TEXTGEN_TYPES.OOBA:
124
            case TEXTGEN_TYPES.VLLM:
125
            case TEXTGEN_TYPES.APHRODITE:
126
            case TEXTGEN_TYPES.KOBOLDCPP:
127
            case TEXTGEN_TYPES.LLAMACPP:
128
            case TEXTGEN_TYPES.INFERMATICAI:
129
            case TEXTGEN_TYPES.OPENROUTER:
130
            case TEXTGEN_TYPES.FEATHERLESS:
131
                url += '/v1/models';
132
                break;
133
            case TEXTGEN_TYPES.DREAMGEN:
134
                url += '/api/openai/v1/models';
135
                break;
136
            case TEXTGEN_TYPES.MANCER:
137
                url += '/oai/v1/models';
138
                break;
139
            case TEXTGEN_TYPES.TABBY:
140
                url += '/v1/model/list';
141
                break;
142
            case TEXTGEN_TYPES.TOGETHERAI:
143
                url += '/api/models?&info';
144
                break;
145
            case TEXTGEN_TYPES.OLLAMA:
146
                url += '/api/tags';
147
                break;
148
            case TEXTGEN_TYPES.HUGGINGFACE:
149
                url += '/info';
150
                break;
151
        }
152

153
        const modelsReply = await fetch(url, args);
154
        const isPossiblyLmStudio = modelsReply.headers.get('x-powered-by') === 'Express';
155

156
        if (!modelsReply.ok) {
157
            console.error('Models endpoint is offline.');
158
            return response.sendStatus(400);
159
        }
160

161
        /** @type {any} */
162
        let data = await modelsReply.json();
163

164
        // Rewrap to OAI-like response
165
        if (apiType === TEXTGEN_TYPES.TOGETHERAI && Array.isArray(data)) {
166
            data = { data: data.map(x => ({ id: x.name, ...x })) };
167
        }
168

169
        if (apiType === TEXTGEN_TYPES.OLLAMA && Array.isArray(data.models)) {
170
            data = { data: data.models.map(x => ({ id: x.name, ...x })) };
171
        }
172

173
        if (apiType === TEXTGEN_TYPES.HUGGINGFACE) {
174
            data = { data: [] };
175
        }
176

177
        if (!Array.isArray(data.data)) {
178
            console.error('Models response is not an array.');
179
            return response.sendStatus(400);
180
        }
181

182
        const modelIds = data.data.map(x => x.id);
183
        console.info('Models available:', modelIds);
184

185
        // Set result to the first model ID
186
        result = modelIds[0] || 'Valid';
187

188
        if (apiType === TEXTGEN_TYPES.OOBA && !isPossiblyLmStudio) {
189
            try {
190
                const modelInfoUrl = baseUrl + '/v1/internal/model/info';
191
                const modelInfoReply = await fetch(modelInfoUrl, args);
192

193
                if (modelInfoReply.ok) {
194
                    /** @type {any} */
195
                    const modelInfo = await modelInfoReply.json();
196
                    console.debug('Ooba model info:', modelInfo);
197

198
                    const modelName = modelInfo?.model_name;
199
                    result = modelName || result;
200
                    response.setHeader('x-supports-tokenization', 'true');
201
                }
202
            } catch (error) {
203
                console.error(`Failed to get Ooba model info: ${error}`);
204
            }
205
        } else if (apiType === TEXTGEN_TYPES.TABBY) {
206
            try {
207
                const modelInfoUrl = baseUrl + '/v1/model';
208
                const modelInfoReply = await fetch(modelInfoUrl, args);
209

210
                if (modelInfoReply.ok) {
211
                    /** @type {any} */
212
                    const modelInfo = await modelInfoReply.json();
213
                    console.debug('Tabby model info:', modelInfo);
214

215
                    const modelName = modelInfo?.id;
216
                    result = modelName || result;
217
                } else {
218
                    // TabbyAPI returns an error 400 if a model isn't loaded
219

220
                    result = 'None';
221
                }
222
            } catch (error) {
223
                console.error(`Failed to get TabbyAPI model info: ${error}`);
224
            }
225
        }
226

227
        return response.send({ result, data: data.data });
228
    } catch (error) {
229
        console.error(error);
230
        return response.sendStatus(500);
231
    }
232
});
233

234
router.post('/props', async function (request, response) {
235
    if (!request.body.api_server) return response.sendStatus(400);
236

237
    try {
238
        const baseUrl = trimV1(request.body.api_server);
239
        const args = {
240
            headers: {},
241
        };
242

243
        setAdditionalHeaders(request, args, baseUrl);
244

245
        const apiType = request.body.api_type;
246
        let propsUrl = baseUrl + '/props';
247
        if (apiType === TEXTGEN_TYPES.LLAMACPP && request.body.model) {
248
            propsUrl += `?model=${encodeURIComponent(request.body.model)}`;
249
            console.debug(`Querying llama-server props with model parameter: ${request.body.model}`);
250
        }
251
        const propsReply = await fetch(propsUrl, args);
252

253
        if (!propsReply.ok) {
254
            return response.sendStatus(400);
255
        }
256

257
        /** @type {any} */
258
        const props = await propsReply.json();
259
        // TEMPORARY: llama.cpp's /props endpoint has a bug which replaces the last newline with a \0
260
        if (apiType === TEXTGEN_TYPES.LLAMACPP && props.chat_template && props.chat_template.endsWith('\u0000')) {
261
            props.chat_template = props.chat_template.slice(0, -1) + '\n';
262
        }
263
        props.chat_template_hash = createHash('sha256').update(props.chat_template).digest('hex');
264
        console.debug(`Model properties: ${JSON.stringify(props)}`);
265
        return response.send(props);
266
    } catch (error) {
267
        console.error(error);
268
        return response.sendStatus(500);
269
    }
270
});
271

272
router.post('/generate', async function (request, response) {
273
    if (!request.body) return response.sendStatus(400);
274

275
    try {
276
        if (request.body.api_server.indexOf('localhost') !== -1) {
277
            request.body.api_server = request.body.api_server.replace('localhost', '127.0.0.1');
278
        }
279

280
        const apiType = request.body.api_type;
281
        const baseUrl = request.body.api_server;
282
        console.debug(request.body);
283

284
        const controller = new AbortController();
285
        request.socket.removeAllListeners('close');
286
        request.socket.on('close', async function () {
287
            if (request.body.api_type === TEXTGEN_TYPES.KOBOLDCPP && !response.writableEnded) {
288
                await abortKoboldCppRequest(request, trimV1(baseUrl));
289
            }
290

291
            controller.abort();
292
        });
293

294
        let url = trimV1(baseUrl);
295

296
        switch (request.body.api_type) {
297
            case TEXTGEN_TYPES.GENERIC:
298
            case TEXTGEN_TYPES.VLLM:
299
            case TEXTGEN_TYPES.FEATHERLESS:
300
            case TEXTGEN_TYPES.APHRODITE:
301
            case TEXTGEN_TYPES.OOBA:
302
            case TEXTGEN_TYPES.TABBY:
303
            case TEXTGEN_TYPES.KOBOLDCPP:
304
            case TEXTGEN_TYPES.TOGETHERAI:
305
            case TEXTGEN_TYPES.INFERMATICAI:
306
            case TEXTGEN_TYPES.HUGGINGFACE:
307
                url += '/v1/completions';
308
                break;
309
            case TEXTGEN_TYPES.DREAMGEN:
310
                url += '/api/openai/v1/completions';
311
                break;
312
            case TEXTGEN_TYPES.MANCER:
313
                url += '/oai/v1/completions';
314
                break;
315
            case TEXTGEN_TYPES.LLAMACPP:
316
                url += '/completion';
317
                break;
318
            case TEXTGEN_TYPES.OLLAMA:
319
                url += '/api/generate';
320
                break;
321
            case TEXTGEN_TYPES.OPENROUTER:
322
                url += '/v1/chat/completions';
323
                break;
324
        }
325

326
        const args = {
327
            method: 'POST',
328
            body: JSON.stringify(request.body),
329
            headers: { 'Content-Type': 'application/json' },
330
            signal: controller.signal,
331
            timeout: 0,
332
        };
333

334
        setAdditionalHeaders(request, args, baseUrl);
335

336
        if (request.body.api_type === TEXTGEN_TYPES.TOGETHERAI) {
337
            request.body = _.pickBy(request.body, (_, key) => TOGETHERAI_KEYS.includes(key));
338
            args.body = JSON.stringify(request.body);
339
        }
340

341
        if (request.body.api_type === TEXTGEN_TYPES.INFERMATICAI) {
342
            request.body = _.pickBy(request.body, (_, key) => INFERMATICAI_KEYS.includes(key));
343
            args.body = JSON.stringify(request.body);
344
        }
345

346
        if (request.body.api_type === TEXTGEN_TYPES.FEATHERLESS) {
347
            request.body = _.pickBy(request.body, (_, key) => FEATHERLESS_KEYS.includes(key));
348
            args.body = JSON.stringify(request.body);
349
        }
350

351
        if (request.body.api_type === TEXTGEN_TYPES.DREAMGEN) {
352
            args.body = JSON.stringify(request.body);
353
        }
354

355
        if (request.body.api_type === TEXTGEN_TYPES.GENERIC) {
356
            request.body = _.pickBy(request.body, (_, key) => OPENAI_KEYS.includes(key));
357
            if (Array.isArray(request.body.stop)) { request.body.stop = request.body.stop.slice(0, 4); }
358
            args.body = JSON.stringify(request.body);
359
        }
360

361
        if (request.body.api_type === TEXTGEN_TYPES.OPENROUTER) {
362
            if (Array.isArray(request.body.provider) && request.body.provider.length > 0) {
363
                request.body.provider = {
364
                    allow_fallbacks: request.body.allow_fallbacks ?? true,
365
                    order: request.body.provider,
366
                };
367
            } else {
368
                delete request.body.provider;
369
            }
370

371
            if (Array.isArray(request.body.quantizations) && request.body.quantizations.length > 0) {
372
                request.body.provider ??= {};
373
                request.body.provider.quantizations = request.body.quantizations;
374
            }
375

376
            request.body = _.pickBy(request.body, (_, key) => OPENROUTER_KEYS.includes(key));
377
            args.body = JSON.stringify(request.body);
378
        }
379

380
        if (request.body.api_type === TEXTGEN_TYPES.VLLM) {
381
            request.body = _.pickBy(request.body, (_, key) => VLLM_KEYS.includes(key));
382
            args.body = JSON.stringify(request.body);
383
        }
384

385
        if (request.body.api_type === TEXTGEN_TYPES.OLLAMA) {
386
            const keepAlive = Number(getConfigValue('ollama.keepAlive', -1, 'number'));
387
            const numBatch = Number(getConfigValue('ollama.batchSize', -1, 'number'));
388
            if (numBatch > 0) {
389
                request.body.num_batch = numBatch;
390
            }
391
            args.body = JSON.stringify({
392
                model: request.body.model,
393
                prompt: request.body.prompt,
394
                stream: request.body.stream ?? false,
395
                keep_alive: keepAlive,
396
                raw: true,
397
                options: _.pickBy(request.body, (_, key) => OLLAMA_KEYS.includes(key)),
398
            });
399
        }
400

401
        if (request.body.api_type === TEXTGEN_TYPES.OLLAMA && request.body.stream) {
402
            const stream = await fetch(url, args);
403
            parseOllamaStream(stream, request, response);
404
        } else if (request.body.stream) {
405
            const completionsStream = await fetch(url, args);
406
            // Pipe remote SSE stream to Express response
407
            await forwardFetchResponse(completionsStream, response);
408
        } else {
409
            const completionsReply = await fetch(url, args);
410

411
            if (completionsReply.ok) {
412
                /** @type {any} */
413
                const data = await completionsReply.json();
414
                console.debug('Endpoint response:', data);
415

416
                // Map InfermaticAI response to OAI completions format
417
                if (apiType === TEXTGEN_TYPES.INFERMATICAI) {
418
                    data.choices = (data?.choices || []).map(choice => ({ text: choice?.message?.content || choice.text, logprobs: choice?.logprobs, index: choice?.index }));
419
                }
420

421
                return response.send(data);
422
            } else {
423
                const text = await completionsReply.text();
424
                const errorBody = { error: true, status: completionsReply.status, response: text };
425

426
                return !response.headersSent
427
                    ? response.send(errorBody)
428
                    : response.end();
429
            }
430
        }
431
    } catch (error) {
432
        const status = error?.status ?? error?.code ?? 'UNKNOWN';
433
        const text = error?.error ?? error?.statusText ?? error?.message ?? 'Unknown error on /generate endpoint';
434
        let value = { error: true, status: status, response: text };
435
        console.error('Endpoint error:', error);
436

437
        return !response.headersSent
438
            ? response.send(value)
439
            : response.end();
440
    }
441
});
442

443
const ollama = express.Router();
444

445
ollama.post('/download', async function (request, response) {
446
    try {
447
        if (!request.body.name || !request.body.api_server) return response.sendStatus(400);
448

449
        const name = request.body.name;
450
        const url = String(request.body.api_server).replace(/\/$/, '');
451
        console.debug('Pulling Ollama model:', name);
452

453
        const fetchResponse = await fetch(`${url}/api/pull`, {
454
            method: 'POST',
455
            headers: { 'Content-Type': 'application/json' },
456
            body: JSON.stringify({
457
                name: name,
458
                stream: false,
459
            }),
460
        });
461

462
        if (!fetchResponse.ok) {
463
            console.error('Download error:', fetchResponse.status, fetchResponse.statusText);
464
            return response.status(500).send({ error: true });
465
        }
466

467
        console.debug('Ollama pull response:', await fetchResponse.json());
468
        return response.send({ ok: true });
469
    } catch (error) {
470
        console.error(error);
471
        return response.sendStatus(500);
472
    }
473
});
474

475
ollama.post('/caption-image', async function (request, response) {
476
    try {
477
        if (!request.body.server_url || !request.body.model) {
478
            return response.sendStatus(400);
479
        }
480

481
        console.debug('Ollama caption request:', request.body);
482
        const baseUrl = trimV1(request.body.server_url);
483

484
        const fetchResponse = await fetch(`${baseUrl}/api/generate`, {
485
            method: 'POST',
486
            headers: { 'Content-Type': 'application/json' },
487
            body: JSON.stringify({
488
                model: request.body.model,
489
                prompt: request.body.prompt,
490
                images: [request.body.image],
491
                stream: false,
492
            }),
493
        });
494

495
        if (!fetchResponse.ok) {
496
            const errorText = await fetchResponse.text();
497
            console.error('Ollama caption error:', fetchResponse.status, fetchResponse.statusText, errorText);
498
            return response.status(500).send({ error: true });
499
        }
500

501
        /** @type {any} */
502
        const data = await fetchResponse.json();
503
        console.debug('Ollama caption response:', data);
504

505
        const caption = data?.response || '';
506

507
        if (!caption) {
508
            console.error('Ollama caption is empty.');
509
            return response.status(500).send({ error: true });
510
        }
511

512
        return response.send({ caption });
513
    } catch (error) {
514
        console.error(error);
515
        return response.sendStatus(500);
516
    }
517
});
518

519
const llamacpp = express.Router();
520

521
llamacpp.post('/props', async function (request, response) {
522
    try {
523
        if (!request.body.server_url) {
524
            return response.sendStatus(400);
525
        }
526

527
        console.debug('LlamaCpp props request:', request.body);
528
        const baseUrl = trimV1(request.body.server_url);
529

530
        const fetchResponse = await fetch(`${baseUrl}/props`, {
531
            method: 'GET',
532
        });
533

534
        if (!fetchResponse.ok) {
535
            console.error('LlamaCpp props error:', fetchResponse.status, fetchResponse.statusText);
536
            return response.status(500).send({ error: true });
537
        }
538

539
        const data = await fetchResponse.json();
540
        console.debug('LlamaCpp props response:', data);
541

542
        return response.send(data);
543
    } catch (error) {
544
        console.error(error);
545
        return response.sendStatus(500);
546
    }
547
});
548

549
llamacpp.post('/slots', async function (request, response) {
550
    try {
551
        if (!request.body.server_url) {
552
            return response.sendStatus(400);
553
        }
554
        if (!/^(erase|info|restore|save)$/.test(request.body.action)) {
555
            return response.sendStatus(400);
556
        }
557

558
        console.debug('LlamaCpp slots request:', request.body);
559
        const baseUrl = trimV1(request.body.server_url);
560

561
        let fetchResponse;
562
        if (request.body.action === 'info') {
563
            fetchResponse = await fetch(`${baseUrl}/slots`, {
564
                method: 'GET',
565
            });
566
        } else {
567
            if (!/^\d+$/.test(request.body.id_slot)) {
568
                return response.sendStatus(400);
569
            }
570
            if (request.body.action !== 'erase' && !request.body.filename) {
571
                return response.sendStatus(400);
572
            }
573

574
            fetchResponse = await fetch(`${baseUrl}/slots/${request.body.id_slot}?action=${request.body.action}`, {
575
                method: 'POST',
576
                headers: { 'Content-Type': 'application/json' },
577
                body: JSON.stringify({
578
                    filename: request.body.action !== 'erase' ? `${request.body.filename}` : undefined,
579
                }),
580
            });
581
        }
582

583
        if (!fetchResponse.ok) {
584
            console.error('LlamaCpp slots error:', fetchResponse.status, fetchResponse.statusText);
585
            return response.status(500).send({ error: true });
586
        }
587

588
        const data = await fetchResponse.json();
589
        console.debug('LlamaCpp slots response:', data);
590

591
        return response.send(data);
592
    } catch (error) {
593
        console.error(error);
594
        return response.sendStatus(500);
595
    }
596
});
597

598
const tabby = express.Router();
599

600
tabby.post('/download', async function (request, response) {
601
    try {
602
        const baseUrl = String(request.body.api_server).replace(/\/$/, '');
603

604
        const args = {
605
            method: 'POST',
606
            headers: { 'Content-Type': 'application/json' },
607
            body: JSON.stringify(request.body),
608
            timeout: 0,
609
        };
610

611
        setAdditionalHeaders(request, args, baseUrl);
612

613
        // Check key permissions
614
        const permissionResponse = await fetch(`${baseUrl}/v1/auth/permission`, {
615
            headers: args.headers,
616
        });
617

618
        if (permissionResponse.ok) {
619
            /** @type {any} */
620
            const permissionJson = await permissionResponse.json();
621

622
            if (permissionJson.permission !== 'admin') {
623
                return response.status(403).send({ error: true });
624
            }
625
        } else {
626
            console.error('API Permission error:', permissionResponse.status, permissionResponse.statusText);
627
            return response.status(500).send({ error: true });
628
        }
629

630
        const fetchResponse = await fetch(`${baseUrl}/v1/download`, args);
631

632
        if (!fetchResponse.ok) {
633
            console.error('Download error:', fetchResponse.status, fetchResponse.statusText);
634
            return response.status(500).send({ error: true });
635
        }
636

637
        return response.send({ ok: true });
638
    } catch (error) {
639
        console.error(error);
640
        return response.sendStatus(500);
641
    }
642
});
643

644
router.use('/ollama', ollama);
645
router.use('/llamacpp', llamacpp);
646
router.use('/tabby', tabby);