Add support for optional telemetry plugin

packages/jupyter-ai/jupyter_ai/handlers.py

@@ -249,6 +249,7 @@ def broadcast_message(self, message: Message):
                 ):
                     stream_message: AgentStreamMessage = history_message
                     stream_message.body += chunk.content
+                    stream_message.metadata = chunk.metadata
                     stream_message.complete = chunk.stream_complete
                     break
         elif isinstance(message, PendingMessage):

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

@@ -74,10 +74,6 @@ function sortMessages(
 export function ChatMessageHeader(props: ChatMessageHeaderProps): JSX.Element {
   const collaborators = useCollaboratorsContext();
 
-  if (props.message.type === 'agent-stream' && props.message.complete) {
-    console.log(props.message.metadata);
-  }
-
   const sharedStyles: SxProps<Theme> = {
     height: '24px',
     width: '24px'
@@ -228,8 +224,9 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
               sx={{ marginBottom: 3 }}
             />
             <RendermimeMarkdown
-              rmRegistry={props.rmRegistry}
               markdownStr={message.body}
+              rmRegistry={props.rmRegistry}
+              parentMessage={message}
               complete={
                 message.type === 'agent-stream' ? !!message.complete : true
               }

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

@@ -19,13 +19,18 @@ import { SelectionContextProvider } from '../contexts/selection-context';
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
 import { CollaboratorsContextProvider } from '../contexts/collaborators-context';
-import { IJaiCompletionProvider, IJaiMessageFooter } from '../tokens';
+import {
+  IJaiCompletionProvider,
+  IJaiMessageFooter,
+  IJaiTelemetryHandler
+} from '../tokens';
 import {
   ActiveCellContextProvider,
   ActiveCellManager
 } from '../contexts/active-cell-context';
 import { ScrollContainer } from './scroll-container';
 import { TooltippedIconButton } from './mui-extras/tooltipped-icon-button';
+import { TelemetryContextProvider } from '../contexts/telemetry-context';
 
 type ChatBodyProps = {
   chatHandler: ChatHandler;
@@ -178,6 +183,7 @@ export type ChatProps = {
   activeCellManager: ActiveCellManager;
   focusInputSignal: ISignal<unknown, void>;
   messageFooter: IJaiMessageFooter | null;
+  telemetryHandler: IJaiTelemetryHandler | null;
 };
 
 enum ChatView {
@@ -201,69 +207,71 @@ export function Chat(props: ChatProps): JSX.Element {
           <ActiveCellContextProvider
             activeCellManager={props.activeCellManager}
           >
-            <Box
-              // root box should not include padding as it offsets the vertical
-              // scrollbar to the left
-              sx={{
-                width: '100%',
-                height: '100%',
-                boxSizing: 'border-box',
-                background: 'var(--jp-layout-color0)',
-                display: 'flex',
-                flexDirection: 'column'
-              }}
-            >
-              {/* top bar */}
-              <Box sx={{ display: 'flex', justifyContent: 'space-between' }}>
-                {view !== ChatView.Chat ? (
-                  <IconButton onClick={() => setView(ChatView.Chat)}>
-                    <ArrowBackIcon />
-                  </IconButton>
-                ) : (
-                  <Box />
-                )}
-                {view === ChatView.Chat ? (
-                  <Box sx={{ display: 'flex' }}>
-                    {!showWelcomeMessage && (
-                      <TooltippedIconButton
-                        onClick={() =>
-                          props.chatHandler.sendMessage({ type: 'clear' })
-                        }
-                        tooltip="New chat"
-                      >
-                        <AddIcon />
-                      </TooltippedIconButton>
-                    )}
-                    <IconButton onClick={() => openSettingsView()}>
-                      <SettingsIcon />
+            <TelemetryContextProvider telemetryHandler={props.telemetryHandler}>
+              <Box
+                // root box should not include padding as it offsets the vertical
+                // scrollbar to the left
+                sx={{
+                  width: '100%',
+                  height: '100%',
+                  boxSizing: 'border-box',
+                  background: 'var(--jp-layout-color0)',
+                  display: 'flex',
+                  flexDirection: 'column'
+                }}
+              >
+                {/* top bar */}
+                <Box sx={{ display: 'flex', justifyContent: 'space-between' }}>
+                  {view !== ChatView.Chat ? (
+                    <IconButton onClick={() => setView(ChatView.Chat)}>
+                      <ArrowBackIcon />
                     </IconButton>
-                  </Box>
-                ) : (
-                  <Box />
+                  ) : (
+                    <Box />
+                  )}
+                  {view === ChatView.Chat ? (
+                    <Box sx={{ display: 'flex' }}>
+                      {!showWelcomeMessage && (
+                        <TooltippedIconButton
+                          onClick={() =>
+                            props.chatHandler.sendMessage({ type: 'clear' })
+                          }
+                          tooltip="New chat"
+                        >
+                          <AddIcon />
+                        </TooltippedIconButton>
+                      )}
+                      <IconButton onClick={() => openSettingsView()}>
+                        <SettingsIcon />
+                      </IconButton>
+                    </Box>
+                  ) : (
+                    <Box />
+                  )}
+                </Box>
+                {/* body */}
+                {view === ChatView.Chat && (
+                  <ChatBody
+                    chatHandler={props.chatHandler}
+                    openSettingsView={openSettingsView}
+                    showWelcomeMessage={showWelcomeMessage}
+                    setShowWelcomeMessage={setShowWelcomeMessage}
+                    rmRegistry={props.rmRegistry}
+                    focusInputSignal={props.focusInputSignal}
+                    messageFooter={props.messageFooter}
+                  />
+                )}
+                {view === ChatView.Settings && (
+                  <ChatSettings
+                    rmRegistry={props.rmRegistry}
+                    completionProvider={props.completionProvider}
+                    openInlineCompleterSettings={
+                      props.openInlineCompleterSettings
+                    }
+                  />
                 )}
               </Box>
-              {/* body */}
-              {view === ChatView.Chat && (
-                <ChatBody
-                  chatHandler={props.chatHandler}
-                  openSettingsView={openSettingsView}
-                  showWelcomeMessage={showWelcomeMessage}
-                  setShowWelcomeMessage={setShowWelcomeMessage}
-                  rmRegistry={props.rmRegistry}
-                  focusInputSignal={props.focusInputSignal}
-                  messageFooter={props.messageFooter}
-                />
-              )}
-              {view === ChatView.Settings && (
-                <ChatSettings
-                  rmRegistry={props.rmRegistry}
-                  completionProvider={props.completionProvider}
-                  openInlineCompleterSettings={
-                    props.openInlineCompleterSettings
-                  }
-                />
-              )}
-            </Box>
+            </TelemetryContextProvider>
           </ActiveCellContextProvider>
         </CollaboratorsContextProvider>
       </SelectionContextProvider>

packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx

@@ -1,8 +1,10 @@
 import React from 'react';
 import { Box } from '@mui/material';
-import { addAboveIcon, addBelowIcon } from '@jupyterlab/ui-components';
-
-import { CopyButton } from './copy-button';
+import {
+  addAboveIcon,
+  addBelowIcon,
+  copyIcon
+} from '@jupyterlab/ui-components';
 import { replaceCellIcon } from '../../icons';
 
 import {
@@ -11,20 +13,29 @@ import {
 } from '../../contexts/active-cell-context';
 import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button';
 import { useReplace } from '../../hooks/use-replace';
+import { useCopy } from '../../hooks/use-copy';
+import { AiService } from '../../handler';
+import { useTelemetry } from '../../contexts/telemetry-context';
+import { TelemetryEvent } from '../../tokens';
 
 export type CodeToolbarProps = {
   /**
    * The content of the Markdown code block this component is attached to.
    */
-  content: string;
+  code: string;
+  /**
+   * Parent message which contains the code referenced by `content`.
+   */
+  parentMessage?: AiService.ChatMessage;
 };
 
 export function CodeToolbar(props: CodeToolbarProps): JSX.Element {
   const activeCell = useActiveCellContext();
-  const sharedToolbarButtonProps = {
-    content: props.content,
+  const sharedToolbarButtonProps: ToolbarButtonProps = {
+    code: props.code,
     activeCellManager: activeCell.manager,
-    activeCellExists: activeCell.exists
+    activeCellExists: activeCell.exists,
+    parentMessage: props.parentMessage
   };
 
   return (
@@ -41,27 +52,68 @@ export function CodeToolbar(props: CodeToolbarProps): JSX.Element {
     >
       <InsertAboveButton {...sharedToolbarButtonProps} />
       <InsertBelowButton {...sharedToolbarButtonProps} />
-      <ReplaceButton value={props.content} />
-      <CopyButton value={props.content} />
+      <ReplaceButton {...sharedToolbarButtonProps} />
+      <CopyButton {...sharedToolbarButtonProps} />
     </Box>
   );
 }
 
 type ToolbarButtonProps = {
-  content: string;
+  code: string;
   activeCellExists: boolean;
   activeCellManager: ActiveCellManager;
+  parentMessage?: AiService.ChatMessage;
+  // TODO: parentMessage should always be defined, but this can be undefined
+  // when the code toolbar appears in Markdown help messages in the Settings
+  // UI. The Settings UI should use a different component to render Markdown,
+  // and should never render code toolbars within it.
 };
 
+function buildTelemetryEvent(
+  type: string,
+  props: ToolbarButtonProps
+): TelemetryEvent {
+  const charCount = props.code.length;
+  // number of lines = number of newlines + 1
+  const lineCount = (props.code.match(/\n/g) ?? []).length + 1;
+
+  return {
+    type,
+    message: {
+      id: props.parentMessage?.id ?? '',
+      type: props.parentMessage?.type ?? 'human',
+      time: props.parentMessage?.time ?? 0,
+      metadata:
+        props.parentMessage && 'metadata' in props.parentMessage
+          ? props.parentMessage.metadata
+          : {}
+    },
+    code: {
+      charCount,
+      lineCount
+    }
+  };
+}
+
 function InsertAboveButton(props: ToolbarButtonProps) {
+  const telemetryHandler = useTelemetry();
   const tooltip = props.activeCellExists
     ? 'Insert above active cell'
     : 'Insert above active cell (no active cell)';
 
   return (
     <TooltippedIconButton
       tooltip={tooltip}
-      onClick={() => props.activeCellManager.insertAbove(props.content)}
+      onClick={() => {
+        props.activeCellManager.insertAbove(props.code);
+
+        try {
+          telemetryHandler.onEvent(buildTelemetryEvent('insert-above', props));
+        } catch (e) {
+          console.error(e);
+          return;
+        }
+      }}
       disabled={!props.activeCellExists}
     >
       <addAboveIcon.react height="16px" width="16px" />
@@ -70,6 +122,7 @@ function InsertAboveButton(props: ToolbarButtonProps) {
 }
 
 function InsertBelowButton(props: ToolbarButtonProps) {
+  const telemetryHandler = useTelemetry();
   const tooltip = props.activeCellExists
     ? 'Insert below active cell'
     : 'Insert below active cell (no active cell)';
@@ -78,23 +131,67 @@ function InsertBelowButton(props: ToolbarButtonProps) {
     <TooltippedIconButton
       tooltip={tooltip}
       disabled={!props.activeCellExists}
-      onClick={() => props.activeCellManager.insertBelow(props.content)}
+      onClick={() => {
+        props.activeCellManager.insertBelow(props.code);
+
+        try {
+          telemetryHandler.onEvent(buildTelemetryEvent('insert-below', props));
+        } catch (e) {
+          console.error(e);
+          return;
+        }
+      }}
     >
       <addBelowIcon.react height="16px" width="16px" />
     </TooltippedIconButton>
   );
 }
 
-function ReplaceButton(props: { value: string }) {
+function ReplaceButton(props: ToolbarButtonProps) {
+  const telemetryHandler = useTelemetry();
   const { replace, replaceDisabled, replaceLabel } = useReplace();
 
   return (
     <TooltippedIconButton
       tooltip={replaceLabel}
       disabled={replaceDisabled}
-      onClick={() => replace(props.value)}
+      onClick={() => {
+        replace(props.code);
+
+        try {
+          telemetryHandler.onEvent(buildTelemetryEvent('replace', props));
+        } catch (e) {
+          console.error(e);
+          return;
+        }
+      }}
     >
       <replaceCellIcon.react height="16px" width="16px" />
     </TooltippedIconButton>
   );
 }
+
+export function CopyButton(props: ToolbarButtonProps): JSX.Element {
+  const telemetryHandler = useTelemetry();
+  const { copy, copyLabel } = useCopy();
+
+  return (
+    <TooltippedIconButton
+      tooltip={copyLabel}
+      placement="top"
+      onClick={() => {
+        copy(props.code);
+
+        try {
+          telemetryHandler.onEvent(buildTelemetryEvent('copy', props));
+        } catch (e) {
+          console.error(e);
+          return;
+        }
+      }}
+      aria-label="Copy to clipboard"
+    >
+      <copyIcon.react height="16px" width="16px" />
+    </TooltippedIconButton>
+  );
+}