From ea2cbc55a7479c98aeaad1c2aeea705e667358cd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 26 Mar 2025 15:04:01 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix:=20S3=20Download=20Stream=20?= =?UTF-8?q?with=20Key=20Extraction=20and=20Blob=20Storage=20Encoding=20for?= =?UTF-8?q?=20Vision=20(#6557)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/services/Files/S3/crud.js | 27 ++++++++++-- api/server/services/Files/images/encode.js | 50 +++++++++++++++++++--- api/typedefs.js | 11 +++++ 3 files changed, 79 insertions(+), 9 deletions(-) diff --git a/api/server/services/Files/S3/crud.js b/api/server/services/Files/S3/crud.js index 06f9116b69..a9ddcf455e 100644 --- a/api/server/services/Files/S3/crud.js +++ b/api/server/services/Files/S3/crud.js @@ -135,20 +135,41 @@ async function uploadFileToS3({ req, file, file_id, basePath = defaultBasePath } } } +/** + * Extracts the S3 key from a full S3 URL. + * + * @param {string} s3Url - The full S3 URL + * @returns {string} The S3 key + */ +function extractKeyFromS3Url(s3Url) { + try { + // Parse the URL + const url = new URL(s3Url); + // Extract the path from the URL, removing the leading slash + let key = url.pathname.substring(1); + + return key; + } catch (error) { + throw new Error(`Failed to extract key from S3 URL: ${error.message}`); + } +} + /** * Retrieves a readable stream for a file stored in S3. * + * @param {ServerRequest} req - Server request object. * @param {string} filePath - The S3 key of the file. * @returns {Promise} */ -async function getS3FileStream(filePath) { - const params = { Bucket: bucketName, Key: filePath }; +async function getS3FileStream(_req, filePath) { try { + const Key = extractKeyFromS3Url(filePath); + const params = { Bucket: bucketName, Key }; const s3 = initializeS3(); const data = await s3.send(new GetObjectCommand(params)); return data.Body; // Returns a Node.js ReadableStream. } catch (error) { - logger.error('[getS3FileStream] Error retrieving S3 file stream:', error.message); + logger.error('[getS3FileStream] Error retrieving S3 file stream:', error); throw error; } } diff --git a/api/server/services/Files/images/encode.js b/api/server/services/Files/images/encode.js index 707632fb6a..a23ce4cf2c 100644 --- a/api/server/services/Files/images/encode.js +++ b/api/server/services/Files/images/encode.js @@ -37,17 +37,21 @@ const base64Only = new Set([ EModelEndpoint.bedrock, ]); +const blobStorageSources = new Set([FileSources.azure, FileSources.s3]); + /** * Encodes and formats the given files. * @param {Express.Request} req - The request object. * @param {Array} files - The array of files to encode and format. * @param {EModelEndpoint} [endpoint] - Optional: The endpoint for the image. * @param {string} [mode] - Optional: The endpoint mode for the image. - * @returns {Promise} - A promise that resolves to the result object containing the encoded images and file details. + * @returns {Promise<{ text: string; files: MongoFile[]; image_urls: MessageContentImageUrl[] }>} - A promise that resolves to the result object containing the encoded images and file details. */ async function encodeAndFormat(req, files, endpoint, mode) { const promises = []; + /** @type {Record, 'prepareImagePayload' | 'getDownloadStream'>>} */ const encodingMethods = {}; + /** @type {{ text: string; files: MongoFile[]; image_urls: MessageContentImageUrl[] }} */ const result = { text: '', files: [], @@ -59,6 +63,7 @@ async function encodeAndFormat(req, files, endpoint, mode) { } for (let file of files) { + /** @type {FileSources} */ const source = file.source ?? FileSources.local; if (source === FileSources.text && file.text) { result.text += `${!result.text ? 'Attached document(s):\n```md' : '\n\n---\n\n'}# "${file.filename}"\n${file.text}\n`; @@ -70,18 +75,51 @@ async function encodeAndFormat(req, files, endpoint, mode) { } if (!encodingMethods[source]) { - const { prepareImagePayload } = getStrategyFunctions(source); + const { prepareImagePayload, getDownloadStream } = getStrategyFunctions(source); if (!prepareImagePayload) { throw new Error(`Encoding function not implemented for ${source}`); } - encodingMethods[source] = prepareImagePayload; + encodingMethods[source] = { prepareImagePayload, getDownloadStream }; } - const preparePayload = encodingMethods[source]; + const preparePayload = encodingMethods[source].prepareImagePayload; + /* We need to fetch the image and convert it to base64 if we are using S3/Azure Blob storage. */ + if (blobStorageSources.has(source)) { + try { + const downloadStream = encodingMethods[source].getDownloadStream; + const stream = await downloadStream(req, file.filepath); + const streamPromise = new Promise((resolve, reject) => { + /** @type {Uint8Array[]} */ + const chunks = []; + stream.on('readable', () => { + let chunk; + while (null !== (chunk = stream.read())) { + chunks.push(chunk); + } + }); - /* Google & Anthropic don't support passing URLs to payload */ - if (source !== FileSources.local && base64Only.has(endpoint)) { + stream.on('end', () => { + const buffer = Buffer.concat(chunks); + const base64Data = buffer.toString('base64'); + resolve(base64Data); + }); + stream.on('error', (error) => { + reject(error); + }); + }); + const base64Data = await streamPromise; + promises.push([file, base64Data]); + } catch (error) { + logger.error( + `Error processing blob storage file stream for ${file.name} base64 payload:`, + error, + ); + continue; + } + + /* Google & Anthropic don't support passing URLs to payload */ + } else if (source !== FileSources.local && base64Only.has(endpoint)) { const [_file, imageURL] = await preparePayload(req, file); promises.push([_file, await fetchImageToBase64(imageURL)]); continue; diff --git a/api/typedefs.js b/api/typedefs.js index 21c4f1fecc..74d9ef5209 100644 --- a/api/typedefs.js +++ b/api/typedefs.js @@ -403,6 +403,12 @@ * @memberof typedefs */ +/** + * @exports MessageContentImageUrl + * @typedef {import('librechat-data-provider').Agents.MessageContentImageUrl} MessageContentImageUrl + * @memberof typedefs + */ + /** Prompts */ /** * @exports TPrompt @@ -759,6 +765,11 @@ * @typedef {import('mongoose').Schema} MongooseSchema * @memberof typedefs */ +/** + * @exports MongoFile + * @typedef {import('@librechat/data-schemas').IMongoFile} MongoFile + * @memberof typedefs + */ /** * @exports ObjectId