Add OpenRouter as a text completion source

This commit is contained in:
Cohee
2024-02-25 22:47:07 +02:00
parent 9e5505a7d4
commit f962ad5c02
12 changed files with 201 additions and 24 deletions

View File

@ -4,7 +4,7 @@ const _ = require('lodash');
const Readable = require('stream').Readable;
const { jsonParser } = require('../../express-common');
const { TEXTGEN_TYPES, TOGETHERAI_KEYS, OLLAMA_KEYS, INFERMATICAI_KEYS } = require('../../constants');
const { TEXTGEN_TYPES, TOGETHERAI_KEYS, OLLAMA_KEYS, INFERMATICAI_KEYS, OPENROUTER_KEYS } = require('../../constants');
const { forwardFetchResponse, trimV1 } = require('../../util');
const { setAdditionalHeaders } = require('../../additional-headers');
@ -107,6 +107,7 @@ router.post('/status', jsonParser, async function (request, response) {
case TEXTGEN_TYPES.KOBOLDCPP:
case TEXTGEN_TYPES.LLAMACPP:
case TEXTGEN_TYPES.INFERMATICAI:
case TEXTGEN_TYPES.OPENROUTER:
url += '/v1/models';
break;
case TEXTGEN_TYPES.MANCER:
@ -209,6 +210,7 @@ router.post('/generate', jsonParser, async function (request, response) {
request.body.api_server = request.body.api_server.replace('localhost', '127.0.0.1');
}
const apiType = request.body.api_type;
const baseUrl = request.body.api_server;
console.log(request.body);
@ -245,6 +247,9 @@ router.post('/generate', jsonParser, async function (request, response) {
case TEXTGEN_TYPES.OLLAMA:
url += '/api/generate';
break;
case TEXTGEN_TYPES.OPENROUTER:
url += '/v1/chat/completions';
break;
}
}
@ -268,6 +273,11 @@ router.post('/generate', jsonParser, async function (request, response) {
args.body = JSON.stringify(request.body);
}
if (request.body.api_type === TEXTGEN_TYPES.OPENROUTER) {
request.body = _.pickBy(request.body, (_, key) => OPENROUTER_KEYS.includes(key));
args.body = JSON.stringify(request.body);
}
if (request.body.api_type === TEXTGEN_TYPES.OLLAMA) {
args.body = JSON.stringify({
model: request.body.model,
@ -300,7 +310,7 @@ router.post('/generate', jsonParser, async function (request, response) {
}
// Map InfermaticAI response to OAI completions format
if (completionsReply.url.includes('https://api.totalgpt.ai')) {
if (apiType === TEXTGEN_TYPES.INFERMATICAI) {
data['choices'] = (data?.choices || []).map(choice => ({ text: choice.message.content }));
}