Skip to main content

AWS SageMakerEndpoint

LangChain.js supports integration with AWS SageMaker-hosted endpoints. Check Amazon SageMaker JumpStart for a list of available models, and how to deploy your own.

Setup

You'll need to install the official SageMaker SDK as a peer dependency:

npm install @aws-sdk/client-sagemaker-runtime
npm install @langchain/community @langchain/core

Usage

import {
SageMakerEndpoint,
SageMakerLLMContentHandler,
} from "@langchain/community/llms/sagemaker_endpoint";

interface ResponseJsonInterface {
generation: {
content: string;
};
}

// Custom for whatever model you'll be using
class LLama213BHandler implements SageMakerLLMContentHandler {
contentType = "application/json";

accepts = "application/json";

async transformInput(
prompt: string,
modelKwargs: Record<string, unknown>
): Promise<Uint8Array> {
const payload = {
inputs: [[{ role: "user", content: prompt }]],
parameters: modelKwargs,
};

const stringifiedPayload = JSON.stringify(payload);

return new TextEncoder().encode(stringifiedPayload);
}

async transformOutput(output: Uint8Array): Promise<string> {
const response_json = JSON.parse(
new TextDecoder("utf-8").decode(output)
) as ResponseJsonInterface[];
const content = response_json[0]?.generation.content ?? "";
return content;
}
}

const contentHandler = new LLama213BHandler();

const model = new SageMakerEndpoint({
endpointName: "aws-llama-2-13b-chat",
modelKwargs: {
temperature: 0.5,
max_new_tokens: 700,
top_p: 0.9,
},
endpointKwargs: {
CustomAttributes: "accept_eula=true",
},
contentHandler,
clientOptions: {
region: "YOUR AWS ENDPOINT REGION",
credentials: {
accessKeyId: "YOUR AWS ACCESS ID",
secretAccessKey: "YOUR AWS SECRET ACCESS KEY",
},
},
});

const res = await model.invoke(
"Hello, my name is John Doe, tell me a joke about llamas "
);

console.log(res);

/*
[
{
content: "Hello, John Doe! Here's a llama joke for you:
Why did the llama become a gardener?
Because it was great at llama-scaping!"
}
]
*/

API Reference:


Was this page helpful?


You can also leave detailed feedback on GitHub.