ChatGPT-Next-Web/app/api/bedrock/models.ts

281 lines
8.2 KiB
TypeScript

import {
Message,
validateMessageOrder,
processDocumentContent,
BedrockTextBlock,
BedrockImageBlock,
BedrockDocumentBlock,
} from "./utils";
export interface ConverseRequest {
modelId: string;
messages: Message[];
inferenceConfig?: {
maxTokens?: number;
temperature?: number;
topP?: number;
};
system?: string;
tools?: Array<{
type: "function";
function: {
name: string;
description: string;
parameters: {
type: string;
properties: Record<string, any>;
required: string[];
};
};
}>;
}
interface ContentItem {
type: string;
text?: string;
image_url?: {
url: string;
};
document?: {
format: string;
name: string;
source: {
bytes: string;
};
};
}
type ProcessedContent =
| ContentItem
| BedrockTextBlock
| BedrockImageBlock
| BedrockDocumentBlock
| {
type: string;
source: { type: string; media_type: string; data: string };
};
// Helper function to format request body based on model type
export function formatRequestBody(request: ConverseRequest) {
const baseModel = request.modelId;
const messages = validateMessageOrder(request.messages).map((msg) => ({
role: msg.role,
content: Array.isArray(msg.content)
? msg.content.map((item: ContentItem) => {
if (item.type === "image_url" && item.image_url?.url) {
// If it's a base64 image URL
const base64Match = item.image_url.url.match(
/^data:image\/([a-zA-Z]*);base64,([^"]*)$/,
);
if (base64Match) {
return {
type: "image",
source: {
type: "base64",
media_type: `image/${base64Match[1]}`,
data: base64Match[2],
},
};
}
// If it's not a base64 URL, return as is
return item;
}
if ("document" in item) {
try {
return processDocumentContent(item);
} catch (error) {
console.error("Error processing document:", error);
return {
type: "text",
text: `[Document: ${item.document?.name || "Unknown"}]`,
};
}
}
return { type: "text", text: item.text };
})
: [{ type: "text", text: msg.content }],
}));
const systemPrompt = request.system
? [{ type: "text", text: request.system }]
: undefined;
const baseConfig = {
max_tokens: request.inferenceConfig?.maxTokens || 2048,
temperature: request.inferenceConfig?.temperature || 0.7,
top_p: request.inferenceConfig?.topP || 0.9,
};
if (baseModel.startsWith("anthropic.claude")) {
return {
messages,
system: systemPrompt,
anthropic_version: "bedrock-2023-05-31",
...baseConfig,
...(request.tools && { tools: request.tools }),
};
} else if (
baseModel.startsWith("meta.llama") ||
baseModel.startsWith("mistral.")
) {
return {
messages: messages.map((m) => ({
role: m.role,
content: Array.isArray(m.content)
? m.content.map((c: ProcessedContent) => {
if ("text" in c) return { type: "text", text: c.text || "" };
if ("image_url" in c)
return {
type: "text",
text: `[Image: ${c.image_url?.url || "URL not provided"}]`,
};
if ("document" in c)
return {
type: "text",
text: `[Document: ${c.document?.name || "Unknown"}]`,
};
return { type: "text", text: "" };
})
: [{ type: "text", text: m.content }],
})),
...baseConfig,
stop_sequences: ["\n\nHuman:", "\n\nAssistant:"],
};
} else if (baseModel.startsWith("amazon.titan")) {
const formattedText = messages.map((m) => ({
role: m.role,
content: [
{
type: "text",
text: `${m.role === "user" ? "Human" : "Assistant"}: ${
Array.isArray(m.content)
? m.content
.map((c: ProcessedContent) => {
if ("text" in c) return c.text || "";
if ("image_url" in c)
return `[Image: ${
c.image_url?.url || "URL not provided"
}]`;
if ("document" in c)
return `[Document: ${c.document?.name || "Unknown"}]`;
return "";
})
.join("")
: m.content
}`,
},
],
}));
return {
messages: formattedText,
textGenerationConfig: {
maxTokenCount: baseConfig.max_tokens,
temperature: baseConfig.temperature,
topP: baseConfig.top_p,
stopSequences: ["Human:", "Assistant:"],
},
};
}
throw new Error(`Unsupported model: ${baseModel}`);
}
// Helper function to parse and format response based on model type
export function parseModelResponse(responseBody: string, modelId: string): any {
const baseModel = modelId;
try {
const response = JSON.parse(responseBody);
// Common response format for all models
const formatResponse = (content: string | any[]) => ({
role: "assistant",
content: Array.isArray(content)
? content.map((item) => {
if (typeof item === "string") {
return { type: "text", text: item };
}
// Handle different content types
if ("text" in item) {
return { type: "text", text: item.text || "" };
}
if ("image" in item) {
return {
type: "image_url",
image_url: {
url: `data:image/${
item.source?.media_type || "image/png"
};base64,${item.source?.data || ""}`,
},
};
}
// Document responses are converted to text
if ("document" in item) {
return {
type: "text",
text: `[Document Content]\n${item.text || ""}`,
};
}
return { type: "text", text: item.text || "" };
})
: [{ type: "text", text: content }],
stop_reason: response.stop_reason || response.stopReason || "end_turn",
usage: response.usage || {
input_tokens: 0,
output_tokens: 0,
total_tokens: 0,
},
});
if (baseModel.startsWith("anthropic.claude")) {
// Handle the new Converse API response format
if (response.output?.message) {
return {
role: response.output.message.role,
content: response.output.message.content.map((item: any) => {
if ("text" in item) return { type: "text", text: item.text || "" };
if ("image" in item) {
return {
type: "image_url",
image_url: {
url: `data:${item.source?.media_type || "image/png"};base64,${
item.source?.data || ""
}`,
},
};
}
return { type: "text", text: item.text || "" };
}),
stop_reason: response.stopReason,
usage: response.usage,
};
}
// Fallback for older format
return formatResponse(
response.content ||
(response.completion
? [{ type: "text", text: response.completion }]
: []),
);
} else if (baseModel.startsWith("meta.llama")) {
return formatResponse(response.generation || response.completion || "");
} else if (baseModel.startsWith("amazon.titan")) {
return formatResponse(response.results?.[0]?.outputText || "");
} else if (baseModel.startsWith("mistral.")) {
return formatResponse(
response.outputs?.[0]?.text || response.response || "",
);
}
throw new Error(`Unsupported model: ${baseModel}`);
} catch (e) {
console.error("[Bedrock] Failed to parse response:", e);
// Return raw text as fallback
return {
role: "assistant",
content: [{ type: "text", text: responseBody }],
};
}
}