[spike] Giving a Cell Output Context

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<ChatInputProps> = ({
     variableManager,
     notebookTracker,
     renderMimeRegistry,
+    app
 }) => {
 
     const [input, setInput] = useState(initialContent);
@@ -164,6 +167,10 @@ const ChatInput: React.FC<ChatInputProps> = ({
     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 (
         <div 
             className={classNames("chat-input-container")}

mito-ai/src/Extensions/AiChat/ChatMessage/ChatMessage.tsx

@@ -58,7 +58,8 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
     rejectAICode,
     onUpdateMessage,
     variableManager,
-    codeReviewStatus
+    codeReviewStatus,
+    app
 }): JSX.Element | null => {
     const [isEditing, setIsEditing] = useState(false);
 
@@ -94,6 +95,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                 variableManager={variableManager}
                 notebookTracker={notebookTracker}
                 renderMimeRegistry={renderMimeRegistry}
+                app={app}
             />
         );
     }

mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx

@@ -104,6 +104,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         }));
     };
 
+
     useEffect(() => {
         const initializeChatHistory = async () => {
           try {
@@ -801,6 +802,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
                         variableManager={variableManager}
                         notebookTracker={notebookTracker}
                         renderMimeRegistry={renderMimeRegistry}
+                        app={app}
                     />
                     {agentModeEnabled &&
                         <>

mito-ai/src/Extensions/InlineCompleter/index.ts

@@ -164,5 +164,6 @@ export const completionPlugin: JupyterFrontEndPlugin<void> = {
       variableManager
     });
     completionManager.registerInlineProvider(provider);
+
   }
 };

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<string> - 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<string> => {
+    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);
+        }
+    }
+});

mito-ai/src/utils/notebook.tsx

@@ -1,6 +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;
@@ -27,6 +30,25 @@ export const getCellCodeByID = (notebookTracker: INotebookTracker, codeCellID: s
     return cell?.model.sharedModel.source
 }
 
+export const getCellOutputByID = async (app: JupyterFrontEnd, notebookTracker: INotebookTracker, codeCellID: string | undefined): Promise<string | undefined> => {
+    if (codeCellID === undefined) {
+        return undefined
+    }
+
+    const notebook = notebookTracker.currentWidget?.content;
+    const cell = notebook?.widgets.find(cell => cell.model.id === codeCellID);
+    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'
+}
+
 export const writeCodeToCellByID = (
     notebookTracker: INotebookTracker, 
     code: string | undefined, 

package.json

@@ -0,0 +1,5 @@
+{
+  "dependencies": {
+    "html2canvas": "^1.4.1"
+  }
+}