🔧 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

client/src/data-provider/queries.ts

@@ -21,7 +21,7 @@ import type {
   TEndpointsConfig,
   TCheckUserKeyResponse,
 } from 'librechat-data-provider';
-import { findPageForConversation } from '~/utils';
+import { findPageForConversation, addFileToCache } from '~/utils';
 
 export const useGetFiles = <TData = TFile[] | boolean>(
   config?: UseQueryOptions<TFile[], unknown, TData>,
@@ -326,15 +326,29 @@ export const useGetAssistantDocsQuery = (
 };
 
 export const useFileDownload = (userId: string, filepath: string): QueryObserverResult<string> => {
+  const queryClient = useQueryClient();
   return useQuery(
     [QueryKeys.fileDownload, filepath],
     async () => {
       if (!userId) {
         console.warn('No user ID provided for file download');
       }
-      const blob = await dataService.getFileDownload(userId, filepath);
-      const downloadUrl = window.URL.createObjectURL(blob);
-      return downloadUrl;
+      const response = await dataService.getFileDownload(userId, filepath);
+      const blob = response.data;
+      const downloadURL = window.URL.createObjectURL(blob);
+      try {
+        const metadata: TFile | undefined = JSON.parse(response.headers['x-file-metadata']);
+        if (!metadata) {
+          console.warn('No metadata found for file download', response.headers);
+          return downloadURL;
+        }
+
+        addFileToCache(queryClient, metadata);
+      } catch (e) {
+        console.error('Error parsing file metadata, skipped updating file query cache', e);
+      }
+
+      return downloadURL;
     },
     {
       enabled: false,

client/src/hooks/SSE/useContentHandler.ts

@@ -1,13 +1,18 @@
+import { useCallback, useMemo } from 'react';
 import { ContentTypes } from 'librechat-data-provider';
+import { useQueryClient } from '@tanstack/react-query';
+
 import type {
   Text,
   TMessage,
+  ImageFile,
   TSubmission,
   ContentPart,
+  PartMetadata,
   TContentData,
   TMessageContentParts,
 } from 'librechat-data-provider';
-import { useCallback, useMemo } from 'react';
+import { addFileToCache } from '~/utils';
 
 type TUseContentHandler = {
   setMessages: (messages: TMessage[]) => void;
@@ -20,6 +25,7 @@ type TContentHandler = {
 };
 
 export default function useContentHandler({ setMessages, getMessages }: TUseContentHandler) {
+  const queryClient = useQueryClient();
   const messageMap = useMemo(() => new Map<string, TMessage>(), []);
   return useCallback(
     ({ data, submission }: TContentHandler) => {
@@ -47,10 +53,14 @@ export default function useContentHandler({ setMessages, getMessages }: TUseCont
       }
 
       // TODO: handle streaming for non-text
-      const textPart: Text | string = data[ContentTypes.TEXT];
+      const textPart: Text | string | undefined = data[ContentTypes.TEXT];
       const part: ContentPart =
         textPart && typeof textPart === 'string' ? { value: textPart } : data[type];
 
+      if (type === ContentTypes.IMAGE_FILE) {
+        addFileToCache(queryClient, part as ImageFile & PartMetadata);
+      }
+
       /* spreading the content array to avoid mutation */
       response.content = [...(response.content ?? [])];
 
@@ -68,6 +78,6 @@ export default function useContentHandler({ setMessages, getMessages }: TUseCont
 
       setMessages([...messages, response]);
     },
-    [getMessages, messageMap, setMessages],
+    [queryClient, getMessages, messageMap, setMessages],
   );
 }

client/src/utils/files.ts

@@ -1,4 +1,6 @@
-import { excelMimeTypes } from 'librechat-data-provider';
+import { excelMimeTypes, QueryKeys } from 'librechat-data-provider';
+import type { QueryClient } from '@tanstack/react-query';
+import type { TFile } from 'librechat-data-provider';
 import SheetPaths from '~/components/svg/Files/SheetPaths';
 import TextPaths from '~/components/svg/Files/TextPaths';
 import FilePaths from '~/components/svg/Files/FilePaths';
@@ -128,3 +130,32 @@ export function formatDate(dateString) {
 
   return `${day} ${month} ${year}`;
 }
+
+/**
+ * Adds a file to the query cache
+ */
+export function addFileToCache(queryClient: QueryClient, newfile: TFile) {
+  const currentFiles = queryClient.getQueryData<TFile[]>([QueryKeys.files]);
+
+  if (!currentFiles) {
+    console.warn('No current files found in cache, skipped updating file query cache');
+    return;
+  }
+
+  const fileIndex = currentFiles.findIndex((file) => file.file_id === newfile.file_id);
+
+  if (fileIndex > -1) {
+    console.warn('File already exists in cache, skipped updating file query cache');
+    return;
+  }
+
+  queryClient.setQueryData<TFile[]>(
+    [QueryKeys.files],
+    [
+      {
+        ...newfile,
+      },
+      ...currentFiles,
+    ],
+  );
+}