355 lines
14 KiB
JavaScript
355 lines
14 KiB
JavaScript
import { AIMessage, AIMessageChunk, } from "@langchain/core/messages";
|
||
import { ChatGenerationChunk } from "@langchain/core/outputs";
|
||
import { _makeMessageChunkFromAnthropicEvent, extractToken, extractToolCallChunk, extractToolUseContent, extractToolCalls, formatMessagesForAnthropic, } from "./anthropic.js";
|
||
/**
|
||
* format messages for Cohere Command-R and CommandR+ via AWS Bedrock.
|
||
*
|
||
* @param messages messages The base messages to format as a prompt.
|
||
*
|
||
* @returns The formatted prompt for Cohere.
|
||
*
|
||
* `system`: user system prompts. Overrides the default preamble for search query generation. Has no effect on tool use generations.\
|
||
* `message`: (Required) Text input for the model to respond to.\
|
||
* `chatHistory`: A list of previous messages between the user and the model, meant to give the model conversational context for responding to the user's message.\
|
||
* The following are required fields.
|
||
* - `role` - The role for the message. Valid values are USER or CHATBOT.\
|
||
* - `message` – Text contents of the message.\
|
||
*
|
||
* The following is example JSON for the chat_history field.\
|
||
* "chat_history": [
|
||
* {"role": "USER", "message": "Who discovered gravity?"},
|
||
* {"role": "CHATBOT", "message": "The man who is widely credited with discovering gravity is Sir Isaac Newton"}]\
|
||
*
|
||
* docs: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html
|
||
*/
|
||
function formatMessagesForCohere(messages) {
|
||
const systemMessages = messages.filter((system) => system._getType() === "system");
|
||
const system = systemMessages
|
||
.filter((m) => typeof m.content === "string")
|
||
.map((m) => m.content)
|
||
.join("\n\n");
|
||
const conversationMessages = messages.filter((message) => message._getType() !== "system");
|
||
const questionContent = conversationMessages.slice(-1);
|
||
if (!questionContent.length || questionContent[0]._getType() !== "human") {
|
||
throw new Error("question message content must be a human message.");
|
||
}
|
||
if (typeof questionContent[0].content !== "string") {
|
||
throw new Error("question message content must be a string.");
|
||
}
|
||
const formattedMessage = questionContent[0].content;
|
||
const formattedChatHistories = conversationMessages
|
||
.slice(0, -1)
|
||
.map((message) => {
|
||
let role;
|
||
switch (message._getType()) {
|
||
case "human":
|
||
role = "USER";
|
||
break;
|
||
case "ai":
|
||
role = "CHATBOT";
|
||
break;
|
||
case "system":
|
||
throw new Error("chat_history can not include system prompts.");
|
||
default:
|
||
throw new Error(`Message type "${message._getType()}" is not supported.`);
|
||
}
|
||
if (typeof message.content !== "string") {
|
||
throw new Error("message content must be a string.");
|
||
}
|
||
return {
|
||
role,
|
||
message: message.content,
|
||
};
|
||
});
|
||
return {
|
||
chatHistory: formattedChatHistories,
|
||
message: formattedMessage,
|
||
system,
|
||
};
|
||
}
|
||
/**
|
||
* A helper class used within the `Bedrock` class. It is responsible for
|
||
* preparing the input and output for the Bedrock service. It formats the
|
||
* input prompt based on the provider (e.g., "anthropic", "ai21",
|
||
* "amazon") and extracts the generated text from the service response.
|
||
*/
|
||
export class BedrockLLMInputOutputAdapter {
|
||
/** Adapter class to prepare the inputs from Langchain to a format
|
||
that LLM model expects. Also, provides a helper function to extract
|
||
the generated text from the model response. */
|
||
static prepareInput(provider, prompt, maxTokens = 50, temperature = 0, stopSequences = undefined, modelKwargs = {}, bedrockMethod = "invoke", guardrailConfig = undefined) {
|
||
const inputBody = {};
|
||
if (provider === "anthropic") {
|
||
inputBody.prompt = prompt;
|
||
inputBody.max_tokens_to_sample = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
inputBody.stop_sequences = stopSequences;
|
||
}
|
||
else if (provider === "ai21") {
|
||
inputBody.prompt = prompt;
|
||
inputBody.maxTokens = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
inputBody.stopSequences = stopSequences;
|
||
}
|
||
else if (provider === "meta") {
|
||
inputBody.prompt = prompt;
|
||
inputBody.max_gen_len = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
}
|
||
else if (provider === "amazon") {
|
||
inputBody.inputText = prompt;
|
||
inputBody.textGenerationConfig = {
|
||
maxTokenCount: maxTokens,
|
||
temperature,
|
||
};
|
||
}
|
||
else if (provider === "cohere") {
|
||
inputBody.prompt = prompt;
|
||
inputBody.max_tokens = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
inputBody.stop_sequences = stopSequences;
|
||
if (bedrockMethod === "invoke-with-response-stream") {
|
||
inputBody.stream = true;
|
||
}
|
||
}
|
||
else if (provider === "mistral") {
|
||
inputBody.prompt = prompt;
|
||
inputBody.max_tokens = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
inputBody.stop = stopSequences;
|
||
}
|
||
if (guardrailConfig &&
|
||
guardrailConfig.tagSuffix &&
|
||
guardrailConfig.streamProcessingMode) {
|
||
inputBody["amazon-bedrock-guardrailConfig"] = guardrailConfig;
|
||
}
|
||
return { ...inputBody, ...modelKwargs };
|
||
}
|
||
static prepareMessagesInput(provider, messages, maxTokens = 1024, temperature = 0, stopSequences = undefined, modelKwargs = {}, guardrailConfig = undefined, tools = []) {
|
||
const inputBody = {};
|
||
if (provider === "anthropic") {
|
||
const { system, messages: formattedMessages } = formatMessagesForAnthropic(messages);
|
||
if (system !== undefined) {
|
||
inputBody.system = system;
|
||
}
|
||
inputBody.anthropic_version = "bedrock-2023-05-31";
|
||
inputBody.messages = formattedMessages;
|
||
inputBody.max_tokens = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
inputBody.stop_sequences = stopSequences;
|
||
if (tools.length > 0) {
|
||
inputBody.tools = tools;
|
||
}
|
||
return { ...inputBody, ...modelKwargs };
|
||
}
|
||
else if (provider === "cohere") {
|
||
const { system, message: formattedMessage, chatHistory: formattedChatHistories, } = formatMessagesForCohere(messages);
|
||
if (system !== undefined && system.length > 0) {
|
||
inputBody.preamble = system;
|
||
}
|
||
inputBody.message = formattedMessage;
|
||
inputBody.chat_history = formattedChatHistories;
|
||
inputBody.max_tokens = maxTokens;
|
||
inputBody.temperature = temperature;
|
||
inputBody.stop_sequences = stopSequences;
|
||
}
|
||
else {
|
||
throw new Error("The messages API is currently only supported by Anthropic or Cohere");
|
||
}
|
||
if (guardrailConfig &&
|
||
guardrailConfig.tagSuffix &&
|
||
guardrailConfig.streamProcessingMode) {
|
||
inputBody["amazon-bedrock-guardrailConfig"] = guardrailConfig;
|
||
}
|
||
return { ...inputBody, ...modelKwargs };
|
||
}
|
||
/**
|
||
* Extracts the generated text from the service response.
|
||
* @param provider The provider name.
|
||
* @param responseBody The response body from the service.
|
||
* @returns The generated text.
|
||
*/
|
||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||
static prepareOutput(provider, responseBody) {
|
||
if (provider === "anthropic") {
|
||
return responseBody.completion;
|
||
}
|
||
else if (provider === "ai21") {
|
||
return responseBody?.completions?.[0]?.data?.text ?? "";
|
||
}
|
||
else if (provider === "cohere") {
|
||
return responseBody?.generations?.[0]?.text ?? responseBody?.text ?? "";
|
||
}
|
||
else if (provider === "meta") {
|
||
return responseBody.generation;
|
||
}
|
||
else if (provider === "mistral") {
|
||
return responseBody?.outputs?.[0]?.text;
|
||
}
|
||
// I haven't been able to get a response with more than one result in it.
|
||
return responseBody.results?.[0]?.outputText;
|
||
}
|
||
static prepareMessagesOutput(provider,
|
||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||
response, fields) {
|
||
const responseBody = response ?? {};
|
||
if (provider === "anthropic") {
|
||
if (responseBody.type === "message") {
|
||
return parseMessage(responseBody);
|
||
}
|
||
else if (responseBody.type === "message_start") {
|
||
return parseMessage(responseBody.message, true);
|
||
}
|
||
const chunk = _makeMessageChunkFromAnthropicEvent(response, {
|
||
coerceContentToString: fields?.coerceContentToString,
|
||
});
|
||
if (!chunk)
|
||
return undefined;
|
||
const newToolCallChunk = extractToolCallChunk(chunk);
|
||
let toolUseContent;
|
||
const extractedContent = extractToolUseContent(chunk, undefined);
|
||
if (extractedContent) {
|
||
toolUseContent = extractedContent.toolUseContent;
|
||
}
|
||
// Filter partial `tool_use` content, and only add `tool_use` chunks if complete JSON available.
|
||
const chunkContent = Array.isArray(chunk.content)
|
||
? chunk.content.filter((c) => c.type !== "tool_use")
|
||
: chunk.content;
|
||
if (Array.isArray(chunkContent) && toolUseContent) {
|
||
chunkContent.push(toolUseContent);
|
||
}
|
||
// Extract the text content token for text field and runManager.
|
||
const token = extractToken(chunk);
|
||
return new ChatGenerationChunk({
|
||
message: new AIMessageChunk({
|
||
content: chunkContent,
|
||
additional_kwargs: chunk.additional_kwargs,
|
||
tool_call_chunks: newToolCallChunk ? [newToolCallChunk] : undefined,
|
||
usage_metadata: chunk.usage_metadata,
|
||
response_metadata: chunk.response_metadata,
|
||
}),
|
||
// Backwards compatibility
|
||
generationInfo: { ...chunk.response_metadata },
|
||
text: token ?? "",
|
||
});
|
||
}
|
||
else if (provider === "cohere") {
|
||
if (responseBody.event_type === "stream-start") {
|
||
return parseMessageCohere(responseBody.message, true);
|
||
}
|
||
else if (responseBody.event_type === "text-generation" &&
|
||
typeof responseBody?.text === "string") {
|
||
return new ChatGenerationChunk({
|
||
message: new AIMessageChunk({
|
||
content: responseBody.text,
|
||
}),
|
||
text: responseBody.text,
|
||
});
|
||
}
|
||
else if (responseBody.event_type === "search-queries-generation") {
|
||
return parseMessageCohere(responseBody);
|
||
}
|
||
else if (responseBody.event_type === "stream-end" &&
|
||
responseBody.response !== undefined &&
|
||
responseBody["amazon-bedrock-invocationMetrics"] !== undefined) {
|
||
return new ChatGenerationChunk({
|
||
message: new AIMessageChunk({ content: "" }),
|
||
text: "",
|
||
generationInfo: {
|
||
response: responseBody.response,
|
||
"amazon-bedrock-invocationMetrics": responseBody["amazon-bedrock-invocationMetrics"],
|
||
},
|
||
});
|
||
}
|
||
else {
|
||
if (responseBody.finish_reason === "COMPLETE" ||
|
||
responseBody.finish_reason === "MAX_TOKENS") {
|
||
return parseMessageCohere(responseBody);
|
||
}
|
||
else {
|
||
return undefined;
|
||
}
|
||
}
|
||
}
|
||
else {
|
||
throw new Error("The messages API is currently only supported by Anthropic or Cohere.");
|
||
}
|
||
}
|
||
}
|
||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||
function parseMessage(responseBody, asChunk) {
|
||
const { content, id, ...generationInfo } = responseBody;
|
||
let parsedContent;
|
||
if (Array.isArray(content) &&
|
||
content.length === 1 &&
|
||
content[0].type === "text") {
|
||
parsedContent = content[0].text;
|
||
}
|
||
else if (Array.isArray(content) && content.length === 0) {
|
||
parsedContent = "";
|
||
}
|
||
else {
|
||
parsedContent = content;
|
||
}
|
||
if (asChunk) {
|
||
return new ChatGenerationChunk({
|
||
message: new AIMessageChunk({
|
||
content: parsedContent,
|
||
additional_kwargs: { id },
|
||
}),
|
||
text: typeof parsedContent === "string" ? parsedContent : "",
|
||
generationInfo,
|
||
});
|
||
}
|
||
else {
|
||
// TODO: we are throwing away here the text response, as the interface of this method returns only one
|
||
const toolCalls = extractToolCalls(responseBody.content);
|
||
if (toolCalls.length > 0) {
|
||
return {
|
||
message: new AIMessage({
|
||
content: "",
|
||
additional_kwargs: { id },
|
||
tool_calls: toolCalls,
|
||
}),
|
||
text: typeof parsedContent === "string" ? parsedContent : "",
|
||
generationInfo,
|
||
};
|
||
}
|
||
return {
|
||
message: new AIMessage({
|
||
content: parsedContent,
|
||
additional_kwargs: { id },
|
||
tool_calls: toolCalls,
|
||
}),
|
||
text: typeof parsedContent === "string" ? parsedContent : "",
|
||
generationInfo,
|
||
};
|
||
}
|
||
}
|
||
function parseMessageCohere(
|
||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||
responseBody, asChunk) {
|
||
const { text, ...generationInfo } = responseBody;
|
||
let parsedContent = text;
|
||
if (typeof text !== "string") {
|
||
parsedContent = "";
|
||
}
|
||
if (asChunk) {
|
||
return new ChatGenerationChunk({
|
||
message: new AIMessageChunk({
|
||
content: parsedContent,
|
||
}),
|
||
text: parsedContent,
|
||
generationInfo,
|
||
});
|
||
}
|
||
else {
|
||
return {
|
||
message: new AIMessage({
|
||
content: parsedContent,
|
||
}),
|
||
text: parsedContent,
|
||
generationInfo,
|
||
};
|
||
}
|
||
}
|