@jupyterlite/ai 0.8.0 → 0.9.0-a0

package/src/components/token-usage-display.tsx

@@ -0,0 +1,130 @@
+import { ReactWidget, UseSignal } from '@jupyterlab/ui-components';
+import React from 'react';
+import { ISignal } from '@lumino/signaling';
+import { AISettingsModel } from '../models/settings-model';
+import { ITokenUsage } from '../tokens';
+
+/**
+ * Props for the TokenUsageDisplay component.
+ */
+export interface ITokenUsageDisplayProps {
+  /**
+   * The token usage changed signal
+   */
+  tokenUsageChanged: ISignal<any, ITokenUsage>;
+
+  /**
+   * The settings model instance for configuration options
+   */
+  settingsModel: AISettingsModel;
+}
+
+/**
+ * React component that displays token usage information.
+ * Shows input/output token counts with up/down arrows.
+ * Only renders when token usage display is enabled in settings.
+ */
+export const TokenUsageDisplay: React.FC<ITokenUsageDisplayProps> = ({
+  tokenUsageChanged,
+  settingsModel
+}) => {
+  return (
+    <UseSignal signal={settingsModel.stateChanged} initialArgs={undefined}>
+      {() => {
+        const config = settingsModel.config;
+        if (!config.showTokenUsage) {
+          return null;
+        }
+
+        return (
+          <UseSignal signal={tokenUsageChanged}>
+            {(_, tokenUsage: ITokenUsage | null | undefined) => {
+              if (!tokenUsage) {
+                return null;
+              }
+
+              const total = tokenUsage.inputTokens + tokenUsage.outputTokens;
+              if (total === 0) {
+                return null;
+              }
+
+              return (
+                <div
+                  style={{
+                    display: 'flex',
+                    alignItems: 'center',
+                    gap: '6px',
+                    fontSize: '12px',
+                    color: 'var(--jp-ui-font-color2)',
+                    padding: '4px 8px',
+                    backgroundColor: 'var(--jp-layout-color1)',
+                    border: '1px solid var(--jp-border-color1)',
+                    borderRadius: '4px',
+                    whiteSpace: 'nowrap'
+                  }}
+                  title={`Token Usage - Sent: ${tokenUsage.inputTokens.toLocaleString()}, Received: ${tokenUsage.outputTokens.toLocaleString()}, Total: ${total.toLocaleString()}`}
+                >
+                  <span
+                    style={{
+                      display: 'flex',
+                      alignItems: 'center',
+                      gap: '2px'
+                    }}
+                  >
+                    <span>↑</span>
+                    <span>{tokenUsage.inputTokens.toLocaleString()}</span>
+                  </span>
+                  <span
+                    style={{
+                      display: 'flex',
+                      alignItems: 'center',
+                      gap: '2px'
+                    }}
+                  >
+                    <span>↓</span>
+                    <span>{tokenUsage.outputTokens.toLocaleString()}</span>
+                  </span>
+                </div>
+              );
+            }}
+          </UseSignal>
+        );
+      }}
+    </UseSignal>
+  );
+};
+
+/**
+ * JupyterLab widget wrapper for the TokenUsageDisplay component.
+ * Extends ReactWidget to integrate with the JupyterLab widget system.
+ */
+export class TokenUsageWidget extends ReactWidget {
+  /**
+   * Creates a new TokenUsageWidget instance.
+   * @param options - Configuration options containing required models
+   */
+  constructor(options: {
+    tokenUsageChanged: ISignal<any, ITokenUsage>;
+    settingsModel: AISettingsModel;
+  }) {
+    super();
+    this._tokenUsageChanged = options.tokenUsageChanged;
+    this._settingsModel = options.settingsModel;
+  }
+
+  /**
+   * Renders the React component within the widget.
+   * @returns The TokenUsageDisplay React element
+   */
+  protected render(): React.ReactElement {
+    return (
+      <TokenUsageDisplay
+        tokenUsageChanged={this._tokenUsageChanged}
+        settingsModel={this._settingsModel}
+      />
+    );
+  }
+
+  private _tokenUsageChanged: ISignal<any, ITokenUsage>;
+  private _settingsModel: AISettingsModel;
+}

