Add optional configurable message footer (jupyterlab#942)

* add configurable message footer token

* add custom footer docs

docs/source/developers/index.md

@@ -391,3 +391,50 @@ custom = "custom_package:CustomChatHandler"
 
 Then, install your package so that Jupyter AI adds custom chat handlers
 to the existing chat handlers.
+
+## Custom message footer
+
+You can provide a custom message footer that will be rendered under each message
+in the UI. To do so, you need to write or install a labextension containing a
+plugin that provides the `IJaiMessageFooter` token. This plugin should return a
+`IJaiMessageFooter` object, which defines the custom footer to be rendered.
+
+The `IJaiMessageFooter` object contains a single property `component`, which
+should reference a React component that defines the custom message footer.
+Jupyter AI will render this component under each chat message, passing the
+component a `message` prop with the definition of each chat message as an
+object. The `message` prop takes the type `AiService.ChatMessage`, where
+`AiService` is imported from `@jupyter-ai/core/handler`.
+
+Here is a reference plugin that shows some custom text under each agent message:
+
+```tsx
+import React from 'react';
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+import { IJaiMessageFooter, IJaiMessageFooterProps } from '@jupyter-ai/core/tokens';
+
+export const footerPlugin: JupyterFrontEndPlugin<IJaiMessageFooter> = {
+  id: '@your-org/your-package:custom-footer',
+  autoStart: true,
+  requires: [],
+  provides: IJaiMessageFooter,
+  activate: (app: JupyterFrontEnd): IJaiMessageFooter => {
+    return {
+      component: MessageFooter
+    };
+  }
+};
+
+function MessageFooter(props: IJaiMessageFooterProps) {
+  if (props.message.type !== 'agent' && props.message.type !== 'agent-stream') {
+    return null;
+  }
+
+  return (
+    <div>This is a test footer that renders under each agent message.</div>
+  );
+}
+```

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

@@ -10,10 +10,12 @@ import { AiService } from '../handler';
 import { RendermimeMarkdown } from './rendermime-markdown';
 import { useCollaboratorsContext } from '../contexts/collaborators-context';
 import { ChatMessageMenu } from './chat-messages/chat-message-menu';
+import { IJaiMessageFooter } from '../tokens';
 
 type ChatMessagesProps = {
   rmRegistry: IRenderMimeRegistry;
   messages: AiService.ChatMessage[];
+  messageFooter: IJaiMessageFooter | null;
 };
 
 type ChatMessageHeaderProps = {
@@ -215,6 +217,9 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
                 message.type === 'agent-stream' ? !!message.complete : true
               }
             />
+            {props.messageFooter && (
+              <props.messageFooter.component message={message} />
+            )}
           </Box>
         );
       })}

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

@@ -18,7 +18,7 @@ import { SelectionContextProvider } from '../contexts/selection-context';
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
 import { CollaboratorsContextProvider } from '../contexts/collaborators-context';
