From c9e7d4ac18afee96fa2b13506592ab92bf6577b4 Mon Sep 17 00:00:00 2001
From: Yuichi Oneda <ohneda@gmail.com>
Date: Tue, 28 May 2024 15:10:33 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=91=20fix(export):=20Export=20Issue=20?=
 =?UTF-8?q?with=20New=20Chat=20(#2777)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 🚑 fix: re-fetch messages when exporting

* Revert "🚑 fix: re-fetch messages when exporting"

This reverts commit 693b86e955093a346defeb0e581a53fe5ad11592.

* 🚑 fix: use the same logic to get export data as useChatHelper

* refactor(useExportConversation): use query cache to build messages tree on request

* chore: organize imports

---------

Co-authored-by: Danny Avila <danny@librechat.ai>
---
 .../Conversations/useExportConversation.ts    | 34 +++++++++++--------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/client/src/hooks/Conversations/useExportConversation.ts b/client/src/hooks/Conversations/useExportConversation.ts
index 5c9750e4da..8cd2133d06 100644
--- a/client/src/hooks/Conversations/useExportConversation.ts
+++ b/client/src/hooks/Conversations/useExportConversation.ts
@@ -1,7 +1,9 @@
 import download from 'downloadjs';
+import { useCallback } from 'react';
 import exportFromJSON from 'export-from-json';
-import { useGetMessagesByConvoId } from 'librechat-data-provider/react-query';
+import { useQueryClient } from '@tanstack/react-query';
 import {
+  QueryKeys,
   ContentTypes,
   ToolCallTypes,
   imageGenTools,
@@ -16,6 +18,7 @@ import type {
 import useBuildMessageTree from '~/hooks/Messages/useBuildMessageTree';
 import { useScreenshot } from '~/hooks/ScreenshotContext';
 import { cleanupPreset, buildTree } from '~/utils';
+import { useParams } from 'react-router-dom';
 
 export default function useExportConversation({
   conversation,
@@ -32,24 +35,25 @@ export default function useExportConversation({
   exportBranches: boolean | 'indeterminate';
   recursive: boolean | 'indeterminate';
 }) {
+  const queryClient = useQueryClient();
   const { captureScreenshot } = useScreenshot();
   const buildMessageTree = useBuildMessageTree();
-  const { data: messagesTree = null } = useGetMessagesByConvoId(
-    conversation?.conversationId ?? '',
-    {
-      select: (data) => {
-        const dataTree = buildTree({ messages: data });
-        return dataTree?.length === 0 ? null : dataTree ?? null;
-      },
-    },
-  );
+
+  const { conversationId: paramId } = useParams();
+
+  const getMessageTree = useCallback(() => {
+    const queryParam = paramId === 'new' ? paramId : conversation?.conversationId ?? paramId ?? '';
+    const messages = queryClient.getQueryData<TMessage[]>([QueryKeys.messages, queryParam]) ?? [];
+    const dataTree = buildTree({ messages });
+    return dataTree?.length === 0 ? null : dataTree ?? null;
+  }, [paramId, conversation?.conversationId, queryClient]);
 
   const getMessageText = (message: TMessage, format = 'text') => {
     if (!message) {
       return '';
     }
 
-    const formatText = (sender, text) => {
+    const formatText = (sender: string, text: string) => {
       if (format === 'text') {
         return `>> ${sender}:\n${text}`;
       }
@@ -149,7 +153,7 @@ export default function useExportConversation({
     const messages = await buildMessageTree({
       messageId: conversation?.conversationId,
       message: null,
-      messages: messagesTree,
+      messages: getMessageTree(),
       branches: !!exportBranches,
       recursive: false,
     });
@@ -224,7 +228,7 @@ export default function useExportConversation({
     const messages = await buildMessageTree({
       messageId: conversation?.conversationId,
       message: null,
-      messages: messagesTree,
+      messages: getMessageTree(),
       branches: false,
       recursive: false,
     });
@@ -280,7 +284,7 @@ export default function useExportConversation({
     const messages = await buildMessageTree({
       messageId: conversation?.conversationId,
       message: null,
-      messages: messagesTree,
+      messages: getMessageTree(),
       branches: false,
       recursive: false,
     });
@@ -332,7 +336,7 @@ export default function useExportConversation({
     const messages = await buildMessageTree({
       messageId: conversation?.conversationId,
       message: null,
-      messages: messagesTree,
+      messages: getMessageTree(),
       branches: !!exportBranches,
       recursive: !!recursive,
     });