mirror of
https://github.com/ChatGPTNextWeb/ChatGPT-Next-Web.git
synced 2025-05-22 21:50:16 +09:00
622 lines
16 KiB
TypeScript
622 lines
16 KiB
TypeScript
import SHA256 from "crypto-js/sha256";
|
|
import HmacSHA256 from "crypto-js/hmac-sha256";
|
|
import Hex from "crypto-js/enc-hex";
|
|
import Utf8 from "crypto-js/enc-utf8";
|
|
import { AES, enc, lib, PBKDF2, mode, pad, algo } from "crypto-js";
|
|
|
|
// Types and Interfaces
|
|
export interface BedrockCredentials {
|
|
region: string;
|
|
accessKeyId: string;
|
|
secretAccessKey: string;
|
|
}
|
|
|
|
// Type definitions for better type safety
|
|
type ParsedEvent = Record<string, any>;
|
|
type EventResult = ParsedEvent[];
|
|
|
|
// Encryption utilities
|
|
function generateSalt(): string {
|
|
const salt = lib.WordArray.random(128 / 8);
|
|
return salt.toString(enc.Base64);
|
|
}
|
|
|
|
function generateIV(): string {
|
|
const iv = lib.WordArray.random(128 / 8);
|
|
return iv.toString(enc.Base64);
|
|
}
|
|
|
|
function deriveKey(password: string, salt: string): lib.WordArray {
|
|
// Use PBKDF2 with SHA256 for key derivation
|
|
return PBKDF2(password, salt, {
|
|
keySize: 256 / 32,
|
|
iterations: 10000,
|
|
hasher: algo.SHA256,
|
|
});
|
|
}
|
|
|
|
// Using a dot as separator since it's not used in Base64
|
|
const SEPARATOR = ".";
|
|
|
|
export function encrypt(data: string, encryptionKey: string): string {
|
|
if (!data) return "";
|
|
if (!encryptionKey) {
|
|
throw new Error("Encryption key is required for AWS credential encryption");
|
|
}
|
|
try {
|
|
// Generate salt and IV
|
|
const salt = generateSalt();
|
|
const iv = generateIV();
|
|
|
|
// Derive key using PBKDF2
|
|
const key = deriveKey(encryptionKey, salt);
|
|
|
|
// Encrypt the data
|
|
const encrypted = AES.encrypt(data, key, {
|
|
iv: enc.Base64.parse(iv),
|
|
mode: mode.CBC,
|
|
padding: pad.Pkcs7,
|
|
});
|
|
|
|
// Combine salt, IV, and encrypted data
|
|
// Format: salt.iv.encryptedData
|
|
return [salt, iv, encrypted.toString()].join(SEPARATOR);
|
|
} catch (error) {
|
|
throw new Error("Failed to encrypt AWS credentials");
|
|
}
|
|
}
|
|
|
|
export function decrypt(encryptedData: string, encryptionKey: string): string {
|
|
if (!encryptedData) return "";
|
|
if (!encryptionKey) {
|
|
throw new Error("Encryption key is required for AWS credential decryption");
|
|
}
|
|
try {
|
|
let components = encryptedData.split(SEPARATOR);
|
|
const [salt, iv, data] = components;
|
|
// For new format, use the provided salt and IV
|
|
const key = deriveKey(encryptionKey, salt);
|
|
const decrypted = AES.decrypt(data, key, {
|
|
iv: enc.Base64.parse(iv),
|
|
mode: mode.CBC,
|
|
padding: pad.Pkcs7,
|
|
});
|
|
|
|
const result = decrypted.toString(enc.Utf8);
|
|
if (!result) {
|
|
throw new Error("Failed to decrypt AWS credentials");
|
|
}
|
|
return result;
|
|
} catch (error) {
|
|
throw new Error("Failed to decrypt AWS credentials");
|
|
}
|
|
}
|
|
|
|
export function maskSensitiveValue(value: string): string {
|
|
if (!value) return "";
|
|
if (value.length <= 4) return value;
|
|
// Use constant-time operations to prevent timing attacks
|
|
const masked = Buffer.alloc(value.length - 4, "*").toString();
|
|
return value.slice(0, 2) + masked + value.slice(-2);
|
|
}
|
|
|
|
// AWS Signing
|
|
export interface SignParams {
|
|
method: string;
|
|
url: string;
|
|
region: string;
|
|
accessKeyId: string;
|
|
secretAccessKey: string;
|
|
body: string | object;
|
|
service: string;
|
|
headers?: Record<string, string>;
|
|
isStreaming?: boolean;
|
|
}
|
|
|
|
function hmac(
|
|
key: string | CryptoJS.lib.WordArray,
|
|
data: string,
|
|
): CryptoJS.lib.WordArray {
|
|
if (typeof key === "string") {
|
|
key = Utf8.parse(key);
|
|
}
|
|
return HmacSHA256(data, key);
|
|
}
|
|
|
|
function getSigningKey(
|
|
secretKey: string,
|
|
dateStamp: string,
|
|
region: string,
|
|
service: string,
|
|
): CryptoJS.lib.WordArray {
|
|
const kDate = hmac("AWS4" + secretKey, dateStamp);
|
|
const kRegion = hmac(kDate, region);
|
|
const kService = hmac(kRegion, service);
|
|
const kSigning = hmac(kService, "aws4_request");
|
|
return kSigning;
|
|
}
|
|
|
|
function normalizeHeaderValue(value: string): string {
|
|
return value.replace(/\s+/g, " ").trim();
|
|
}
|
|
|
|
function encodeRFC3986(str: string): string {
|
|
return encodeURIComponent(str)
|
|
.replace(
|
|
/[!'()*]/g,
|
|
(c) => "%" + c.charCodeAt(0).toString(16).toUpperCase(),
|
|
)
|
|
.replace(/[-_.~]/g, (c) => c);
|
|
}
|
|
|
|
function getCanonicalUri(path: string): string {
|
|
if (!path || path === "/") return "/";
|
|
|
|
return (
|
|
"/" +
|
|
path
|
|
.split("/")
|
|
.map((segment) => {
|
|
if (!segment) return "";
|
|
if (segment === "invoke-with-response-stream") return segment;
|
|
|
|
if (segment.includes("model/")) {
|
|
return segment
|
|
.split(/(model\/)/)
|
|
.map((part) => {
|
|
if (part === "model/") return part;
|
|
return part
|
|
.split(/([.:])/g)
|
|
.map((subpart, i) =>
|
|
i % 2 === 1 ? subpart : encodeRFC3986(subpart),
|
|
)
|
|
.join("");
|
|
})
|
|
.join("");
|
|
}
|
|
|
|
return encodeRFC3986(segment);
|
|
})
|
|
.join("/")
|
|
);
|
|
}
|
|
|
|
export async function sign({
|
|
method,
|
|
url,
|
|
region,
|
|
accessKeyId,
|
|
secretAccessKey,
|
|
body,
|
|
service,
|
|
headers: customHeaders = {},
|
|
isStreaming = true,
|
|
}: SignParams): Promise<Record<string, string>> {
|
|
try {
|
|
const endpoint = new URL(url);
|
|
const canonicalUri = getCanonicalUri(endpoint.pathname.slice(1));
|
|
const canonicalQueryString = endpoint.search.slice(1);
|
|
|
|
const now = new Date();
|
|
const amzDate = now.toISOString().replace(/[:-]|\.\d{3}/g, "");
|
|
const dateStamp = amzDate.slice(0, 8);
|
|
|
|
const bodyString = typeof body === "string" ? body : JSON.stringify(body);
|
|
const payloadHash = SHA256(bodyString).toString(Hex);
|
|
|
|
const headers: Record<string, string> = {
|
|
accept: isStreaming
|
|
? "application/vnd.amazon.eventstream"
|
|
: "application/json",
|
|
"content-type": "application/json",
|
|
host: endpoint.host,
|
|
"x-amz-content-sha256": payloadHash,
|
|
"x-amz-date": amzDate,
|
|
...customHeaders,
|
|
};
|
|
|
|
if (isStreaming) {
|
|
headers["x-amzn-bedrock-accept"] = "*/*";
|
|
}
|
|
|
|
const sortedHeaderKeys = Object.keys(headers).sort((a, b) =>
|
|
a.toLowerCase().localeCompare(b.toLowerCase()),
|
|
);
|
|
|
|
const canonicalHeaders = sortedHeaderKeys
|
|
.map(
|
|
(key) => `${key.toLowerCase()}:${normalizeHeaderValue(headers[key])}\n`,
|
|
)
|
|
.join("");
|
|
|
|
const signedHeaders = sortedHeaderKeys
|
|
.map((key) => key.toLowerCase())
|
|
.join(";");
|
|
|
|
const canonicalRequest = [
|
|
method.toUpperCase(),
|
|
canonicalUri,
|
|
canonicalQueryString,
|
|
canonicalHeaders,
|
|
signedHeaders,
|
|
payloadHash,
|
|
].join("\n");
|
|
|
|
const algorithm = "AWS4-HMAC-SHA256";
|
|
const credentialScope = `${dateStamp}/${region}/${service}/aws4_request`;
|
|
const stringToSign = [
|
|
algorithm,
|
|
amzDate,
|
|
credentialScope,
|
|
SHA256(canonicalRequest).toString(Hex),
|
|
].join("\n");
|
|
|
|
const signingKey = getSigningKey(
|
|
secretAccessKey,
|
|
dateStamp,
|
|
region,
|
|
service,
|
|
);
|
|
const signature = hmac(signingKey, stringToSign).toString(Hex);
|
|
|
|
const authorization = [
|
|
`${algorithm} Credential=${accessKeyId}/${credentialScope}`,
|
|
`SignedHeaders=${signedHeaders}`,
|
|
`Signature=${signature}`,
|
|
].join(", ");
|
|
|
|
return {
|
|
...headers,
|
|
Authorization: authorization,
|
|
};
|
|
} catch (error) {
|
|
console.error("[AWS Signing Error]:", error);
|
|
throw new Error("Failed to sign AWS request");
|
|
}
|
|
}
|
|
|
|
// Bedrock utilities
|
|
function decodeBase64(base64String: string): string {
|
|
try {
|
|
return Buffer.from(base64String, "base64").toString("utf-8");
|
|
} catch (e) {
|
|
console.error("[Base64 Decode Error]:", e);
|
|
return "";
|
|
}
|
|
}
|
|
|
|
export function parseEventData(chunk: Uint8Array): EventResult {
|
|
const decoder = new TextDecoder();
|
|
const text = decoder.decode(chunk);
|
|
const results: EventResult = [];
|
|
|
|
try {
|
|
// First try to parse as regular JSON
|
|
const parsed = JSON.parse(text);
|
|
if (parsed.bytes) {
|
|
const decoded = decodeBase64(parsed.bytes);
|
|
try {
|
|
const decodedJson = JSON.parse(decoded);
|
|
results.push(decodedJson);
|
|
} catch (e) {
|
|
results.push({ output: decoded });
|
|
}
|
|
return results;
|
|
}
|
|
|
|
if (typeof parsed.body === "string") {
|
|
try {
|
|
const parsedBody = JSON.parse(parsed.body);
|
|
results.push(parsedBody);
|
|
} catch (e) {
|
|
results.push({ output: parsed.body });
|
|
}
|
|
return results;
|
|
}
|
|
|
|
results.push(parsed.body || parsed);
|
|
return results;
|
|
} catch (e) {
|
|
// If regular JSON parse fails, try to extract event content
|
|
const eventRegex = /:event-type[^\{]+(\{[^\}]+\})/g;
|
|
let match;
|
|
|
|
while ((match = eventRegex.exec(text)) !== null) {
|
|
try {
|
|
const eventData = match[1];
|
|
const parsed = JSON.parse(eventData);
|
|
|
|
if (parsed.bytes) {
|
|
const decoded = decodeBase64(parsed.bytes);
|
|
try {
|
|
const decodedJson = JSON.parse(decoded);
|
|
if (decodedJson.choices?.[0]?.message?.content) {
|
|
results.push({ output: decodedJson.choices[0].message.content });
|
|
} else {
|
|
results.push(decodedJson);
|
|
}
|
|
} catch (e) {
|
|
results.push({ output: decoded });
|
|
}
|
|
} else {
|
|
results.push(parsed);
|
|
}
|
|
} catch (e) {
|
|
console.debug("[Event Parse Warning]:", e);
|
|
}
|
|
}
|
|
|
|
// If no events were found, try to extract clean text
|
|
if (results.length === 0) {
|
|
// Remove event metadata markers and clean the text
|
|
const cleanText = text
|
|
.replace(/\{KG[^:]+:event-type[^}]+\}/g, "") // Remove event markers
|
|
.replace(/[\x00-\x1F\x7F-\x9F\uFEFF]/g, "") // Remove control characters
|
|
.trim();
|
|
|
|
if (cleanText) {
|
|
results.push({ output: cleanText });
|
|
}
|
|
}
|
|
}
|
|
|
|
return results;
|
|
}
|
|
|
|
export function processMessage(
|
|
data: ParsedEvent,
|
|
remainText: string,
|
|
runTools: any[],
|
|
index: number,
|
|
): { remainText: string; index: number } {
|
|
if (!data) return { remainText, index };
|
|
|
|
try {
|
|
// Handle Nova's tool calls
|
|
// console.log("processMessage data=========================",data);
|
|
if (
|
|
data.stopReason === "tool_use" &&
|
|
data.output?.message?.content?.[0]?.toolUse
|
|
) {
|
|
const toolUse = data.output.message.content[0].toolUse;
|
|
index += 1;
|
|
runTools.push({
|
|
id: `tool-${Date.now()}`,
|
|
type: "function",
|
|
function: {
|
|
name: toolUse.name,
|
|
arguments: JSON.stringify(toolUse.input),
|
|
},
|
|
});
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle Nova's text content
|
|
if (data.output?.message?.content?.[0]?.text) {
|
|
remainText += data.output.message.content[0].text;
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle Nova's messageStart event
|
|
if (data.messageStart) {
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle Nova's contentBlockDelta event
|
|
if (data.contentBlockDelta) {
|
|
if (data.contentBlockDelta.delta?.text) {
|
|
remainText += data.contentBlockDelta.delta.text;
|
|
}
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle Nova's contentBlockStop event
|
|
if (data.contentBlockStop) {
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle Nova's messageStop event
|
|
if (data.messageStop) {
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle message_start event (for other models)
|
|
if (data.type === "message_start") {
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle content_block_start event (for other models)
|
|
if (data.type === "content_block_start") {
|
|
if (data.content_block?.type === "tool_use") {
|
|
index += 1;
|
|
runTools.push({
|
|
id: data.content_block.id,
|
|
type: "function",
|
|
function: {
|
|
name: data.content_block.name,
|
|
arguments: "",
|
|
},
|
|
});
|
|
}
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle content_block_delta event (for other models)
|
|
if (data.type === "content_block_delta") {
|
|
if (data.delta?.type === "input_json_delta" && runTools[index]) {
|
|
runTools[index].function.arguments += data.delta.partial_json;
|
|
} else if (data.delta?.type === "text_delta") {
|
|
const newText = data.delta.text || "";
|
|
remainText += newText;
|
|
}
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle tool calls for other models
|
|
if (data.choices?.[0]?.message?.tool_calls) {
|
|
for (const toolCall of data.choices[0].message.tool_calls) {
|
|
index += 1;
|
|
runTools.push({
|
|
id: toolCall.id || `tool-${Date.now()}`,
|
|
type: "function",
|
|
function: {
|
|
name: toolCall.function?.name,
|
|
arguments: toolCall.function?.arguments || "",
|
|
},
|
|
});
|
|
}
|
|
return { remainText, index };
|
|
}
|
|
|
|
// Handle various response formats
|
|
let newText = "";
|
|
if (data.delta?.text) {
|
|
newText = data.delta.text;
|
|
} else if (data.choices?.[0]?.message?.content) {
|
|
newText = data.choices[0].message.content;
|
|
} else if (data.content?.[0]?.text) {
|
|
newText = data.content[0].text;
|
|
} else if (data.generation) {
|
|
newText = data.generation;
|
|
} else if (data.outputText) {
|
|
newText = data.outputText;
|
|
} else if (data.response) {
|
|
newText = data.response;
|
|
} else if (data.output) {
|
|
newText = data.output;
|
|
}
|
|
|
|
// Only append if we have new text
|
|
if (newText) {
|
|
remainText += newText;
|
|
}
|
|
} catch (e) {
|
|
console.error("[Bedrock Request] parse error", e);
|
|
}
|
|
|
|
return { remainText, index };
|
|
}
|
|
|
|
export function processChunks(
|
|
chunks: Uint8Array[],
|
|
pendingChunk: Uint8Array | null,
|
|
remainText: string,
|
|
runTools: any[],
|
|
index: number,
|
|
): {
|
|
chunks: Uint8Array[];
|
|
pendingChunk: Uint8Array | null;
|
|
remainText: string;
|
|
index: number;
|
|
} {
|
|
let currentText = remainText;
|
|
let currentIndex = index;
|
|
|
|
while (chunks.length > 0) {
|
|
const chunk = chunks[0];
|
|
try {
|
|
// If there's a pending chunk, try to merge it with the current chunk
|
|
let chunkToProcess = chunk;
|
|
if (pendingChunk) {
|
|
const mergedChunk = new Uint8Array(pendingChunk.length + chunk.length);
|
|
mergedChunk.set(pendingChunk);
|
|
mergedChunk.set(chunk, pendingChunk.length);
|
|
chunkToProcess = mergedChunk;
|
|
pendingChunk = null;
|
|
}
|
|
|
|
// Try to process the chunk
|
|
const parsedEvents = parseEventData(chunkToProcess);
|
|
if (parsedEvents.length > 0) {
|
|
// Process each event in the chunk
|
|
for (const parsed of parsedEvents) {
|
|
const result = processMessage(
|
|
parsed,
|
|
currentText,
|
|
runTools,
|
|
currentIndex,
|
|
);
|
|
currentText = result.remainText;
|
|
currentIndex = result.index;
|
|
}
|
|
chunks.shift(); // Remove processed chunk
|
|
} else {
|
|
// If parsing fails, it might be an incomplete chunk
|
|
pendingChunk = chunkToProcess;
|
|
chunks.shift();
|
|
}
|
|
} catch (e) {
|
|
console.error("[Chunk Process Error]:", e);
|
|
chunks.shift(); // Remove error chunk
|
|
pendingChunk = null; // Reset pending chunk on error
|
|
}
|
|
}
|
|
|
|
return {
|
|
chunks,
|
|
pendingChunk,
|
|
remainText: currentText,
|
|
index: currentIndex,
|
|
};
|
|
}
|
|
|
|
export function getBedrockEndpoint(
|
|
region: string,
|
|
modelId: string,
|
|
shouldStream: boolean,
|
|
): string {
|
|
if (!region || !modelId) {
|
|
throw new Error("Region and model ID are required for Bedrock endpoint");
|
|
}
|
|
const baseEndpoint = `https://bedrock-runtime.${region}.amazonaws.com`;
|
|
const endpoint =
|
|
shouldStream === false
|
|
? `${baseEndpoint}/model/${modelId}/invoke`
|
|
: `${baseEndpoint}/model/${modelId}/invoke-with-response-stream`;
|
|
return endpoint;
|
|
}
|
|
|
|
export function extractMessage(res: any, modelId: string = ""): string {
|
|
if (!res) {
|
|
console.error("[AWS Extract Error] extractMessage Empty response");
|
|
return "";
|
|
}
|
|
|
|
let message = "";
|
|
|
|
// Handle Nova model response format
|
|
if (modelId.toLowerCase().includes("nova")) {
|
|
if (res.output?.message?.content?.[0]?.text) {
|
|
message = res.output.message.content[0].text;
|
|
} else {
|
|
message = res.output || "";
|
|
}
|
|
}
|
|
// Handle Mistral model response format
|
|
else if (modelId.toLowerCase().includes("mistral")) {
|
|
if (res.choices?.[0]?.message?.content) {
|
|
message = res.choices[0].message.content;
|
|
} else {
|
|
message = res.output || "";
|
|
}
|
|
}
|
|
// Handle Llama model response format
|
|
else if (modelId.toLowerCase().includes("llama")) {
|
|
message = res?.generation || "";
|
|
}
|
|
// Handle Titan model response format
|
|
else if (modelId.toLowerCase().includes("titan")) {
|
|
message = res?.outputText || "";
|
|
}
|
|
// Handle Claude and other models
|
|
else if (res.content?.[0]?.text) {
|
|
message = res.content[0].text;
|
|
}
|
|
// Handle other response formats
|
|
else {
|
|
message = res.output || res.response || res.message || "";
|
|
}
|
|
|
|
return message;
|
|
}
|