feat: download image of workflow

PySpur-Dev · Feb 13, 2025 · 83601f7 · 83601f7
1 parent e26b061
commit 83601f7
Show file tree

Hide file tree

Showing 6 changed files with 139 additions and 7 deletions.
diff --git a/docs/introduction.mdx b/docs/introduction.mdx
@@ -15,7 +15,7 @@ description: 'PySpur is a Python based AI Agents Builder with Graph UI.'
 
 Pyspur encompasses three design principles:
 
-## 1. Rapid Feedback Loops:
+## 1. Rapid Feedback Loops
 
 * **Build and test in one place:** Build and test your workflow in real-time within the canvas editor, providing you with instant feedback on your changes.
 * **Vendor abstraction**: Not 20 different nodes for 20 different vendors. For example, a unified LLM call node that is vendor-agnostic, enabling you to jump between models quickly to see which one is best for your use case.

diff --git a/frontend/package-lock.json b/frontend/package-lock.json
diff --git a/frontend/package.json b/frontend/package.json
@@ -43,6 +43,7 @@
         "codemirror": "^6.0.1",
         "date-fns": "^4.1.0",
         "framer-motion": "^11.13.5",
+        "html-to-image": "^1.11.12",
         "install": "^0.13.0",
         "lodash.isequal": "^4.5.0",
         "lucide-react": "^0.446.0",

diff --git a/frontend/src/components/Header.tsx b/frontend/src/components/Header.tsx
@@ -31,6 +31,7 @@ import HelpModal from './modals/HelpModal'
 import RunModal from './modals/RunModal'
 import SettingsCard from './modals/SettingsModal'
 import { RunResponse } from '../types/api_types/runSchemas'
