Add support for optional telemetry plugin (#1018)

* remove console log accidentally merged with #1013 * set metadata on stream messages in the chat_history array * implement support for optional telemetry plugin * anonymize message & code details * export telemetry hook from NPM package entry point
jupyterlab · Oct 7, 2024 · 097dbe4 · 097dbe4
1 parent 636d5e9
commit 097dbe4
Show file tree

Hide file tree

Showing 11 changed files with 322 additions and 113 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/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):

diff --git a/packages/jupyter-ai/src/components/chat-messages.tsx b/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
  }

diff --git a/packages/jupyter-ai/src/components/chat.tsx b/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>

diff --git a/packages/jupyter-ai/src/components/code-blocks/code-toolbar.tsx b/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>
+ );
+}