package/src/components/tool-select.tsx

@@ -0,0 +1,218 @@
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+
+import BuildIcon from '@mui/icons-material/Build';
+
+import CheckIcon from '@mui/icons-material/Check';
+
+import { Menu, MenuItem, Tooltip, Typography } from '@mui/material';
+
+import React, { useCallback, useEffect, useState } from 'react';
+
+import { INamedTool, IToolRegistry } from '../tokens';
+import { AIChatModel } from '../chat-model';
+
+const SELECT_ITEM_CLASS = 'jp-AIToolSelect-item';
+
+/**
+ * Properties for the tool select component.
+ */
+export interface IToolSelectProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The tool registry to get available tools from.
+   */
+  toolRegistry: IToolRegistry;
+
+  /**
+   * Whether tools are enabled.
+   */
+  toolsEnabled: boolean;
+
+  /**
+   * Function to handle tool selection changes.
+   */
+  onToolSelectionChange: (selectedToolNames: string[]) => void;
+}
+
+/**
+ * The tool select component for choosing AI tools.
+ */
+export function ToolSelect(props: IToolSelectProps): JSX.Element {
+  const { toolRegistry, onToolSelectionChange, toolsEnabled } = props;
+
+  const [selectedToolNames, setSelectedToolNames] = useState<string[]>([]);
+  const [tools, setTools] = useState<INamedTool[]>(
+    toolRegistry?.namedTools || []
+  );
+  const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLElement | null>(null);
+  const [menuOpen, setMenuOpen] = useState(false);
+
+  const openMenu = useCallback((el: HTMLElement | null) => {
+    setMenuAnchorEl(el);
+    setMenuOpen(true);
+  }, []);
+
+  const closeMenu = useCallback(() => {
+    setMenuOpen(false);
+  }, []);
+
+  const toggleTool = useCallback(
+    (toolName: string) => {
+      const currentToolNames = [...selectedToolNames];
+      const index = currentToolNames.indexOf(toolName);
+
+      if (index !== -1) {
+        // Remove tool
+        currentToolNames.splice(index, 1);
+      } else {
+        // Add tool
+        currentToolNames.push(toolName);
+      }
+
+      setSelectedToolNames(currentToolNames);
+      onToolSelectionChange(currentToolNames);
+    },
+    [selectedToolNames, onToolSelectionChange]
+  );
+
+  // Update tools when registry changes
+  useEffect(() => {
+    const updateTools = () => {
+      const newTools = toolRegistry?.namedTools || [];
+      setTools(newTools);
+    };
+
+    if (toolRegistry) {
+      updateTools();
+      toolRegistry.toolsChanged.connect(updateTools);
+      return () => {
+        toolRegistry.toolsChanged.disconnect(updateTools);
+      };
+    }
+  }, [toolRegistry]);
+
+  // Initialize selected tools to all tools by default
+  useEffect(() => {
+    if (tools.length > 0 && selectedToolNames.length === 0) {
+      const defaultToolNames = tools.map(tool => tool.name);
+      setSelectedToolNames(defaultToolNames);
+      onToolSelectionChange(defaultToolNames);
+    }
+  }, [tools, selectedToolNames.length, onToolSelectionChange]);
+
+  // Don't render if tools are disabled or no tools available
+  if (!toolsEnabled || tools.length === 0) {
+    return <></>;
+  }
+
+  return (
+    <>
+      <TooltippedButton
+        onClick={e => {
+          openMenu(e.currentTarget);
+        }}
+        tooltip={`Tools (${selectedToolNames.length}/${tools.length} selected)`}
+        buttonProps={{
+          size: 'small',
+          variant: selectedToolNames.length > 0 ? 'contained' : 'outlined',
+          color: 'primary',
+          title: 'Select AI Tools',
+          onKeyDown: e => {
+            if (e.key !== 'Enter' && e.key !== ' ') {
+              return;
+            }
+            openMenu(e.currentTarget);
+            // Stop propagation to prevent sending message
+            e.stopPropagation();
+          }
+        }}
+        sx={
+          selectedToolNames.length === 0
+            ? { backgroundColor: 'var(--jp-layout-color3)' }
+            : {}
+        }
+      >
+        <BuildIcon />
+      </TooltippedButton>
+
+      <Menu
+        open={menuOpen}
+        onClose={closeMenu}
+        anchorEl={menuAnchorEl}
+        anchorOrigin={{
+          vertical: 'top',
+          horizontal: 'right'
+        }}
+        transformOrigin={{
+          vertical: 'bottom',
+          horizontal: 'right'
+        }}
+        sx={{
+          '& .MuiMenuItem-root': {
+            padding: '0.5em',
+            paddingRight: '2em'
+          }
+        }}
+      >
+        {tools.map(namedTool => (
+          <Tooltip
+            key={namedTool.name}
+            title={namedTool.tool.description || namedTool.name}
+            placement="left"
+          >
+            <MenuItem
+              className={SELECT_ITEM_CLASS}
+              onClick={e => {
+                toggleTool(namedTool.name);
+                // Prevent sending message on tool selection
+                e.stopPropagation();
+              }}
+            >
+              {selectedToolNames.includes(namedTool.name) ? (
+                <CheckIcon
+                  sx={{
+                    marginRight: '8px',
+                    color: 'var(--jp-brand-color1, #2196F3)'
+                  }}
+                />
+              ) : (
+                <div style={{ width: '24px', marginRight: '8px' }} />
+              )}
+              <Typography variant="body2">{namedTool.name}</Typography>
+            </MenuItem>
+          </Tooltip>
+        ))}
+      </Menu>
+    </>
+  );
+}
+
+/**
+ * Factory function returning the toolbar item for tool selection.
+ */
+export function createToolSelectItem(
+  toolRegistry: IToolRegistry,
+  toolsEnabled: boolean = true
+): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const onToolSelectionChange = (tools: string[]) => {
+        const chatContext = props.model
+          .chatContext as AIChatModel.IAIChatContext;
+        if (!chatContext.agentManager) {
+          return;
+        }
+        chatContext.agentManager.setSelectedTools(tools);
+      };
+
+      const toolSelectProps: IToolSelectProps = {
+        ...props,
+        toolRegistry,
+        onToolSelectionChange,
+        toolsEnabled
+      };
+      return <ToolSelect {...toolSelectProps} />;
+    },
+    position: 1
+  };
+}

