From 96914387a6cd774d7745f2308dbc5484f842a43a Mon Sep 17 00:00:00 2001
From: Wentao Lyu <35-wentao.lyu@users.noreply.git.stereye.tech>
Date: Thu, 6 Apr 2023 02:48:32 +0800
Subject: [PATCH] feat: export to screenshot

---
 client/package.json                           |  3 +
 client/src/App.jsx                            |  8 ++-
 client/src/components/Messages/index.jsx      | 10 +++-
 .../Nav/ExportConversation/ExportModel.jsx    | 56 +++++++++++--------
 client/src/utils/screenshotContext.jsx        | 21 +++++++
 5 files changed, 73 insertions(+), 25 deletions(-)
 create mode 100644 client/src/utils/screenshotContext.jsx

diff --git a/client/package.json b/client/package.json
index f0f477bac..27227c2be 100644
--- a/client/package.json
+++ b/client/package.json
@@ -40,8 +40,10 @@
     "clsx": "^1.2.1",
     "copy-to-clipboard": "^3.3.3",
     "crypto-browserify": "^3.12.0",
+    "downloadjs": "^1.4.7",
     "esbuild": "0.17.15",
     "export-from-json": "^1.7.2",
+    "html2canvas": "^1.4.1",
     "lodash": "^4.17.21",
     "lucide-react": "^0.113.0",
     "rc-input-number": "^7.4.2",
@@ -65,6 +67,7 @@
     "tailwindcss-animate": "^1.0.5",
     "tailwindcss-radix": "^2.8.0",
     "url": "^0.11.0",
