Merge pull request #2137 from mito-ds/create-notebook

Auto create new notebook

Author: aarondr77

mito-ai/src/Extensions/AiChat/AiChatPlugin.ts

@@ -68,6 +68,7 @@ const AiChatPlugin: JupyterFrontEndPlugin<WidgetTracker> = {
         rendermime,
         contextManager,
         streamlitPreviewManager,
+        documentManager,
       );
       return chatWidget;
     };

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

@@ -11,6 +11,7 @@ import React, { useEffect, useRef } from 'react';
 import { JupyterFrontEnd } from '@jupyterlab/application';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+import { IDocumentManager } from '@jupyterlab/docmanager';
 import { addIcon, historyIcon, deleteIcon, settingsIcon } from '@jupyterlab/ui-components';
 import { ReadonlyPartialJSONObject, UUID } from '@lumino/coreutils';
 
@@ -49,6 +50,7 @@ import {
 import { getActiveCellOutput } from '../../utils/cellOutput';
 import { OperatingSystem } from '../../utils/user';
 import { IStreamlitPreviewManager } from '../AppPreview/StreamlitPreviewPlugin';
+import { ensureNotebookExists } from './utils';
 import { waitForNotebookReady } from '../../utils/waitForNotebookReady';
 import { getBase64EncodedCellOutputInNotebook } from './utils';
 import { logEvent } from '../../restAPI/RestAPI';
@@ -122,6 +124,7 @@ interface IChatTaskpaneProps {
     app: JupyterFrontEnd
     operatingSystem: OperatingSystem
     websocketClient: CompletionWebsocketClient
+    documentManager: IDocumentManager
 }
 
 // Re-export types from hooks for backward compatibility
@@ -142,6 +145,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
     app,
     operatingSystem,
     websocketClient,
+    documentManager,
 }) => {
 
     // User signup state
@@ -374,6 +378,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         // Step 3: No post processing step needed for explaining code. 
     }
 
+
     const sendAgentExecutionMessage = async (
         input: string,
         messageIndex?: number,
@@ -384,12 +389,6 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         // Step 0: reset the state for a new message
         resetForNewMessage()
 
-        const agentTargetNotebookPanel = agentTargetNotebookPanelRef.current
-
-        if (agentTargetNotebookPanel === null) {
-            return
-        }
-
         // Step 1: Add the user's message to the chat history
         const newChatHistoryManager = getDuplicateChatHistoryManager()
 
@@ -400,15 +399,15 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
 
         const agentExecutionMetadata = newChatHistoryManager.addAgentExecutionMessage(
             activeThreadIdRef.current, 
-            agentTargetNotebookPanel,
+            agentTargetNotebookPanelRef.current,
             input,
             additionalContext
         )
         if (messageIndex !== undefined) {
             agentExecutionMetadata.index = messageIndex
         }
 
-        agentExecutionMetadata.base64EncodedActiveCellOutput = await getBase64EncodedCellOutputInNotebook(agentTargetNotebookPanel, sendCellIDOutput)
+        agentExecutionMetadata.base64EncodedActiveCellOutput = await getBase64EncodedCellOutputInNotebook(agentTargetNotebookPanelRef.current, sendCellIDOutput)
 
         setChatHistoryManager(newChatHistoryManager)
         setLoadingStatus('thinking');
@@ -430,6 +429,9 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         // Step 0: reset the state for a new message
         resetForNewMessage()
 
+        // Ensure a notebook exists before proceeding
+        await ensureNotebookExists(notebookTracker, documentManager);
+
         // Enable follow mode when user sends a new message
         setAutoScrollFollowMode(true);
 
@@ -702,6 +704,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
         app,
         streamlitPreviewManager,
         websocketClient,
+        documentManager,
         chatHistoryManagerRef,
         activeThreadIdRef,
         activeRequestControllerRef,

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

@@ -13,6 +13,7 @@ import chatIconSvg from '../../../src/icons/ChatIcon.svg';
 import { IContextManager } from '../ContextManager/ContextManagerPlugin';
 import { IStreamlitPreviewManager } from '../AppPreview/StreamlitPreviewPlugin';
 import { JupyterFrontEnd } from '@jupyterlab/application';
+import { IDocumentManager } from '@jupyterlab/docmanager';
 import { getOperatingSystem, OperatingSystem } from '../../utils/user';
 import type { IChatWidget as IChatWidget } from './token';
 import { Signal, type ISignal } from '@lumino/signaling';
@@ -40,6 +41,7 @@ export class ChatWidget extends ReactWidget implements IChatWidget {
       contextManager: IContextManager;
       streamlitPreviewManager: IStreamlitPreviewManager;
       operatingSystem: OperatingSystem;
+      documentManager: IDocumentManager;
     }
   ) {
     super();
@@ -93,6 +95,7 @@ export class ChatWidget extends ReactWidget implements IChatWidget {
         streamlitPreviewManager={this.options.streamlitPreviewManager}
         operatingSystem={this.options.operatingSystem}
         websocketClient={this.websocketClient}
+        documentManager={this.options.documentManager}
       />
     );
   }
