add rendermime

packages/jupyter-ai/src/components/chat-messages.tsx

@@ -5,11 +5,13 @@ import type { SxProps, Theme } from '@mui/material';
 import 'katex/dist/katex.min.css';
 
 import { AiService } from '../handler';
-import { useCollaboratorsContext } from '../contexts/collaborators-context';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { Jupyternaut } from '../icons';
 import { MarkdownComponent } from './markdown-component';
+import { useCollaboratorsContext } from '../contexts/collaborators-context';
 
 type ChatMessagesProps = {
+  rendermime: IRenderMimeRegistry;
   messages: AiService.ChatMessage[];
 };
 
@@ -19,11 +21,6 @@ type ChatMessageHeaderProps = {
   sx?: SxProps<Theme>;
 };
 
-type NewTabLinkProps = {
-  children: React.ReactNode;
-  href?: string;
-};
-
 export function ChatMessageHeader(props: ChatMessageHeaderProps): JSX.Element {
   const collaborators = useCollaboratorsContext();
 
@@ -130,14 +127,6 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
     }
   }, [props.messages]);
 
-  function NewTabLink(props: NewTabLinkProps) {
-    return (
-      <a href={props.href ?? '#'} target="_blank" rel="noopener noreferrer">
-        {props.children}
-      </a>
-    );
-  }
-
   return (
     <Box
       sx={{
@@ -155,7 +144,7 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
             sx={{ marginBottom: 3 }}
           />
           <MarkdownComponent
-            rendermime={yourRendermimeInstance}
+            rendermime={props.rendermime}
             markdownString={message.body}
           />
         </Box>

packages/jupyter-ai/src/components/chat.tsx

@@ -17,16 +17,19 @@ import {
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
 import { CollaboratorsContextProvider } from '../contexts/collaborators-context';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ScrollContainer } from './scroll-container';
 
 type ChatBodyProps = {
   chatHandler: ChatHandler;
   setChatView: (view: ChatView) => void;
+  renderMimeRegistry: IRenderMimeRegistry;
 };
 
 function ChatBody({
   chatHandler,
-  setChatView: chatViewHandler
+  setChatView: chatViewHandler,
+  renderMimeRegistry
 }: ChatBodyProps): JSX.Element {
   const [messages, setMessages] = useState<AiService.ChatMessage[]>([]);
   const [showWelcomeMessage, setShowWelcomeMessage] = useState<boolean>(false);
@@ -146,7 +149,7 @@ function ChatBody({
   return (
     <>
       <ScrollContainer sx={{ flexGrow: 1 }}>
-        <ChatMessages messages={messages} />
+        <ChatMessages messages={messages} rendermime={renderMimeRegistry} />
       </ScrollContainer>
       <ChatInput
         value={input}
@@ -178,6 +181,7 @@ export type ChatProps = {
   selectionWatcher: SelectionWatcher;
   chatHandler: ChatHandler;
   globalAwareness: Awareness | null;
+  renderMimeRegistry: IRenderMimeRegistry;
   chatView?: ChatView;
 };
 
@@ -224,7 +228,11 @@ export function Chat(props: ChatProps): JSX.Element {
             </Box>
             {/* body */}
             {view === ChatView.Chat && (
-              <ChatBody chatHandler={props.chatHandler} setChatView={setView} />
+              <ChatBody
+                chatHandler={props.chatHandler}
+                setChatView={setView}
+                renderMimeRegistry={props.renderMimeRegistry}
+              />
             )}
             {view === ChatView.Settings && <ChatSettings />}
           </Box>

packages/jupyter-ai/src/components/lumino-component.tsx

@@ -1,25 +1,39 @@
 import React, { useRef, useEffect } from 'react';
+import { MessageLoop } from '@lumino/messaging';
 import { Widget } from '@lumino/widgets';
 
 type LuminoComponentProps = {
   widget: Widget;
 };
 
-export function LuminoComponent(
-  props: LuminoComponentProps
-): React.ReactElement {
+export function LuminoComponent({
+  widget
+}: LuminoComponentProps): React.ReactElement {
   const ref = useRef<HTMLDivElement>(null);
 
   useEffect(() => {
-    if (ref.current) {
-      Widget.attach(props.widget, ref.current);
+    if (!ref.current) {
+      return;
+    }
+
+    try {
+      MessageLoop.sendMessage(widget, Widget.Msg.BeforeAttach);
+      ref.current.appendChild(widget.node);
+      MessageLoop.sendMessage(widget, Widget.Msg.AfterAttach);
+    } catch (e) {
+      console.warn('Exception while attaching Lumino widget:', e);
     }
 
     return () => {
-      // Detach the widget when the component unmounts
-      Widget.detach(props.widget);
+      try {
+        if (widget.isAttached || widget.node.isConnected) {
+          Widget.detach(widget);
+        }
+      } catch (e) {
+        console.warn('Exception while detaching Lumino widget:', e);
+      }
     };
-  }, [props.widget]);
+  }, [widget]);
 
   return <div ref={ref} />;
 }

packages/jupyter-ai/src/components/markdown-component.tsx

@@ -1,5 +1,4 @@
 import React, { useRef, useEffect } from 'react';
-
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { Widget } from '@lumino/widgets';
 import { LuminoComponent } from './lumino-component';
@@ -14,18 +13,14 @@ type MarkdownComponentProps = {
   rendermime: IRenderMimeRegistry;
 };
 
-export class MarkdownWidget extends Widget {
+class MarkdownWidget extends Widget {
   private rendermime: IRenderMimeRegistry;
   private markdownString: string;
 
   constructor(props: MarkdownWidgetProps) {
     super();
     this.rendermime = props.rendermime;
     this.markdownString = props.markdownString;
-    this.onAfterAttach = this.onAfterAttach.bind(this);
-  }
-
-  onAfterAttach(): void {
     this.initializeMarkdownRendering();
   }
 
@@ -38,35 +33,23 @@ export class MarkdownWidget extends Widget {
     await renderer.renderModel(model);
     this.node.appendChild(renderer.node);
   }
-
-  async updateMarkdown(markdownString: string): Promise<void> {
-    this.markdownString = markdownString;
-    // Clear the existing content
-    while (this.node.firstChild) {
-      this.node.removeChild(this.node.firstChild);
-    }
-    // Reinitialize rendering with the new markdown string
-    await this.initializeMarkdownRendering();
-  }
 }
 
-export function MarkdownComponent({
-  markdownString,
-  rendermime
-}: MarkdownComponentProps): React.ReactElement | null {
+export function MarkdownComponent(
+  props: MarkdownComponentProps
+): React.ReactElement | null {
+  const { markdownString, rendermime } = props;
   const widgetRef = useRef<MarkdownWidget | null>(null);
 
   useEffect(() => {
     if (!widgetRef.current) {
       widgetRef.current = new MarkdownWidget({ rendermime, markdownString });
-    } else {
-      widgetRef.current.updateMarkdown(markdownString);
     }
 
     return () => {
       widgetRef.current?.dispose();
     };
-  }, [markdownString, rendermime]);
+  }, []); // Empty dependency array if props are not expected to change
 
   return widgetRef.current ? (
     <LuminoComponent widget={widgetRef.current} />

packages/jupyter-ai/src/index.ts

@@ -12,6 +12,7 @@ import { buildChatSidebar } from './widgets/chat-sidebar';
 import { SelectionWatcher } from './selection-watcher';
 import { ChatHandler } from './chat_handler';
 import { buildErrorWidget } from './widgets/chat-error';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 
 export type DocumentTracker = IWidgetTracker<IDocumentWidget>;
 
@@ -25,7 +26,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
   activate: async (
     app: JupyterFrontEnd,
     globalAwareness: Awareness | null,
-    restorer: ILayoutRestorer | null
+    restorer: ILayoutRestorer | null,
+    renderMimeRegistry: IRenderMimeRegistry
   ) => {
     /**
      * Initialize selection watcher singleton
@@ -43,7 +45,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
       chatWidget = buildChatSidebar(
         selectionWatcher,
         chatHandler,
-        globalAwareness
+        globalAwareness,
+        renderMimeRegistry
       );
     } catch (e) {
       chatWidget = buildErrorWidget();

packages/jupyter-ai/src/widgets/chat-sidebar.tsx

@@ -6,17 +6,20 @@ import { Chat } from '../components/chat';
 import { chatIcon } from '../icons';
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 
 export function buildChatSidebar(
   selectionWatcher: SelectionWatcher,
   chatHandler: ChatHandler,
-  globalAwareness: Awareness | null
+  globalAwareness: Awareness | null,
+  renderMimeRegistry: IRenderMimeRegistry
 ): ReactWidget {
   const ChatWidget = ReactWidget.create(
     <Chat
       selectionWatcher={selectionWatcher}
       chatHandler={chatHandler}
       globalAwareness={globalAwareness}
+      renderMimeRegistry={renderMimeRegistry}
     />
   );
   ChatWidget.id = 'jupyter-ai::chat';