package/src/icons.ts

@@ -1,18 +1,16 @@
-/*
- * Copyright (c) Jupyter Development Team.
- * Distributed under the terms of the Modified BSD License.
- */
-
 import { LabIcon } from '@jupyterlab/ui-components';
 
-/**
- * This icon is based on the jupyternaut icon from Jupyter AI:
- * https://github.com/jupyterlab/jupyter-ai/blob/main/packages/jupyter-ai/style/icons/jupyternaut.svg
- * With a small tweak for the colors to match the JupyterLite icon.
- */
-import jupyternautLiteSvg from '../style/icons/jupyternaut-lite.svg';
+import labaiIconSvg from '../style/icons/jupyternaut-lite.svg';
+
+export const labaiIcon = new LabIcon({
+  name: '@jupyterlite/ai:icon',
+  svgstr: labaiIconSvg
+});
 
-export const jupyternautLiteIcon = new LabIcon({
-  name: '@jupyterlite/ai:jupyternaut-lite',
-  svgstr: jupyternautLiteSvg
+export const jupyternautIcon = new LabIcon({
+  name: '@jupyterlite/ai:jupyternaut',
+  svgstr: labaiIconSvg
 });
+
+const AI_AVATAR_BASE64 = btoa(jupyternautIcon.svgstr);
+export const AI_AVATAR = `data:image/svg+xml;base64,${AI_AVATAR_BASE64}`;