🔧 fix: Improve Assistants File Citation & Download Handling (#2248)

* fix(processMessages): properly handle assistant file citations and add sources list

* feat: improve file download UX by making any downloaded files accessible within the app post-download

* refactor(processOpenAIImageOutput): correctly handle two different outputs for images since OpenAI generates a file in their storage, shares filepath for image rendering

* refactor: create `addFileToCache` helper to use across frontend

* refactor: add ImageFile parts to cache on processing content stream

api/server/routes/files/files.js

@@ -104,6 +104,7 @@ router.get('/download/:userId/:filepath', async (req, res) => {
     const setHeaders = () => {
       res.setHeader('Content-Disposition', `attachment; filename="${file.filename}"`);
       res.setHeader('Content-Type', 'application/octet-stream');
+      res.setHeader('X-File-Metadata', JSON.stringify(file));
     };
 
     /** @type {{ body: import('stream').PassThrough } | undefined} */

api/server/services/AssistantService.js

@@ -1,4 +1,3 @@
-const path = require('path');
 const { klona } = require('klona');
 const {
   StepTypes,
@@ -233,14 +232,9 @@ function createInProgressHandler(openai, thread_id, messages) {
               file_id,
               basename: `${file_id}.png`,
             });
-            // toolCall.asset_pointer = file.filepath;
-            const prelimImage = {
-              file_id,
-              filename: path.basename(file.filepath),
-              filepath: file.filepath,
-              height: file.height,
-              width: file.width,
-            };
+
+            const prelimImage = file;
+
             // check if every key has a value before adding to content
             const prelimImageKeys = Object.keys(prelimImage);
             const validImageFile = prelimImageKeys.every((key) => prelimImage[key]);

api/server/services/Files/process.js

@@ -9,6 +9,8 @@ const {
   imageExtRegex,
   EModelEndpoint,
   mergeFileConfig,
+  hostImageIdSuffix,
+  hostImageNamePrefix,
 } = require('librechat-data-provider');
 const { convertToWebP, resizeAndConvert } = require('~/server/services/Files/images');
 const { initializeClient } = require('~/server/services/Endpoints/assistants');
@@ -309,7 +311,7 @@ const processFileUpload = async ({ req, res, file, metadata }) => {
  * @param {OpenAI} params.openai - The OpenAI client instance.
  * @param {string} params.file_id - The ID of the file to retrieve.
  * @param {string} params.userId - The user ID.
- * @param {string} params.filename - The name of the file.
+ * @param {string} [params.filename] - The name of the file. `undefined` for `file_citation` annotations.
  * @param {boolean} [params.saveFile=false] - Whether to save the file metadata to the database.
  * @param {boolean} [params.updateUsage=false] - Whether to update file usage in database.
  */
@@ -322,18 +324,23 @@ const processOpenAIFile = async ({
   updateUsage = false,
 }) => {
   const _file = await openai.files.retrieve(file_id);
-  const filepath = `${openai.baseURL}/files/${userId}/${file_id}/${filename}`;
+  const originalName = filename ?? (_file.filename ? path.basename(_file.filename) : undefined);
+  const filepath = `${openai.baseURL}/files/${userId}/${file_id}${
+    originalName ? `/${originalName}` : ''
+  }`;
+  const type = mime.getType(originalName ?? file_id);
+
   const file = {
     ..._file,
+    type,
     file_id,
     filepath,
     usage: 1,
-    filename,
     user: userId,
+    context: _file.purpose,
     source: FileSources.openai,
     model: openai.req.body.model,
-    type: mime.getType(filename),
-    context: FileContext.assistants_output,
+    filename: originalName ?? file_id,
   };
 
   if (saveFile) {
@@ -360,18 +367,32 @@ const processOpenAIFile = async ({
  * @returns {Promise<MongoFile>} The file metadata.
  */
 const processOpenAIImageOutput = async ({ req, buffer, file_id, filename, fileExt }) => {
+  const currentDate = new Date();
+  const formattedDate = currentDate.toISOString();
   const _file = await convertToWebP(req, buffer, 'high', `${file_id}${fileExt}`);
   const file = {
     ..._file,
-    file_id,
     usage: 1,
-    filename,
     user: req.user.id,
     type: 'image/webp',
+    createdAt: formattedDate,
+    updatedAt: formattedDate,
     source: req.app.locals.fileStrategy,
     context: FileContext.assistants_output,
+    file_id: `${file_id}${hostImageIdSuffix}`,
+    filename: `${hostImageNamePrefix}${filename}`,
   };
   createFile(file, true);
+  createFile(
+    {
+      ...file,
+      file_id,
+      filename,
+      source: FileSources.openai,
+      type: mime.getType(fileExt),
+    },
+    true,
+  );
   return file;
 };
 
@@ -382,7 +403,7 @@ const processOpenAIImageOutput = async ({ req, buffer, file_id, filename, fileEx
  * @param {OpenAIClient} params.openai - The OpenAI client instance.
  * @param {RunClient} params.client - The LibreChat client instance: either refers to `openai` or `streamRunManager`.
  * @param {string} params.file_id - The ID of the file to retrieve.
- * @param {string} params.basename - The basename of the file (if image); e.g., 'image.jpg'.
+ * @param {string} [params.basename] - The basename of the file (if image); e.g., 'image.jpg'. `undefined` for `file_citation` annotations.
  * @param {boolean} [params.unknownType] - Whether the file type is unknown.
  * @returns {Promise<{file_id: string, filepath: string, source: string, bytes?: number, width?: number, height?: number} | null>}
  * - Returns null if `file_id` is not defined; else, the file metadata if successfully retrieved and processed.
@@ -398,14 +419,19 @@ async function retrieveAndProcessFile({
     return null;
   }
 
+  let basename = _basename;
+  const processArgs = { openai, file_id, filename: basename, userId: client.req.user.id };
+
+  // If no basename provided, return only the file metadata
+  if (!basename) {
+    return await processOpenAIFile({ ...processArgs, saveFile: true });
+  }
+
+  const fileExt = path.extname(basename);
   if (client.attachedFileIds?.has(file_id) || client.processedFileIds?.has(file_id)) {
     return processOpenAIFile({ ...processArgs, updateUsage: true });
   }
 
-  let basename = _basename;
-  const fileExt = path.extname(basename);
-  const processArgs = { openai, file_id, filename: basename, userId: client.req.user.id };
-
   /**
    * @returns {Promise<Buffer>} The file data buffer.
    */
@@ -415,11 +441,6 @@ async function retrieveAndProcessFile({
     return Buffer.from(arrayBuffer);
   };
 
-  // If no basename provided, return only the file metadata
-  if (!basename) {
-    return await processOpenAIFile({ ...processArgs, saveFile: true });
-  }
-
   let dataBuffer;
   if (unknownType || !fileExt || imageExtRegex.test(basename)) {
     try {

api/server/services/Runs/StreamRunManager.js

@@ -1,4 +1,3 @@
-const path = require('path');
 const {
   StepTypes,
   ContentTypes,
@@ -222,14 +221,9 @@ class StreamRunManager {
       file_id,
       basename: `${file_id}.png`,
     });
-    // toolCall.asset_pointer = file.filepath;
-    const prelimImage = {
-      file_id,
-      filename: path.basename(file.filepath),
-      filepath: file.filepath,
-      height: file.height,
-      width: file.width,
-    };
+
+    const prelimImage = file;
+
     // check if every key has a value before adding to content
     const prelimImageKeys = Object.keys(prelimImage);
     const validImageFile = prelimImageKeys.every((key) => prelimImage[key]);

api/server/services/Threads/manage.js

@@ -549,6 +549,7 @@ async function processMessages({ openai, client, messages = [] }) {
 
   let text = '';
   let edited = false;
+  const sources = [];
   for (const message of sorted) {
     message.files = [];
     for (const content of message.content) {
@@ -588,6 +589,17 @@ async function processMessages({ openai, client, messages = [] }) {
         const file_id = annotationType?.file_id;
         const alreadyProcessed = client.processedFileIds.has(file_id);
 
+        const replaceCurrentAnnotation = (replacement = '') => {
+          currentText = replaceAnnotation(
+            currentText,
+            annotation.start_index,
+            annotation.end_index,
+            annotation.text,
+            replacement,
+          );
+          edited = true;
+        };
+
         if (alreadyProcessed) {
           const { file_id } = annotationType || {};
           file = await retrieveAndProcessFile({ openai, client, file_id, unknownType: true });
@@ -599,6 +611,7 @@ async function processMessages({ openai, client, messages = [] }) {
             file_id,
             basename,
           });
+          replaceCurrentAnnotation(file.filepath);
         } else if (type === AnnotationTypes.FILE_CITATION) {
           file = await retrieveAndProcessFile({
             openai,
@@ -606,17 +619,8 @@ async function processMessages({ openai, client, messages = [] }) {
             file_id,
             unknownType: true,
           });
-        }
-
-        if (file.filepath) {
-          currentText = replaceAnnotation(
-            currentText,
-            annotation.start_index,
-            annotation.end_index,
-            annotation.text,
-            file.filepath,
-          );
-          edited = true;
+          sources.push(file.filename);
+          replaceCurrentAnnotation(`^${sources.length}^`);
         }
 
         text += currentText + ' ';
@@ -631,6 +635,13 @@ async function processMessages({ openai, client, messages = [] }) {
     }
   }
 
+  if (sources.length) {
+    text += '\n\n';
+    for (let i = 0; i < sources.length; i++) {
+      text += `^${i + 1}.^ ${sources[i]}${i === sources.length - 1 ? '' : '\n'}`;
+    }
+  }
+
   return { messages: sorted, text, edited };
 }