Blame Raw
Cohee · 51ad27fb · · 265 lines (9.8 KB)
2 contributors
1import { TEXTGEN_TYPES, OPENROUTER_HEADERS, FEATHERLESS_HEADERS } from './constants.js';
2import { SECRET_KEYS, readSecret } from './endpoints/secrets.js';
3import { getConfigValue } from './util.js';
4
5/**
6 * Gets the headers for the Mancer API.
7 * @param {import('./users.js').UserDirectoryList} directories User directories
8 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
9 * @returns {object} Headers for the request
10 */
11function getMancerHeaders(directories, secretId = null) {
12 const apiKey = readSecret(directories, SECRET_KEYS.MANCER, secretId);
13
14 return apiKey ? ({
15 'X-API-KEY': apiKey,
16 'Authorization': `Bearer ${apiKey}`,
17 }) : {};
18}
19
20/**
21 * Gets the headers for the TogetherAI API.
22 * @param {import('./users.js').UserDirectoryList} directories User directories
23 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
24 * @returns {object} Headers for the request
25 */
26function getTogetherAIHeaders(directories, secretId = null) {
27 const apiKey = readSecret(directories, SECRET_KEYS.TOGETHERAI, secretId);
28
29 return apiKey ? ({
30 'Authorization': `Bearer ${apiKey}`,
31 }) : {};
32}
33
34/**
35 * Gets the headers for the InfermaticAI API.
36 * @param {import('./users.js').UserDirectoryList} directories User directories
37 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
38 * @returns {object} Headers for the request
39 */
40function getInfermaticAIHeaders(directories, secretId = null) {
41 const apiKey = readSecret(directories, SECRET_KEYS.INFERMATICAI, secretId);
42
43 return apiKey ? ({
44 'Authorization': `Bearer ${apiKey}`,
45 }) : {};
46}
47
48/**
49 * Gets the headers for the DreamGen API.
50 * @param {import('./users.js').UserDirectoryList} directories User directories
51 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
52 * @returns {object} Headers for the request
53 */
54function getDreamGenHeaders(directories, secretId = null) {
55 const apiKey = readSecret(directories, SECRET_KEYS.DREAMGEN, secretId);
56
57 return apiKey ? ({
58 'Authorization': `Bearer ${apiKey}`,
59 }) : {};
60}
61
62/**
63 * Gets the headers for the OpenRouter API.
64 * @param {import('./users.js').UserDirectoryList} directories User directories
65 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
66 * @returns {object} Headers for the request
67 */
68function getOpenRouterHeaders(directories, secretId = null) {
69 const apiKey = readSecret(directories, SECRET_KEYS.OPENROUTER, secretId);
70 const baseHeaders = { ...OPENROUTER_HEADERS };
71
72 return apiKey ? Object.assign(baseHeaders, { 'Authorization': `Bearer ${apiKey}` }) : baseHeaders;
73}
74
75/**
76 * Gets the headers for the vLLM API.
77 * @param {import('./users.js').UserDirectoryList} directories User directories
78 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
79 * @returns {object} Headers for the request
80 */
81function getVllmHeaders(directories, secretId = null) {
82 const apiKey = readSecret(directories, SECRET_KEYS.VLLM, secretId);
83
84 return apiKey ? ({
85 'Authorization': `Bearer ${apiKey}`,
86 }) : {};
87}
88
89/**
90 * Gets the headers for the Aphrodite API.
91 * @param {import('./users.js').UserDirectoryList} directories User directories
92 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
93 * @returns {object} Headers for the request
94 */
95function getAphroditeHeaders(directories, secretId = null) {
96 const apiKey = readSecret(directories, SECRET_KEYS.APHRODITE, secretId);
97
98 return apiKey ? ({
99 'X-API-KEY': apiKey,
100 'Authorization': `Bearer ${apiKey}`,
101 }) : {};
102}
103
104/**
105 * Gets the headers for the Tabby API.
106 * @param {import('./users.js').UserDirectoryList} directories User directories
107 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
108 * @returns {object} Headers for the request
109 */
110function getTabbyHeaders(directories, secretId = null) {
111 const apiKey = readSecret(directories, SECRET_KEYS.TABBY, secretId);
112
113 return apiKey ? ({
114 'x-api-key': apiKey,
115 'Authorization': `Bearer ${apiKey}`,
116 }) : {};
117}
118
119/**
120 * Gets the headers for the LlamaCPP API.
121 * @param {import('./users.js').UserDirectoryList} directories User directories
122 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
123 * @returns {object} Headers for the request
124 */
125function getLlamaCppHeaders(directories, secretId = null) {
126 const apiKey = readSecret(directories, SECRET_KEYS.LLAMACPP, secretId);
127
128 return apiKey ? ({
129 'Authorization': `Bearer ${apiKey}`,
130 }) : {};
131}
132
133/**
134 * Gets the headers for the Ooba API.
135 * @param {import('./users.js').UserDirectoryList} directories User directories
136 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
137 * @returns {object} Headers for the request
138 */
139function getOobaHeaders(directories, secretId = null) {
140 const apiKey = readSecret(directories, SECRET_KEYS.OOBA, secretId);
141
142 return apiKey ? ({
143 'Authorization': `Bearer ${apiKey}`,
144 }) : {};
145}
146
147/**
148 * Gets the headers for the KoboldCpp API.
149 * @param {import('./users.js').UserDirectoryList} directories
150 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
151 * @returns {object} Headers for the request
152 */
153function getKoboldCppHeaders(directories, secretId = null) {
154 const apiKey = readSecret(directories, SECRET_KEYS.KOBOLDCPP, secretId);
155
156 return apiKey ? ({
157 'Authorization': `Bearer ${apiKey}`,
158 }) : {};
159}
160
161/**
162 * Gets the headers for the Featherless API.
163 * @param {import('./users.js').UserDirectoryList} directories
164 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
165 * @returns {object} Headers for the request
166 */
167function getFeatherlessHeaders(directories, secretId = null) {
168 const apiKey = readSecret(directories, SECRET_KEYS.FEATHERLESS, secretId);
169 const baseHeaders = { ...FEATHERLESS_HEADERS };
170
171 return apiKey ? Object.assign(baseHeaders, { 'Authorization': `Bearer ${apiKey}` }) : baseHeaders;
172}
173
174/**
175 * Gets the headers for the HuggingFace API.
176 * @param {import('./users.js').UserDirectoryList} directories
177 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
178 * @returns {object} Headers for the request
179 */
180function getHuggingFaceHeaders(directories, secretId = null) {
181 const apiKey = readSecret(directories, SECRET_KEYS.HUGGINGFACE, secretId);
182
183 return apiKey ? ({
184 'Authorization': `Bearer ${apiKey}`,
185 }) : {};
186}
187
188/**
189 * Gets the headers for the Generic text completion API.
190 * @param {import('./users.js').UserDirectoryList} directories
191 * @param {string|null} secretId Secret ID for the request (optional, used to determine which secret to use)
192 * @returns {object} Headers for the request
193 */
194function getGenericHeaders(directories, secretId = null) {
195 const apiKey = readSecret(directories, SECRET_KEYS.GENERIC, secretId);
196
197 return apiKey ? ({
198 'Authorization': `Bearer ${apiKey}`,
199 }) : {};
200}
201
202export function getOverrideHeaders(urlHost) {
203 const requestOverrides = getConfigValue('requestOverrides', []);
204 const overrideHeaders = requestOverrides?.find((e) => e.hosts?.includes(urlHost))?.headers;
205 if (overrideHeaders && urlHost) {
206 return overrideHeaders;
207 } else {
208 return {};
209 }
210}
211
212/**
213 * Sets additional headers for the request.
214 * @param {import('express').Request} request Original request body
215 * @param {object} args New request arguments
216 * @param {string|null} server API server for new request
217 */
218export function setAdditionalHeaders(request, args, server) {
219 setAdditionalHeadersByType(args.headers, request.body.api_type, server, request.user.directories, request.body.secret_id);
220}
221
222/**
223 *
224 * @param {object} requestHeaders Request headers
225 * @param {string} type API type
226 * @param {string|null} server API server for new request
227 * @param {import('./users.js').UserDirectoryList} directories User directories
228 * @param {string|null} secretId Secret ID for the request (optional, used for some API types to determine which secret to use)
229 */
230export function setAdditionalHeadersByType(requestHeaders, type, server, directories, secretId = null) {
231 const headerGetters = {
232 [TEXTGEN_TYPES.MANCER]: getMancerHeaders,
233 [TEXTGEN_TYPES.VLLM]: getVllmHeaders,
234 [TEXTGEN_TYPES.APHRODITE]: getAphroditeHeaders,
235 [TEXTGEN_TYPES.TABBY]: getTabbyHeaders,
236 [TEXTGEN_TYPES.TOGETHERAI]: getTogetherAIHeaders,
237 [TEXTGEN_TYPES.OOBA]: getOobaHeaders,
238 [TEXTGEN_TYPES.INFERMATICAI]: getInfermaticAIHeaders,
239 [TEXTGEN_TYPES.DREAMGEN]: getDreamGenHeaders,
240 [TEXTGEN_TYPES.OPENROUTER]: getOpenRouterHeaders,
241 [TEXTGEN_TYPES.KOBOLDCPP]: getKoboldCppHeaders,
242 [TEXTGEN_TYPES.LLAMACPP]: getLlamaCppHeaders,
243 [TEXTGEN_TYPES.FEATHERLESS]: getFeatherlessHeaders,
244 [TEXTGEN_TYPES.HUGGINGFACE]: getHuggingFaceHeaders,
245 [TEXTGEN_TYPES.GENERIC]: getGenericHeaders,
246 };
247
248 const getHeaders = headerGetters[type];
249 const headers = getHeaders ? getHeaders(directories, secretId) : {};
250
251 if (typeof server === 'string' && server.length > 0) {
252 try {
253 const url = new URL(server);
254 const overrideHeaders = getOverrideHeaders(url.host);
255
256 if (overrideHeaders && Object.keys(overrideHeaders).length > 0) {
257 Object.assign(headers, overrideHeaders);
258 }
259 } catch {
260 // Do nothing
261 }
262 }
263
264 Object.assign(requestHeaders, headers);
265}