@@ -118,6 +121,7 @@ export function buildChatWidget(
   renderMimeRegistry: IRenderMimeRegistry,
   contextManager: IContextManager,
   streamlitPreviewManager: IStreamlitPreviewManager,
+  documentManager: IDocumentManager,
 ): ChatWidget {
   // Get the operating system here so we don't have to do it each time the chat changes.
   // The operating system won't change, duh.
@@ -129,7 +133,8 @@ export function buildChatWidget(
     renderMimeRegistry,
     contextManager,
     streamlitPreviewManager: streamlitPreviewManager,
-    operatingSystem
+    operatingSystem,
+    documentManager
   });
   chatWidget.id = 'mito_ai';
   return chatWidget;

mito-ai/src/Extensions/AiChat/hooks/useAgentExecution.ts

@@ -6,6 +6,7 @@
 import { useRef, useState } from 'react';
 import { JupyterFrontEnd } from '@jupyterlab/application';
 import { INotebookTracker } from '@jupyterlab/notebook';
+import { IDocumentManager } from '@jupyterlab/docmanager';
 import { UUID } from '@lumino/coreutils';
 import { IStreamlitPreviewManager } from '../../AppPreview/StreamlitPreviewPlugin';
 import { CompletionWebsocketClient } from '../../../websockets/completions/CompletionsWebsocketClient';
@@ -18,6 +19,7 @@ import { getCodeBlockFromMessage } from '../../../utils/strings';
 import { getAIOptimizedCellsInNotebookPanel, setActiveCellByIDInNotebookPanel } from '../../../utils/notebook';
 import { AgentReviewStatus } from '../ChatTaskpane';
 import { LoadingStatus } from './useChatState';
+import { ensureNotebookExists } from '../utils';
 
 export type AgentExecutionStatus = 'working' | 'stopping' | 'idle';
 
@@ -28,6 +30,7 @@ interface UseAgentExecutionProps {
     app: JupyterFrontEnd;
     streamlitPreviewManager: IStreamlitPreviewManager;
     websocketClient: CompletionWebsocketClient;
+    documentManager: IDocumentManager;
     chatHistoryManagerRef: React.MutableRefObject<ChatHistoryManager>;
     activeThreadIdRef: React.MutableRefObject<string>;
     activeRequestControllerRef: React.MutableRefObject<AbortController | null>;
@@ -63,6 +66,7 @@ export const useAgentExecution = ({
     app,
     streamlitPreviewManager,
     websocketClient,
+    documentManager,
     chatHistoryManagerRef,
     activeThreadIdRef,
     activeRequestControllerRef,
@@ -133,7 +137,10 @@ export const useAgentExecution = ({
         messageIndex?: number,
         additionalContext?: Array<{ type: string, value: string }>
     ): Promise<void> => {
-        agentTargetNotebookPanelRef.current = notebookTracker.currentWidget;
+        
+        // Ensure a notebook exists before proceeding with agent execution
+        const agentTargetNotebookPanel = await ensureNotebookExists(notebookTracker, documentManager);
+        agentTargetNotebookPanelRef.current = agentTargetNotebookPanel;
 
         agentReview.acceptAllAICode();
         agentReview.setNotebookSnapshotPreAgentExecution(getAIOptimizedCellsInNotebookPanel(agentTargetNotebookPanelRef.current));

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

@@ -4,9 +4,12 @@
  */
 
 import { INotebookTracker, NotebookPanel } from "@jupyterlab/notebook"
+import { IDocumentManager } from "@jupyterlab/docmanager"
 import { getCellIndexByIDInNotebookPanel } from "../../utils/notebook"
 import { getCellOutputByIDInNotebook } from "../../utils/cellOutput"
 import { logEvent } from "../../restAPI/RestAPI"
+import { waitForNotebookReady } from "../../utils/waitForNotebookReady"
+import { setNotebookID } from "../../utils/notebookMetadata"
 
 export const getBase64EncodedCellOutput = async (notebookTracker: INotebookTracker, cellID: string | undefined): Promise<string | undefined> => {
     const notebookPanel = notebookTracker.currentWidget
@@ -33,3 +36,46 @@ export const getBase64EncodedCellOutputInNotebook = async (notebookPanel: Notebo
 
     return undefined
 }
+
+/* 
+    Ensure a notebook exists. If no notebook is open, create a new one.
+    Returns the notebook panel.
+*/
+export const ensureNotebookExists = async (
+    notebookTracker: INotebookTracker,
+    documentManager: IDocumentManager
+): Promise<NotebookPanel> => {
+    // Check if a notebook already exists
+    let notebookPanel = notebookTracker.currentWidget;
+    
+    if (notebookPanel !== null) {
+        return notebookPanel;
+    }
+
+    // No notebook exists, create a new one
+    try {
+        // Create a new notebook model (Contents.IModel has a path property)
+        const model = await documentManager.newUntitled({ type: 'notebook' });
+        
+        // Open the notebook using the path from the model
+        await documentManager.open(model.path);
+        
+        // Wait for the notebook to appear in the tracker and be ready
+        await waitForNotebookReady(notebookTracker);
+        
+        // Get the notebook panel from the tracker
+        notebookPanel = notebookTracker.currentWidget;
+        
+        if (notebookPanel === null) {
+            throw new Error('Failed to get notebook panel after creation');
+        }
+
+        // Set the notebook ID if it doesn't exist
+        setNotebookID(notebookPanel);
+        
+        return notebookPanel;
+    } catch (error) {
+        console.error('Error creating new notebook:', error);
+        throw error;
+    }
+};