Add Sagemaker support (#401)

* work on sagemaker support

* fix sagemaker integration

* remove unnecessary deps

* fix default endpoint

* remove unneeded deps, fixed types

* Use conditional validation for endpoints

This was needed because the discriminated union couldn't handle the legacy case where `host` is undefined.

* add note in readme about aws sagemaker

README.md

@@ -198,6 +198,24 @@ You can then add the generated information and the `authorization` parameter to
 
 ```
 
+### Amazon SageMaker
+
+You can also specify your Amazon SageMaker instance as an endpoint for chat-ui. The config goes like this:
+
+```
+"endpoints": [
+    {
+      "host" : "sagemaker",
+      "url": "", // your aws sagemaker url here
+      "accessKey": "",
+      "secretKey" : "",
+      "sessionToken": "", // optional
+      "weight": 1
+    }
+```
+
+You can get the `accessKey` and `secretKey` from your AWS user, under programmatic access.
+
 #### Client Certificate Authentication (mTLS)
 
 Custom endpoints may require client certificate authentication, depending on how you configure them. To enable mTLS between Chat UI and your custom endpoint, you will need to set the `USE_CLIENT_CERTIFICATE` to `true`, and add the `CERT_PATH` and `KEY_PATH` parameters to your `.env.local`. These parameters should point to the location of the certificate and key files on your local machine. The certificate and key files should be in PEM format. The key file can be encrypted with a passphrase, in which case you will also need to add the `CLIENT_KEY_PASSWORD` parameter to your `.env.local`.

package-lock.json

@@ -11,6 +11,7 @@
 				"@huggingface/hub": "^0.5.1",
 				"@huggingface/inference": "^2.2.0",
 				"autoprefixer": "^10.4.14",
+				"aws4fetch": "^1.0.17",
 				"date-fns": "^2.29.3",
 				"dotenv": "^16.0.3",
 				"highlight.js": "^11.7.0",
@@ -1465,6 +1466,11 @@
 				"postcss": "^8.1.0"
 			}
 		},
+		"node_modules/aws4fetch": {
+			"version": "1.0.17",
+			"resolved": "https://registry.npmjs.org/aws4fetch/-/aws4fetch-1.0.17.tgz",
+			"integrity": "sha512-4IbOvsxqxeOSxI4oA+8xEO8SzBMVlzbSTgGy/EF83rHnQ/aKtP6Sc6YV/k0oiW0mqrcxuThlbDosnvetGOuO+g=="
+		},
 		"node_modules/balanced-match": {
 			"version": "1.0.2",
 			"resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz",

package.json

@@ -43,6 +43,7 @@
 		"@huggingface/hub": "^0.5.1",
 		"@huggingface/inference": "^2.2.0",
 		"autoprefixer": "^10.4.14",
+		"aws4fetch": "^1.0.17",
 		"date-fns": "^2.29.3",
 		"dotenv": "^16.0.3",
 		"highlight.js": "^11.7.0",

src/lib/server/generateFromDefaultEndpoint.ts

@@ -1,9 +1,9 @@
 import { defaultModel } from "$lib/server/models";
 import { modelEndpoint } from "./modelEndpoint";
-import { textGeneration } from "@huggingface/inference";
 import { trimSuffix } from "$lib/utils/trimSuffix";
 import { trimPrefix } from "$lib/utils/trimPrefix";
 import { PUBLIC_SEP_TOKEN } from "$lib/constants/publicSepToken";
+import { AwsClient } from "aws4fetch";
 
 interface Parameters {
 	temperature: number;
@@ -21,24 +21,76 @@ export async function generateFromDefaultEndpoint(
 		return_full_text: false,
 	};
 
-	const endpoint = modelEndpoint(defaultModel);
-	let { generated_text } = await textGeneration(
-		{
-			model: endpoint.url,
+	const randomEndpoint = modelEndpoint(defaultModel);
+
+	const abortController = new AbortController();
+
+	let resp: Response;
+
+	if (randomEndpoint.host === "sagemaker") {
+		const requestParams = JSON.stringify({
+			...newParameters,
 			inputs: prompt,
-			parameters: newParameters,
-		},
-		{
-			fetch: (url, options) =>
-				fetch(url, {
-					...options,
-					headers: { ...options?.headers, Authorization: endpoint.authorization },
-				}),
-		}
-	);
+		});
+
+		const aws = new AwsClient({
+			accessKeyId: randomEndpoint.accessKey,
+			secretAccessKey: randomEndpoint.secretKey,
+			sessionToken: randomEndpoint.sessionToken,
+			service: "sagemaker",
+		});
+
+		resp = await aws.fetch(randomEndpoint.url, {
+			method: "POST",
+			body: requestParams,
+			signal: abortController.signal,
+			headers: {
+				"Content-Type": "application/json",
+			},
+		});
+	} else {
+		resp = await fetch(randomEndpoint.url, {
+			headers: {
+				"Content-Type": "application/json",
+				Authorization: randomEndpoint.authorization,
+			},
+			method: "POST",
+			body: JSON.stringify({
+				...newParameters,
+				inputs: prompt,
+			}),
+			signal: abortController.signal,
+		});
+	}
+
+	if (!resp.ok) {
+		throw new Error(await resp.text());
+	}
+
+	if (!resp.body) {
+		throw new Error("Response body is empty");
+	}
+
+	const decoder = new TextDecoder();
+	const reader = resp.body.getReader();
+
+	let isDone = false;
+	let result = "";
+
+	while (!isDone) {
+		const { done, value } = await reader.read();
+
+		isDone = done;
+		result += decoder.decode(value, { stream: true }); // Convert current chunk to text
+	}
+
+	// Close the reader when done
+	reader.releaseLock();
+
+	const results = await JSON.parse(result);
 
-	generated_text = trimSuffix(
-		trimPrefix(generated_text, "<|startoftext|>"),
+	let generated_text = trimSuffix(
+		trimPrefix(trimPrefix(results[0].generated_text, "<|startoftext|>"), prompt),
 		PUBLIC_SEP_TOKEN
 	).trimEnd();
 

src/lib/server/modelEndpoint.ts

@@ -9,7 +9,7 @@ import {
 	REJECT_UNAUTHORIZED,
 } from "$env/static/private";
 import { sum } from "$lib/utils/sum";
-import type { BackendModel } from "./models";
+import type { BackendModel, Endpoint } from "./models";
 
 import { loadClientCertificates } from "$lib/utils/loadClientCerts";
 
@@ -26,13 +26,10 @@ if (USE_CLIENT_CERTIFICATE === "true") {
 /**
  * Find a random load-balanced endpoint
  */
-export function modelEndpoint(model: BackendModel): {
-	url: string;
-	authorization: string;
-	weight: number;
-} {
+export function modelEndpoint(model: BackendModel): Endpoint {
 	if (!model.endpoints) {
 		return {
+			host: "tgi",
 			url: `${HF_API_ROOT}/${model.name}`,
 			authorization: `Bearer ${HF_ACCESS_TOKEN}`,
 			weight: 1,

src/lib/server/models.ts

@@ -1,6 +1,38 @@
 import { HF_ACCESS_TOKEN, MODELS, OLD_MODELS } from "$env/static/private";
 import { z } from "zod";
 
+const sagemakerEndpoint = z.object({
+	host: z.literal("sagemaker"),
+	url: z.string().url(),
+	accessKey: z.string().min(1),
+	secretKey: z.string().min(1),
+	sessionToken: z.string().optional(),
+});
+
+const tgiEndpoint = z.object({
+	host: z.union([z.literal("tgi"), z.undefined()]),
+	url: z.string().url(),
+	authorization: z.string().min(1).default(`Bearer ${HF_ACCESS_TOKEN}`),
+});
+
+const commonEndpoint = z.object({
+	weight: z.number().int().positive().default(1),
+});
+
+const endpoint = z.lazy(() =>
+	z.union([sagemakerEndpoint.merge(commonEndpoint), tgiEndpoint.merge(commonEndpoint)])
+);
+
+const combinedEndpoint = endpoint.transform((data) => {
+	if (data.host === "tgi" || data.host === undefined) {
+		return tgiEndpoint.merge(commonEndpoint).parse(data);
+	} else if (data.host === "sagemaker") {
+		return sagemakerEndpoint.merge(commonEndpoint).parse(data);
+	} else {
+		throw new Error(`Invalid host: ${data.host}`);
+	}
+});
+
 const modelsRaw = z
 	.array(
 		z.object({
@@ -29,15 +61,7 @@ const modelsRaw = z
 					})
 				)
 				.optional(),
-			endpoints: z
-				.array(
-					z.object({
-						url: z.string().url(),
-						authorization: z.string().min(1).default(`Bearer ${HF_ACCESS_TOKEN}`),
-						weight: z.number().int().positive().default(1),
-					})
-				)
-				.optional(),
+			endpoints: z.array(combinedEndpoint).optional(),
 			parameters: z
 				.object({
 					temperature: z.number().min(0).max(1),
@@ -77,6 +101,7 @@ export const oldModels = OLD_MODELS
 	: [];
 
 export type BackendModel = (typeof models)[0];
+export type Endpoint = z.infer<typeof endpoint>;
 
 export const defaultModel = models[0];
 

src/routes/conversation/[id]/+server.ts

@@ -16,6 +16,7 @@ import type { TextGenerationStreamOutput } from "@huggingface/inference";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";
 import { z } from "zod";
+import { AwsClient } from "aws4fetch";
 
 export async function POST({ request, fetch, locals, params }) {
 	const id = z.string().parse(params.id);
@@ -101,18 +102,42 @@ export async function POST({ request, fetch, locals, params }) {
 
 	const abortController = new AbortController();
 
-	const resp = await fetch(randomEndpoint.url, {
-		headers: {
-			"Content-Type": request.headers.get("Content-Type") ?? "application/json",
-			Authorization: randomEndpoint.authorization,
-		},
-		method: "POST",
-		body: JSON.stringify({
+	let resp: Response;
+	if (randomEndpoint.host === "sagemaker") {
+		const requestParams = JSON.stringify({
 			...json,
 			inputs: prompt,
-		}),
-		signal: abortController.signal,
-	});
+		});
+
+		const aws = new AwsClient({
+			accessKeyId: randomEndpoint.accessKey,
+			secretAccessKey: randomEndpoint.secretKey,
+			sessionToken: randomEndpoint.sessionToken,
+			service: "sagemaker",
+		});
+
+		resp = await aws.fetch(randomEndpoint.url, {
+			method: "POST",
+			body: requestParams,
+			signal: abortController.signal,
+			headers: {
+				"Content-Type": "application/json",
+			},
+		});
+	} else {
+		resp = await fetch(randomEndpoint.url, {
+			headers: {
+				"Content-Type": request.headers.get("Content-Type") ?? "application/json",
+				Authorization: randomEndpoint.authorization,
+			},
+			method: "POST",
+			body: JSON.stringify({
+				...json,
+				inputs: prompt,
+			}),
+			signal: abortController.signal,
+		});
+	}
 
 	if (!resp.body) {
 		throw new Error("Response body is empty");