Pass messages to SageMakerEndpoint

[jmnarloch]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

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",
    },
  },
});

Description

I followed the SageMakerEndpoint documentation to integrate it with my application and as the integration is working correctly I don't appear to find a way to get direct access to passed messages. The content handler appears to only accept the entire prompt input as one string with all messages combined. Is there no way to access the individual messages using the SageMakerEndpoint?

System Info

@langchain/community 0.2.9

[jmnarloch]

I realized that what is missing is the *Chat wrapper that can deal with handling of the messages for specific model. In the Python langchain library the Llama2Chat could be used to wrap the SageMakerEndpoint, but in the JS version I don't believe any wrapper like that exist out of the shelf.

[dosubot[bot]]

You are correct that the SageMakerEndpoint class in the LangChain JS library does not provide a direct way to access individual messages, as it handles the entire prompt input as a single string. However, you can manage individual messages by creating a custom handler or using utility functions to map stored messages to chat messages and vice versa.

Here is an example of how you can create a custom handler to manage individual messages:

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

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

class CustomHandler implements SageMakerLLMContentHandler {
  contentType = "application/json";
  accepts = "application/json";

  async transformInput(
    messages: { role: string; content: string }[],
    modelKwargs: Record<string, unknown>
  ): Promise<Uint8Array> {
    const payload = {
      inputs: [messages],
      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 CustomHandler();

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 messages = [
  { role: "user", content: "Hello, my name is John Doe." },
  { role: "user", content: "Tell me a joke about llamas." },
];

const response = await model.invoke(messages);
console.log(response);

This custom handler allows you to pass individual messages to the SageMakerEndpoint class.

Additionally, you can use utility functions from the @langchain/core/messages module to map stored messages to chat messages and vice versa:

import { mapStoredMessagesToChatMessages, mapChatMessagesToStoredMessages } from "@langchain/core/messages";

// Example usage of the utility functions
const storedMessages = /* some stored messages */;
const chatMessages = mapStoredMessagesToChatMessages(storedMessages);
const backToStoredMessages = mapChatMessagesToStoredMessages(chatMessages);

console.log({ chatMessages, backToStoredMessages });

These utility functions can help in managing and converting messages within your chat sessions ^[1][2].

To continue talking to Dosu, mention @dosu.