+    "use-react-screenshot": "^3.0.0",
     "uuidv4": "^6.2.13"
   },
   "devDependencies": {
diff --git a/client/src/App.jsx b/client/src/App.jsx
index b63e47467..cca1de99a 100644
--- a/client/src/App.jsx
+++ b/client/src/App.jsx
@@ -7,6 +7,8 @@ import store from './store';
 import userAuth from './utils/userAuth';
 import { useRecoilState, useSetRecoilState } from 'recoil';
 
+import { ScreenshotProvider } from './utils/screenshotContext.jsx';
+
 import axios from 'axios';
 
 const router = createBrowserRouter([
@@ -97,4 +99,8 @@ const App = () => {
   else return <div className="flex h-screen"></div>;
 };
 
-export default App;
+export default () => (
+  <ScreenshotProvider>
+    <App />
+  </ScreenshotProvider>
+);
diff --git a/client/src/components/Messages/index.jsx b/client/src/components/Messages/index.jsx
index 3a7f95c3d..dae528eb9 100644
--- a/client/src/components/Messages/index.jsx
+++ b/client/src/components/Messages/index.jsx
@@ -6,6 +6,7 @@ import { CSSTransition } from 'react-transition-group';
 import ScrollToBottom from './ScrollToBottom';
 import MultiMessage from './MultiMessage';
 import MessageHeader from './MessageHeader';
+import { useScreenshot } from '~/utils/screenshotContext.jsx';
 
 import store from '~/store';
 
@@ -23,6 +24,8 @@ export default function Messages({ isSearchView = false }) {
   const conversation = useRecoilValue(store.conversation) || {};
   const { conversationId } = conversation;
 
+  const { screenshotTargetRef } = useScreenshot();
+
   // const models = useRecoilValue(store.models) || [];
   // const modelName = models.find(element => element.model == model)?.name;
 
@@ -84,8 +87,11 @@ export default function Messages({ isSearchView = false }) {
       ref={scrollableRef}
       onScroll={debouncedHandleScroll}
     >
-      <div className="dark:gpt-dark-gray h-full">
-        <div className="dark:gpt-dark-gray flex h-full flex-col items-center text-sm">
+      <div
+        className="dark:gpt-dark-gray h-auto"
+        ref={screenshotTargetRef}
+      >
+        <div className="dark:gpt-dark-gray flex h-auto flex-col items-center text-sm">
           <MessageHeader isSearchView={isSearchView} />
           {_messagesTree === null ? (
             <Spinner />
diff --git a/client/src/components/Nav/ExportConversation/ExportModel.jsx b/client/src/components/Nav/ExportConversation/ExportModel.jsx
index dfcc9cba2..b4c0080c8 100644
--- a/client/src/components/Nav/ExportConversation/ExportModel.jsx
+++ b/client/src/components/Nav/ExportConversation/ExportModel.jsx
@@ -1,6 +1,7 @@
 import React, { useEffect, useState } from 'react';
 import { useRecoilValue, useRecoilCallback } from 'recoil';
 import exportFromJSON from 'export-from-json';
+import download from 'downloadjs';
 import DialogTemplate from '~/components/ui/DialogTemplate.jsx';
 import { Dialog, DialogClose, DialogButton } from '~/components/ui/Dialog.tsx';
 import { Input } from '~/components/ui/Input.tsx';
@@ -8,11 +9,14 @@ import { Label } from '~/components/ui/Label.tsx';
 import { Checkbox } from '~/components/ui/Checkbox.tsx';
 import Dropdown from '~/components/ui/Dropdown';
 import { cn } from '~/utils/';
+import { useScreenshot } from '~/utils/screenshotContext';
 
 import store from '~/store';
 import cleanupPreset from '~/utils/cleanupPreset.js';
 
 export default function ExportModel({ open, onOpenChange }) {
+  const { captureScreenshot } = useScreenshot();
+
   const [filename, setFileName] = useState('');
   const [type, setType] = useState('');
 
@@ -32,7 +36,7 @@ export default function ExportModel({ open, onOpenChange }) {
     []
   );
 
-  const typeOptions = ['text', 'markdown', 'csv', 'json']; //, 'screenshot', 'webpage'];
+  const typeOptions = ['text', 'markdown', 'csv', 'json', 'screenshot']; //,, 'webpage'];
 
   useEffect(() => {
     setFileName(
@@ -105,6 +109,11 @@ export default function ExportModel({ open, onOpenChange }) {
     }
   };
 
+  const exportScreenshot = async () => {
+    const data = await captureScreenshot();
+    download(data, `${filename}.png`, 'image/png');
+  };
+
   const exportCSV = async () => {
     let data = [];
 
@@ -256,6 +265,7 @@ export default function ExportModel({ open, onOpenChange }) {
     else if (type == 'text') exportText();
     else if (type == 'markdown') exportMarkdown();
     else if (type == 'csv') exportCSV();
+    else if (type == 'screenshot') exportScreenshot();
   };
 
   const defaultTextProps =
@@ -311,7 +321,7 @@ export default function ExportModel({ open, onOpenChange }) {
               </div>
             </div>
             <div className="grid w-full gap-6 sm:grid-cols-2">
-              {type !== 'csv' ? (
+              {type !== 'csv' && type !== 'screenshot' ? (
                 <div className="col-span-1 flex flex-col items-start justify-start gap-2">
                   <div className="grid w-full items-center gap-2">
                     <Label
@@ -337,29 +347,31 @@ export default function ExportModel({ open, onOpenChange }) {
                   </div>
                 </div>
               ) : null}
-              <div className="grid w-full items-center gap-2">
-                <Label
-                  htmlFor="exportBranches"
-                  className="text-left text-sm font-medium"
-                >
-                  Export all message branches
-                </Label>
-                <div className="flex h-[40px] w-full items-center space-x-3">
-                  <Checkbox
-                    id="exportBranches"
-                    disabled={!exportBranchesSupport}
-                    checked={exportBranches}
-                    className="focus:ring-opacity-20 dark:border-gray-500 dark:bg-gray-700 dark:text-gray-50 dark:focus:ring-gray-600 dark:focus:ring-opacity-50 dark:focus:ring-offset-0"
-                    onCheckedChange={setExportBranches}
-                  />
-                  <label
+              {type !== 'screenshot' ? (
+                <div className="grid w-full items-center gap-2">
+                  <Label
                     htmlFor="exportBranches"
-                    className="text-sm font-medium leading-none peer-disabled:cursor-not-allowed peer-disabled:opacity-70 dark:text-gray-50"
+                    className="text-left text-sm font-medium"
                   >
-                    {exportBranchesSupport ? 'Enabled' : 'Not Supported'}
-                  </label>
+                    Export all message branches
+                  </Label>
+                  <div className="flex h-[40px] w-full items-center space-x-3">
+                    <Checkbox
+                      id="exportBranches"
+                      disabled={!exportBranchesSupport}
+                      checked={exportBranches}
+                      className="focus:ring-opacity-20 dark:border-gray-500 dark:bg-gray-700 dark:text-gray-50 dark:focus:ring-gray-600 dark:focus:ring-opacity-50 dark:focus:ring-offset-0"
+                      onCheckedChange={setExportBranches}
+                    />
+                    <label
+                      htmlFor="exportBranches"
+                      className="text-sm font-medium leading-none peer-disabled:cursor-not-allowed peer-disabled:opacity-70 dark:text-gray-50"
+                    >
+                      {exportBranchesSupport ? 'Enabled' : 'Not Supported'}
+                    </label>
+                  </div>
                 </div>
-              </div>
+              ) : null}
               {type === 'json' ? (
                 <div className="grid w-full items-center gap-2">
                   <Label
diff --git a/client/src/utils/screenshotContext.jsx b/client/src/utils/screenshotContext.jsx
new file mode 100644
index 000000000..cc0ecfc64
--- /dev/null
+++ b/client/src/utils/screenshotContext.jsx
@@ -0,0 +1,21 @@
+import React, { createContext, useRef, useContext, useCallback } from 'react';
+import { useScreenshot as useScreenshot_ } from 'use-react-screenshot';
+
+const ScreenshotContext = createContext({});
+
+export const useScreenshot = () => {
+  const { ref } = useContext(ScreenshotContext);
+  const [image, takeScreenshot] = useScreenshot_();
+
+  const captureScreenshot = () => {
+    return takeScreenshot(ref.current);
+  };
+
+  return { screenshotTargetRef: ref, captureScreenshot };
+};
+
+export const ScreenshotProvider = ({ children }) => {
+  const ref = useRef(null);
+
+  return <ScreenshotContext.Provider value={{ ref }}>{children}</ScreenshotContext.Provider>;
+};