-import { IJaiCompletionProvider } from '../tokens';
+import { IJaiCompletionProvider, IJaiMessageFooter } from '../tokens';
 import {
   ActiveCellContextProvider,
   ActiveCellManager
@@ -30,6 +30,7 @@ type ChatBodyProps = {
   setChatView: (view: ChatView) => void;
   rmRegistry: IRenderMimeRegistry;
   focusInputSignal: ISignal<unknown, void>;
+  messageFooter: IJaiMessageFooter | null;
 };
 
 /**
@@ -51,7 +52,8 @@ function ChatBody({
   chatHandler,
   focusInputSignal,
   setChatView: chatViewHandler,
-  rmRegistry: renderMimeRegistry
+  rmRegistry: renderMimeRegistry,
+  messageFooter
 }: ChatBodyProps): JSX.Element {
   const [messages, setMessages] = useState<AiService.ChatMessage[]>([
     ...chatHandler.history.messages
@@ -139,7 +141,11 @@ function ChatBody({
   return (
     <>
       <ScrollContainer sx={{ flexGrow: 1 }}>
-        <ChatMessages messages={messages} rmRegistry={renderMimeRegistry} />
+        <ChatMessages
+          messages={messages}
+          rmRegistry={renderMimeRegistry}
+          messageFooter={messageFooter}
+        />
         <PendingMessages messages={pendingMessages} />
       </ScrollContainer>
       <ChatInput
@@ -170,6 +176,7 @@ export type ChatProps = {
   openInlineCompleterSettings: () => void;
   activeCellManager: ActiveCellManager;
   focusInputSignal: ISignal<unknown, void>;
+  messageFooter: IJaiMessageFooter | null;
 };
 
 enum ChatView {
@@ -223,6 +230,7 @@ export function Chat(props: ChatProps): JSX.Element {
                   setChatView={setView}
                   rmRegistry={props.rmRegistry}
                   focusInputSignal={props.focusInputSignal}
+                  messageFooter={props.messageFooter}
                 />
               )}
               {view === ChatView.Settings && (

packages/jupyter-ai/src/index.ts

@@ -18,7 +18,7 @@ import { ChatHandler } from './chat_handler';
 import { buildErrorWidget } from './widgets/chat-error';
 import { completionPlugin } from './completions';
 import { statusItemPlugin } from './status';
-import { IJaiCompletionProvider } from './tokens';
+import { IJaiCompletionProvider, IJaiMessageFooter } from './tokens';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ActiveCellManager } from './contexts/active-cell-context';
 import { Signal } from '@lumino/signaling';
@@ -42,7 +42,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
     IGlobalAwareness,
     ILayoutRestorer,
     IThemeManager,
-    IJaiCompletionProvider
+    IJaiCompletionProvider,
+    IJaiMessageFooter
   ],
   requires: [IRenderMimeRegistry],
   activate: async (
@@ -51,7 +52,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
     globalAwareness: Awareness | null,
     restorer: ILayoutRestorer | null,
     themeManager: IThemeManager | null,
-    completionProvider: IJaiCompletionProvider | null
+    completionProvider: IJaiCompletionProvider | null,
+    messageFooter: IJaiMessageFooter | null
   ) => {
     /**
      * Initialize selection watcher singleton
@@ -88,7 +90,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
         completionProvider,
         openInlineCompleterSettings,
         activeCellManager,
-        focusInputSignal
+        focusInputSignal,
+        messageFooter
       );
     } catch (e) {
       chatWidget = buildErrorWidget(themeManager);

packages/jupyter-ai/src/tokens.ts

@@ -1,6 +1,8 @@
+import React from 'react';
 import { Token } from '@lumino/coreutils';
 import { ISignal } from '@lumino/signaling';
 import type { IRankedMenu } from '@jupyterlab/ui-components';
+import { AiService } from './handler';
 
 export interface IJaiStatusItem {
   addItem(item: IRankedMenu.IItemOptions): void;
@@ -26,3 +28,21 @@ export const IJaiCompletionProvider = new Token<IJaiCompletionProvider>(
   'jupyter_ai:IJaiCompletionProvider',
   'The jupyter-ai inline completion provider API'
 );
+
+export type IJaiMessageFooterProps = {
+  message: AiService.ChatMessage;
+};
+
+export interface IJaiMessageFooter {
+  component: React.FC<IJaiMessageFooterProps>;
+}
+
+/**
+ * The message footer provider token. Another extension should provide this
+ * token to add a footer to each message.
+ */
+
+export const IJaiMessageFooter = new Token<IJaiMessageFooter>(
+  'jupyter_ai:IJaiMessageFooter',
+  'Optional component that is used to render a footer on each Jupyter AI chat message, when provided.'
+);

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

@@ -8,7 +8,7 @@ import { Chat } from '../components/chat';
 import { chatIcon } from '../icons';
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
-import { IJaiCompletionProvider } from '../tokens';
+import { IJaiCompletionProvider, IJaiMessageFooter } from '../tokens';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import type { ActiveCellManager } from '../contexts/active-cell-context';
 
@@ -21,7 +21,8 @@ export function buildChatSidebar(
   completionProvider: IJaiCompletionProvider | null,
   openInlineCompleterSettings: () => void,
   activeCellManager: ActiveCellManager,
-  focusInputSignal: ISignal<unknown, void>
+  focusInputSignal: ISignal<unknown, void>,
+  messageFooter: IJaiMessageFooter | null
 ): ReactWidget {
   const ChatWidget = ReactWidget.create(
     <Chat
@@ -34,6 +35,7 @@ export function buildChatSidebar(
       openInlineCompleterSettings={openInlineCompleterSettings}
       activeCellManager={activeCellManager}
       focusInputSignal={focusInputSignal}
+      messageFooter={messageFooter}
     />
   );
   ChatWidget.id = 'jupyter-ai::chat';