Add support for passing an API key or any other custom token in the authorization header (#579)

* Add support for passing an API key or any other custom token in the authorization header

* Make linter happy

* Fix README as per linter suggestions

* Refactor endpoints to actually parse zod config

* Remove top level env var and simplify header addition

* Skip section on API key or other, remove obsolete comment in endpointTgi.ts and remote CUSTOM_AUTHORIZATION_TOKEN from .env

---------

Co-authored-by: Nathan Sarrazin <[email protected]>

README.md

@@ -397,6 +397,8 @@ You can then add the generated information and the `authorization` parameter to
 ]
 ```
 
+Please note that if `HF_ACCESS_TOKEN` is also set or not empty, it will take precedence.
+
 #### Models hosted on multiple custom endpoints
 
 If the model being hosted will be available on multiple servers/instances add the `weight` parameter to your `.env.local`. The `weight` will be used to determine the probability of requesting a particular endpoint.

src/lib/server/endpoints/aws/endpointAws.ts

@@ -15,15 +15,9 @@ export const endpointAwsParametersSchema = z.object({
 	region: z.string().optional(),
 });
 
-export async function endpointAws({
-	url,
-	accessKey,
-	secretKey,
-	sessionToken,
-	model,
-	region,
-	service,
-}: z.infer<typeof endpointAwsParametersSchema>): Promise<Endpoint> {
+export async function endpointAws(
+	input: z.input<typeof endpointAwsParametersSchema>
+): Promise<Endpoint> {
 	let AwsClient;
 	try {
 		AwsClient = (await import("aws4fetch")).AwsClient;
@@ -31,6 +25,9 @@ export async function endpointAws({
 		throw new Error("Failed to import aws4fetch");
 	}
 
+	const { url, accessKey, secretKey, sessionToken, model, region, service } =
+		endpointAwsParametersSchema.parse(input);
+
 	const aws = new AwsClient({
 		accessKeyId: accessKey,
 		secretAccessKey: secretKey,

src/lib/server/endpoints/llamacpp/endpointLlamacpp.ts

@@ -12,10 +12,10 @@ export const endpointLlamacppParametersSchema = z.object({
 	accessToken: z.string().min(1).default(HF_ACCESS_TOKEN),
 });
 
-export function endpointLlamacpp({
-	url,
-	model,
-}: z.infer<typeof endpointLlamacppParametersSchema>): Endpoint {
+export function endpointLlamacpp(
+	input: z.input<typeof endpointLlamacppParametersSchema>
+): Endpoint {
+	const { url, model } = endpointLlamacppParametersSchema.parse(input);
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,

src/lib/server/endpoints/ollama/endpointOllama.ts

@@ -11,11 +11,9 @@ export const endpointOllamaParametersSchema = z.object({
 	ollamaName: z.string().min(1).optional(),
 });
 
-export function endpointOllama({
-	url,
-	model,
-	ollamaName,
-}: z.infer<typeof endpointOllamaParametersSchema>): Endpoint {
+export function endpointOllama(input: z.input<typeof endpointOllamaParametersSchema>): Endpoint {
+	const { url, model, ollamaName } = endpointOllamaParametersSchema.parse(input);
+
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,

src/lib/server/endpoints/openai/endpointOai.ts

@@ -16,12 +16,10 @@ export const endpointOAIParametersSchema = z.object({
 		.default("chat_completions"),
 });
 
-export async function endpointOai({
-	baseURL,
-	apiKey,
-	completion,
-	model,
-}: z.infer<typeof endpointOAIParametersSchema>): Promise<Endpoint> {
+export async function endpointOai(
+	input: z.input<typeof endpointOAIParametersSchema>
+): Promise<Endpoint> {
+	const { baseURL, apiKey, completion, model } = endpointOAIParametersSchema.parse(input);
 	let OpenAI;
 	try {
 		OpenAI = (await import("openai")).OpenAI;

src/lib/server/endpoints/tgi/endpointTgi.ts

@@ -10,13 +10,11 @@ export const endpointTgiParametersSchema = z.object({
 	type: z.literal("tgi"),
 	url: z.string().url(),
 	accessToken: z.string().default(HF_ACCESS_TOKEN),
+	authorization: z.string().optional(),
 });
 
-export function endpointTgi({
-	url,
-	accessToken,
-	model,
-}: z.infer<typeof endpointTgiParametersSchema>): Endpoint {
+export function endpointTgi(input: z.input<typeof endpointTgiParametersSchema>): Endpoint {
+	const { url, accessToken, model, authorization } = endpointTgiParametersSchema.parse(input);
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,
@@ -33,7 +31,19 @@ export function endpointTgi({
 				inputs: prompt,
 				accessToken,
 			},
-			{ use_cache: false }
+			{
+				use_cache: false,
+				fetch: async (endpointUrl, info) => {
+					if (info && authorization && !accessToken) {
+						// Set authorization header if it is defined and HF_ACCESS_TOKEN is empty
+						info.headers = {
+							...info.headers,
+							Authorization: authorization,
+						};
+					}
+					return fetch(endpointUrl, info);
+				},
+			}
 		);
 	};
 }