From 65fe092a9f30d6be44d67cfa2b510b151ef9ceb6 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Wed, 12 Feb 2025 17:14:20 -0500 Subject: [PATCH 1/2] mito-ai: spike --- .../AiChat/ChatMessage/ChatInput.tsx | 9 ++- .../AiChat/ChatMessage/ChatMessage.tsx | 4 +- .../src/Extensions/AiChat/ChatTaskpane.tsx | 2 + .../src/Extensions/InlineCompleter/index.ts | 1 + mito-ai/src/utils/notebook.tsx | 24 ++++++++ mito-ai/test.py | 61 +++++++++++++++++++ 6 files changed, 99 insertions(+), 2 deletions(-) create mode 100644 mito-ai/test.py diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx index 66d27a1b4..da349f059 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx @@ -3,13 +3,14 @@ import { classNames } from '../../../utils/classNames'; import { IVariableManager } from '../../VariableManager/VariableManagerPlugin'; import ChatDropdown from './ChatDropdown'; import { Variable } from '../../VariableManager/VariableInspector'; -import { getActiveCellID, getCellCodeByID } from '../../../utils/notebook'; +import { getActiveCellID, getCellCodeByID, getCellOutputByID } from '../../../utils/notebook'; import { INotebookTracker } from '@jupyterlab/notebook'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import PythonCode from './PythonCode'; import '../../../../style/ChatInput.css'; import '../../../../style/ChatDropdown.css'; import { useDebouncedFunction } from '../../../hooks/useDebouncedFunction'; +import { JupyterFrontEnd } from '@jupyterlab/application'; interface ChatInputProps { initialContent: string; @@ -20,6 +21,7 @@ interface ChatInputProps { variableManager?: IVariableManager; notebookTracker: INotebookTracker; renderMimeRegistry: IRenderMimeRegistry; + app: JupyterFrontEnd } export interface ExpandedVariable extends Variable { @@ -35,6 +37,7 @@ const ChatInput: React.FC = ({ variableManager, notebookTracker, renderMimeRegistry, + app }) => { const [input, setInput] = useState(initialContent); @@ -164,6 +167,10 @@ const ChatInput: React.FC = ({ const activeCellCodePreview = activeCellCode.split('\n').slice(0, 8).join('\n') + ( activeCellCode.split('\n').length > 8 ? '\n\n# Rest of active cell code...' : '') + + const activeCellOutput = getCellOutputByID(app, notebookTracker, activeCellID) + console.log(activeCellOutput) + return (
= ({ rejectAICode, onUpdateMessage, variableManager, - codeReviewStatus + codeReviewStatus, + app }): JSX.Element | null => { const [isEditing, setIsEditing] = useState(false); @@ -94,6 +95,7 @@ const ChatMessage: React.FC = ({ variableManager={variableManager} notebookTracker={notebookTracker} renderMimeRegistry={renderMimeRegistry} + app={app} /> ); } diff --git a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx index 32e667367..e84f8e89d 100644 --- a/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx @@ -104,6 +104,7 @@ const ChatTaskpane: React.FC = ({ })); }; + useEffect(() => { const initializeChatHistory = async () => { try { @@ -801,6 +802,7 @@ const ChatTaskpane: React.FC = ({ variableManager={variableManager} notebookTracker={notebookTracker} renderMimeRegistry={renderMimeRegistry} + app={app} /> {agentModeEnabled && <> diff --git a/mito-ai/src/Extensions/InlineCompleter/index.ts b/mito-ai/src/Extensions/InlineCompleter/index.ts index 7fd2749e5..1c4374693 100644 --- a/mito-ai/src/Extensions/InlineCompleter/index.ts +++ b/mito-ai/src/Extensions/InlineCompleter/index.ts @@ -164,5 +164,6 @@ export const completionPlugin: JupyterFrontEndPlugin = { variableManager }); completionManager.registerInlineProvider(provider); + } }; diff --git a/mito-ai/src/utils/notebook.tsx b/mito-ai/src/utils/notebook.tsx index be0bd0117..e07958e52 100644 --- a/mito-ai/src/utils/notebook.tsx +++ b/mito-ai/src/utils/notebook.tsx @@ -1,6 +1,7 @@ import { INotebookTracker } from '@jupyterlab/notebook'; import { Cell } from '@jupyterlab/cells'; import { removeMarkdownCodeFormatting } from './strings'; +import { JupyterFrontEnd } from '@jupyterlab/application'; export const getActiveCell = (notebookTracker: INotebookTracker): Cell | undefined => { const notebook = notebookTracker.currentWidget?.content; @@ -27,6 +28,29 @@ export const getCellCodeByID = (notebookTracker: INotebookTracker, codeCellID: s return cell?.model.sharedModel.source } +export const getCellOutputByID = async (app: JupyterFrontEnd, notebookTracker: INotebookTracker, codeCellID: string | undefined): Promise => { + if (codeCellID === undefined) { + return undefined + } + + const notebook = notebookTracker.currentWidget?.content; + const cell = notebook?.widgets.find(cell => cell.model.id === codeCellID); + const outputs = (cell?.model.sharedModel as any)?.outputs; + const outputData = outputs?.[0]?.data; + console.log('outputData', outputData) + if (outputData && outputData['image/png']) { + console.log('image/png') + console.log(outputData['image/png']) + } else if (outputData && outputData['text/html']) { + + console.log('text/html') + console.log(outputData['text/html']) + } else { + console.log('no output data') + } + return 'test' +} + export const writeCodeToCellByID = ( notebookTracker: INotebookTracker, code: string | undefined, diff --git a/mito-ai/test.py b/mito-ai/test.py new file mode 100644 index 000000000..2082bf69a --- /dev/null +++ b/mito-ai/test.py @@ -0,0 +1,61 @@ +import base64 +from openai import OpenAI +import imgkit + +html_plotly = """ + + + +""" + +imgkit.from_string(html_plotly, 'out_ABC.png') + + +# # Create the OpenAI client +# client = OpenAI() + +# # Assume you already obtained the Base64 string from your JupyterLab extension. +# # For example, let's say you stored it in the variable base64_image: +# # (Make sure it does NOT include any extra headers—it should be just the encoded data.) +# base64_image = "" + +# # Now, build your message. Notice we use "data:image/png;base64," for a PNG image. +# response = client.chat.completions.create( +# model="gpt-4o-mini", +# messages=[ +# { +# "role": "user", +# "content": [ +# { +# "type": "text", +# "text": "What is in this image?", +# }, +# { +# "type": "image_url", +# "image_url": {"url": f"data:image/png;base64,{base64_image}"}, +# }, +# ], +# } +# ], +# ) + +# print(response.choices[0]) \ No newline at end of file From 7cb94a6d51cf44c1af43ec5b11e405fb7a3c70d8 Mon Sep 17 00:00:00 2001 From: Aaron Diamond-Reivich Date: Thu, 13 Feb 2025 12:18:23 -0500 Subject: [PATCH 2/2] mito-ai: explore html2canvas --- mito-ai/src/utils/nodeToPng.tsx | 132 ++++++++++++++++++++++++++++++++ mito-ai/src/utils/notebook.tsx | 24 +++--- package.json | 5 ++ 3 files changed, 148 insertions(+), 13 deletions(-) create mode 100644 mito-ai/src/utils/nodeToPng.tsx create mode 100644 package.json diff --git a/mito-ai/src/utils/nodeToPng.tsx b/mito-ai/src/utils/nodeToPng.tsx new file mode 100644 index 000000000..90ae10322 --- /dev/null +++ b/mito-ai/src/utils/nodeToPng.tsx @@ -0,0 +1,132 @@ +import html2canvas from 'html2canvas'; + +/** + * Captures a DOM element as a PNG image with preserved styles. + * + * This utility creates a high-fidelity screenshot of a DOM element by: + * 1. Cloning the target element + * 2. Preserving all computed styles with !important flags + * 3. Capturing the clone using html2canvas + * 4. Converting the result to a base64-encoded PNG + * + * @param node - The DOM element to capture + * @returns Promise - Base64-encoded PNG data (without data URL prefix) + * @throws Error if node is null or capture fails + * + * @example + * ```typescript + * const base64Image = await captureNode(document.querySelector('.my-element')); + * // Use base64Image... + * ``` + */ +export const captureNode = async (node: HTMLElement): Promise => { + try { + if (!node) { + throw new Error('No node provided'); + } + + // Create an off-screen wrapper to hold our clone + const wrapper = createWrapper(node); + + // Create and prepare the clone + const clone = node.cloneNode(true) as HTMLElement; + preserveStyles(node, clone); + + // Position clone for capture + wrapper.appendChild(clone); + document.body.appendChild(wrapper); + + try { + // Perform the capture + const canvas = await html2canvas(clone, getHtml2CanvasOptions(node)); + return canvas.toDataURL('image/png').split(',')[1]; + } finally { + // Clean up + wrapper.parentNode?.removeChild(wrapper); + } + + } catch (error) { + console.error('Capture failed:', error); + throw error; + } +}; + +/** + * Creates an off-screen wrapper element to contain the cloned node. + * + * @param node - Reference node to size the wrapper + * @returns HTMLDivElement - Configured wrapper element + */ +const createWrapper = (node: HTMLElement): HTMLDivElement => { + const wrapper = document.createElement('div'); + wrapper.style.cssText = ` + position: fixed !important; + top: 0 !important; + left: 0 !important; + width: ${node.offsetWidth}px !important; + height: ${node.offsetHeight}px !important; + z-index: -9999 !important; + background: transparent !important; + pointer-events: none !important; + opacity: 0 !important; + `; + return wrapper; +}; + +/** + * Recursively copies all computed styles from a source element to a target element. + * Adds !important to all styles to ensure they're preserved during capture. + * + * @param sourceElement - Element to copy styles from + * @param targetElement - Element to copy styles to + */ +const preserveStyles = (sourceElement: HTMLElement, targetElement: HTMLElement): void => { + // Copy computed styles + const computed = window.getComputedStyle(sourceElement); + let stylesText = ''; + + for (let i = 0; i < computed.length; i++) { + const property = computed[i]; + const value = computed.getPropertyValue(property); + if (value) { + stylesText += `${property}: ${value} !important; `; + } + } + + // Apply styles to target + targetElement.style.cssText += stylesText; + + // Process children recursively + Array.from(sourceElement.children).forEach((sourceChild, index) => { + const targetChild = targetElement.children[index]; + if (sourceChild instanceof HTMLElement && targetChild instanceof HTMLElement) { + preserveStyles(sourceChild, targetChild); + } + }); +}; + +/** + * Configures html2canvas options for optimal capture. + * + * @param node - Reference node for dimensioning + * @returns html2canvas configuration object + */ +const getHtml2CanvasOptions = (node: HTMLElement) => ({ + scale: window.devicePixelRatio, + useCORS: true, + logging: false, + allowTaint: true, + backgroundColor: null, + removeContainer: false, + foreignObjectRendering: true, + width: node.offsetWidth, + height: node.offsetHeight, + onclone: (document: Document) => { + // Re-apply styles to cloned element + const clonedElement = document.body.querySelector('*[data-html2canvas-clone="true"]'); + const originalElement = node.querySelector('*[data-html2canvas-clone="true"]'); + if (clonedElement instanceof HTMLElement && originalElement instanceof HTMLElement) { + preserveStyles(originalElement, clonedElement); + } + } +}); \ No newline at end of file diff --git a/mito-ai/src/utils/notebook.tsx b/mito-ai/src/utils/notebook.tsx index e07958e52..526b6c07c 100644 --- a/mito-ai/src/utils/notebook.tsx +++ b/mito-ai/src/utils/notebook.tsx @@ -1,7 +1,9 @@ import { INotebookTracker } from '@jupyterlab/notebook'; -import { Cell } from '@jupyterlab/cells'; +import { Cell, CodeCell } from '@jupyterlab/cells'; import { removeMarkdownCodeFormatting } from './strings'; import { JupyterFrontEnd } from '@jupyterlab/application'; +import { captureNode } from './nodeToPng'; + export const getActiveCell = (notebookTracker: INotebookTracker): Cell | undefined => { const notebook = notebookTracker.currentWidget?.content; @@ -35,18 +37,14 @@ export const getCellOutputByID = async (app: JupyterFrontEnd, notebookTracker: I const notebook = notebookTracker.currentWidget?.content; const cell = notebook?.widgets.find(cell => cell.model.id === codeCellID); - const outputs = (cell?.model.sharedModel as any)?.outputs; - const outputData = outputs?.[0]?.data; - console.log('outputData', outputData) - if (outputData && outputData['image/png']) { - console.log('image/png') - console.log(outputData['image/png']) - } else if (outputData && outputData['text/html']) { - - console.log('text/html') - console.log(outputData['text/html']) - } else { - console.log('no output data') + if (cell instanceof CodeCell) { + console.log('cell', cell.outputArea?.node); + const outputNode = cell.outputArea?.node; + if (outputNode) { + console.log('outputNode', outputNode) + const image = await captureNode(outputNode); + console.log('image', image) + } } return 'test' } diff --git a/package.json b/package.json new file mode 100644 index 000000000..95c9e363b --- /dev/null +++ b/package.json @@ -0,0 +1,5 @@ +{ + "dependencies": { + "html2canvas": "^1.4.1" + } +}