+import { handleDownloadImage } from './canvas/FlowCanvas'
 
 interface HeaderProps {
     activePage: 'dashboard' | 'workflow' | 'evals' | 'trace' | 'rag'
@@ -403,9 +404,29 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId
                             </Dropdown>
                         </NavbarItem>
                         <NavbarItem className="hidden sm:flex">
-                            <Button isIconOnly radius="full" variant="light" onPress={handleDownloadWorkflow}>
-                                <Icon className="text-foreground/60" icon="solar:download-linear" width={24} />
-                            </Button>
+                            <Dropdown>
+                                <DropdownTrigger>
+                                    <Button isIconOnly radius="full" variant="light">
+                                        <Icon className="text-foreground/60" icon="solar:download-linear" width={24} />
+                                    </Button>
+                                </DropdownTrigger>
+                                <DropdownMenu>
+                                    <DropdownItem
+                                        key="download-json-workflow"
+                                        onPress={handleDownloadWorkflow}
+                                        startContent={<Icon className="text-foreground/60" icon="solar:document-text-linear" width={20} />}
+                                    >
+                                        Download JSON
+                                    </DropdownItem>
+                                    <DropdownItem
+                                        key="download-image-workflow"
+                                        onPress={handleDownloadImage}
+                                        startContent={<Icon className="text-foreground/60" icon="solar:gallery-linear" width={20} />}
+                                    >
+                                        Download Image
+                                    </DropdownItem>
+                                </DropdownMenu>
+                            </Dropdown>
                         </NavbarItem>
                         <NavbarItem className="hidden sm:flex">
                             <Button isIconOnly radius="full" variant="light" onPress={handleDeploy}>
@@ -420,9 +441,29 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId
                         justify="end"
                     >
                         <NavbarItem className="hidden sm:flex">
-                            <Button isIconOnly radius="full" variant="light" onPress={handleDownloadTrace}>
-                                <Icon className="text-foreground/60" icon="solar:download-linear" width={24} />
-                            </Button>
+                            <Dropdown>
+                                <DropdownTrigger>
+                                    <Button isIconOnly radius="full" variant="light">
+                                        <Icon className="text-foreground/60" icon="solar:download-linear" width={24} />
+                                    </Button>
+                                </DropdownTrigger>
+                                <DropdownMenu>
+                                    <DropdownItem
+                                        key="download-json-trace"
+                                        onPress={handleDownloadTrace}
+                                        startContent={<Icon className="text-foreground/60" icon="solar:document-text-linear" width={20} />}
+                                    >
+                                        Download JSON
+                                    </DropdownItem>
+                                    <DropdownItem
+                                        key="download-image-trace"
+                                        onPress={handleDownloadImage}
+                                        startContent={<Icon className="text-foreground/60" icon="solar:gallery-linear" width={20} />}
+                                    >
+                                        Download Image
+                                    </DropdownItem>
+                                </DropdownMenu>
+                            </Dropdown>
                         </NavbarItem>
                         <NavbarItem>
                             <Link href={`/workflows/${associatedWorkflowId}`}>

diff --git a/frontend/src/components/canvas/FlowCanvas.tsx b/frontend/src/components/canvas/FlowCanvas.tsx
@@ -10,6 +10,9 @@ import {
     ConnectionMode,
     Node,
     useReactFlow,
+    getNodesBounds,
+    getViewportForBounds,
+    Viewport,
 } from '@xyflow/react'
 import '@xyflow/react/dist/style.css'
 import { useSelector, useDispatch } from 'react-redux'
@@ -39,9 +42,78 @@ import { onNodeDragOverGroupNode, onNodeDragStopOverGroupNode } from '../nodes/l
 import { MouseEvent as ReactMouseEvent } from 'react'
 import { throttle } from 'lodash'
 import { Icon } from '@iconify/react'
+import { toPng } from 'html-to-image'
 
 // Type definitions
 
+// Create a utility function for downloading the image
+const downloadImage = (dataUrl: string): void => {
+    const a = document.createElement('a')
+    a.href = dataUrl
+    a.download = 'reactflow.png'
+    a.click()
+}
+
+// Update the handleDownloadImage function to use ReactFlow instance
+export const handleDownloadImage = (): void => {
+    // Get the ReactFlow instance
+    const flow = document.querySelector('.react-flow') as HTMLElement
+    const viewportEl = document.querySelector('.react-flow__viewport') as HTMLElement
+    if (!flow || !viewportEl) {
+        console.error('Unable to locate the flow canvas elements!')
+        return
+    }
+
+    // Get the flow dimensions
+    const flowBounds = flow.getBoundingClientRect()
+    const imageWidth = flowBounds.width
+    const imageHeight = flowBounds.height
+
+    // Get all nodes from the DOM to calculate bounds
+    const nodeElements = document.querySelectorAll('.react-flow__node')
+    const nodes = Array.from(nodeElements).map(el => {
+        const bounds = el.getBoundingClientRect()
+        return {
+            id: el.getAttribute('data-id') || '',
+            position: {
+                x: bounds.x - flowBounds.x,
+                y: bounds.y - flowBounds.y
+            },
+            width: bounds.width,
+            height: bounds.height,
+            data: {},
+            type: 'default'
+        } as Node
+    })
+
+    // Calculate the bounds and viewport transform
+    const nodesBounds = getNodesBounds(nodes)
+    const transform = getViewportForBounds(
+        nodesBounds,
+        imageWidth,
+        imageHeight,
+        0.5,
+        2,
+        0 // Adding the missing padding parameter
+    )
+
+    // Generate the image with the calculated transform
+    toPng(viewportEl, {
+        backgroundColor: '#1a365d',
+        width: imageWidth,
+        height: imageHeight,
+        style: {
+            width: `${imageWidth}px`,
+            height: `${imageHeight}px`,
+            transform: `translate(${transform.x}px, ${transform.y}px) scale(${transform.zoom})`,
+        },
+    })
+    .then(downloadImage)
+    .catch((err) => {
+        console.error('Failed to download image', err)
+    })
+}
+
 interface FlowCanvasProps {
     workflowData?: WorkflowCreateRequest
     workflowID?: string

diff --git a/frontend/src/types/html-to-image.d.ts b/frontend/src/types/html-to-image.d.ts
@@ -0,0 +1,11 @@
+declare module 'html-to-image' {
+    export function toPng(
+        node: HTMLElement,
+        options?: {
+            width?: number;
+            height?: number;
+            backgroundColor?: string;
+            style?: Record<string, string>;
+        }
+    ): Promise<string>;
+}