@jupyterlite/ai 0.8.1 → 0.9.0-a1

package/src/completion/completion-provider.ts

@@ -0,0 +1,346 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext,
+  IInlineCompletionList,
+  IInlineCompletionProvider
+} from '@jupyterlab/completer';
+import { NotebookPanel } from '@jupyterlab/notebook';
+import { generateText, LanguageModel } from 'ai';
+import { ISecretsManager } from 'jupyter-secrets-manager';
+
+import { AISettingsModel } from '../models/settings-model';
+import { createCompletionModel } from '../providers/models';
+import { SECRETS_NAMESPACE, type ICompletionProviderRegistry } from '../tokens';
+
+/**
+ * Configuration interface for provider-specific completion behavior
+ */
+export interface IProviderCompletionConfig {
+  /**
+   * Temperature setting for the provider
+   */
+  temperature?: number;
+
+  /**
+   * Whether the provider supports fill-in-the-middle completion
+   */
+  supportsFillInMiddle?: boolean;
+
+  /**
+   * Whether to set filterText for this provider
+   */
+  useFilterText?: boolean;
+
+  /**
+   * Custom prompt formatter for provider-specific requirements
+   */
+  customPromptFormat?: (prompt: string, suffix: string) => string;
+
+  /**
+   * Function to clean up provider-specific artifacts from completion text
+   */
+  cleanupCompletion?: (completion: string) => string;
+}
+
+/**
+ * Default system prompt for code completion
+ */
+const DEFAULT_COMPLETION_SYSTEM_PROMPT = `You are an AI code completion assistant. Complete the given code fragment with appropriate code.
+Rules:
+- Return only the completion text, no explanations or comments
+- Do not include code block markers (\`\`\` or similar)
+- Make completions contextually relevant to the surrounding code and notebook context
+- Follow the language-specific conventions and style guidelines for the detected programming language
+- Keep completions concise but functional
+- Do not repeat the existing code that comes before the cursor
+- Use variables, imports, functions, and other definitions from previous notebook cells when relevant`;
+
+/**
+ * The generic completion provider to register to the completion provider manager.
+ */
+export class AICompletionProvider implements IInlineCompletionProvider {
+  /**
+   * Construct a new completion provider.
+   */
+  constructor(options: AICompletionProvider.IOptions) {
+    Private.setToken(options.token);
+    this._settingsModel = options.settingsModel;
+    this._completionProviderRegistry = options.completionProviderRegistry;
+    this._secretsManager = options.secretsManager;
+    this._settingsModel.stateChanged.connect(() => {
+      this._updateModel();
+    });
+    this._updateModel();
+  }
+
+  /**
+   * The unique identifier of the provider.
+   */
+  readonly identifier = '@jupyterlite/ai:completer';
+
+  /**
+   * Get the current completer name based on settings.
+   */
+  get name(): string {
+    const activeProvider = this._settingsModel.getCompleterProvider();
+    return activeProvider ? `${activeProvider.provider}-completer` : 'none';
+  }
+
+  /**
+   * Get the system prompt for the completion.
+   */
+  get systemPrompt(): string {
+    return DEFAULT_COMPLETION_SYSTEM_PROMPT;
+  }
+
+  /**
+   * Fetch completion items based on the request and context.
+   */
+  async fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ): Promise<IInlineCompletionList> {
+    if (!this._model) {
+      return { items: [] };
+    }
+
+    const { text, offset: cursorOffset } = request;
+    const prompt = text.slice(0, cursorOffset);
+    const suffix = text.slice(cursorOffset);
+
+    // Get current provider settings
+    const activeProvider = this._settingsModel.getCompleterProvider();
+    if (!activeProvider) {
+      return { items: [] };
+    }
+
+    const provider = activeProvider.provider;
+    const providerConfig = this._getProviderCompletionConfig(provider);
+
+    try {
+      let completionPrompt: string;
+
+      // Check if we're in a notebook or file and handle context accordingly
+      if (context.widget instanceof NotebookPanel) {
+        // Extract notebook context with surrounding cells
+        const contextString = this._extractNotebookContext(context, request);
+        completionPrompt = contextString;
+      } else {
+        // For files, use simpler approach
+        completionPrompt = prompt.trim();
+        if (providerConfig.customPromptFormat && suffix.trim()) {
+          completionPrompt = providerConfig.customPromptFormat(prompt, suffix);
+        }
+      }
+
+      const { text: completion } = await generateText({
+        model: this._model,
+        prompt: completionPrompt,
+        system: this.systemPrompt,
+        temperature: providerConfig.temperature || 0.3
+      });
+
+      // Clean up provider-specific artifacts if cleanup function is provided
+      let cleanCompletion = completion;
+      if (providerConfig.cleanupCompletion) {
+        cleanCompletion = providerConfig.cleanupCompletion(completion);
+      }
+
+      const items = [
+        {
+          insertText: cleanCompletion,
+          filterText: providerConfig.useFilterText
+            ? prompt.substring(completionPrompt.length)
+            : undefined
+        }
+      ];
+
+      return { items };
+    } catch (error) {
+      console.error(`Error fetching completions from ${provider}:`, error);
+      return { items: [] };
+    }
+  }
+
+  /**
+   * Update the language model based on current settings.
+   */
+  private async _updateModel(): Promise<void> {
+    const activeProvider = this._settingsModel.getCompleterProvider();
+    if (!activeProvider) {
+      this._model = null;
+      return;
+    }
+
+    const provider = activeProvider.provider;
+    const model = activeProvider.model;
+
+    let apiKey: string;
+    if (this._secretsManager && this._settingsModel.config.useSecretsManager) {
+      apiKey =
+        (
+          await this._secretsManager.get(
+            Private.getToken(),
+            SECRETS_NAMESPACE,
+            `${provider}:apiKey`
+          )
+        )?.value ?? '';
+    } else {
+      apiKey = this._settingsModel.getApiKey(activeProvider.id);
+    }
+
+    try {
+      this._model = createCompletionModel(
+        {
+          provider,
+          model,
+          apiKey
+        },
+        this._completionProviderRegistry
+      );
+    } catch (error) {
+      console.error(`Error creating model for ${provider}:`, error);
+      this._model = null;
+    }
+  }
+
+  /**
+   * Extract context from notebook cells
+   */
+  private _extractNotebookContext(
+    context: IInlineCompletionContext,
+    request: CompletionHandler.IRequest
+  ): string {
+    const { text, offset: cursorOffset } = request;
+    let codeBeforeCursor = text.slice(0, cursorOffset);
+    let codeAfterCursor = text.slice(cursorOffset);
+
+    const notebookPanel = context.widget as NotebookPanel;
+    const notebook = notebookPanel.content;
+    const currentCellIndex = notebook.activeCellIndex;
+    const cells = notebook.widgets;
+
+    // For notebooks, include context from surrounding cells
+    const cellsAbove: string[] = [];
+    const cellsBelow: string[] = [];
+
+    // Get content from cells above current cell
+    for (let i = 0; i < currentCellIndex; i++) {
+      const cell = cells[i];
+      if (cell.model.type === 'code') {
+        const source = cell.model.sharedModel.source;
+        if (source.trim()) {
+          cellsAbove.push(source.trim());
+        }
+      }
+    }
+
+    // Get content from cells below current cell
+    for (let i = currentCellIndex + 1; i < cells.length; i++) {
+      const cell = cells[i];
+      if (cell.model.type === 'code') {
+        const source = cell.model.sharedModel.source;
+        if (source.trim()) {
+          cellsBelow.push(source.trim());
+        }
+      }
+    }
+
+    // Include cells above in the code before cursor
+    if (cellsAbove.length > 0) {
+      const cellsAboveText = cellsAbove
+        .map((cell, index) => `# Cell ${index + 1}:\n${cell}`)
+        .join('\n\n');
+      codeBeforeCursor = `${cellsAboveText}\n\n# Current cell:\n${codeBeforeCursor}`;
+    }
+
+    // Include cells below in the code after cursor
+    if (cellsBelow.length > 0) {
+      const cellsBelowText = cellsBelow
+        .map((cell, index) => `# Cell ${index + 1}:\n${cell}`)
+        .join('\n\n');
+      codeAfterCursor = `${codeAfterCursor}\n\n# Cells below:\n${cellsBelowText}`;
+    }
+
+    const parts: string[] = [];
+
+    // Add code before cursor
+    if (codeBeforeCursor) {
+      parts.push('# Code before cursor:');
+      parts.push(codeBeforeCursor);
+    }
+
+    // Add completion instruction
+    parts.push('# Complete the code at cursor position');
+
+    // Add code after cursor
+    if (codeAfterCursor) {
+      parts.push('# Code after cursor:');
+      parts.push(codeAfterCursor);
+    }
+
+    return parts.length > 1 ? parts.join('\n\n') + '\n\n' : '';
+  }
+
+  /**
+   * Get provider-specific completion configuration from registry
+   */
+  private _getProviderCompletionConfig(
+    provider: string
+  ): IProviderCompletionConfig {
+    const providerInfo =
+      this._completionProviderRegistry?.getProviderInfo(provider);
+    const completionConfig = providerInfo?.customSettings?.completionConfig;
+
+    // Return provider config or default config
+    return (
+      completionConfig || {
+        temperature: 0.3,
+        supportsFillInMiddle: false,
+        useFilterText: false
+      }
+    );
+  }
+
+  private _settingsModel: AISettingsModel;
+  private _completionProviderRegistry?: ICompletionProviderRegistry;
+  private _model: LanguageModel | null = null;
+  private _secretsManager?: ISecretsManager;
+}
+
+export namespace AICompletionProvider {
+  /**
+   * The options for the constructor of the completion provider.
+   */
+  export interface IOptions {
+    /**
+     * The AI settings model.
+     */
+    settingsModel: AISettingsModel;
+    /**
+     * The completion provider registry.
+     */
+    completionProviderRegistry?: ICompletionProviderRegistry;
+    /**
+     * The secrets manager.
+     */
+    secretsManager?: ISecretsManager;
+    /**
+     * The token used to request the secrets manager.
+     */
+    token: symbol;
+  }
+}
+
+namespace Private {
+  /**
+   * The token to use with the secrets manager, setter and getter.
+   */
+  let secretsToken: symbol;
+  export function setToken(value: symbol): void {
+    secretsToken = value;
+  }
+  export function getToken(): symbol {
+    return secretsToken;
+  }
+}

package/src/completion/index.ts

@@ -0,0 +1 @@
+export * from './completion-provider';

package/src/components/clear-button.tsx

@@ -0,0 +1,56 @@
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+
+import ClearIcon from '@mui/icons-material/Clear';
+
+import React from 'react';
+
+import { AIChatModel } from '../chat-model';
+
+/**
+ * Properties of the clear button.
+ */
+export interface IClearButtonProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The function to clear messages.
+   */
+  clearMessages: () => void;
+}
+
+/**
+ * The clear button component.
+ */
+export function ClearButton(props: IClearButtonProps): JSX.Element {
+  const tooltip = 'Clear chat';
+  return (
+    <TooltippedButton
+      onClick={props.clearMessages}
+      tooltip={tooltip}
+      buttonProps={{
+        size: 'small',
+        variant: 'outlined',
+        color: 'secondary',
+        title: tooltip
+      }}
+    >
+      <ClearIcon />
+    </TooltippedButton>
+  );
+}
+
+/**
+ * Factory returning the clear button toolbar item.
+ */
+export function clearItem(): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const { model } = props;
+      const clearMessages = () =>
+        (model.chatContext as AIChatModel.IAIChatContext).clearMessages();
+      const clearProps: IClearButtonProps = { ...props, clearMessages };
+      return ClearButton(clearProps);
+    },
+    position: 0,
+    hidden: false
+  };
+}

package/src/components/index.ts

@@ -0,0 +1,3 @@
+export * from './clear-button';
+export * from './stop-button';
+export * from './tool-select';

package/src/components/model-select.tsx

@@ -0,0 +1,245 @@
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+import CheckIcon from '@mui/icons-material/Check';
+import { Menu, MenuItem, Typography } from '@mui/material';
+import React, { useCallback, useEffect, useState } from 'react';
+import { AISettingsModel } from '../models/settings-model';
+
+const SELECT_ITEM_CLASS = 'labai-model-select-item';
+
+/**
+ * Properties for the model select component.
+ */
+export interface IModelSelectProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The settings model to get available models and current selection from.
+   */
+  settingsModel: AISettingsModel;
+}
+
+/**
+ * The model select component for choosing AI models.
+ */
+export function ModelSelect(props: IModelSelectProps): JSX.Element {
+  const { settingsModel } = props;
+
+  const [currentProvider, setCurrentProvider] = useState<string>('');
+  const [currentModel, setCurrentModel] = useState<string>('');
+  const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLElement | null>(null);
+  const [menuOpen, setMenuOpen] = useState(false);
+
+  // Get configured providers from settings model
+  const configuredProviders = settingsModel.providers;
+
+  const openMenu = useCallback((el: HTMLElement | null) => {
+    setMenuAnchorEl(el);
+    setMenuOpen(true);
+  }, []);
+
+  const closeMenu = useCallback(() => {
+    setMenuOpen(false);
+  }, []);
+
+  const selectModel = useCallback(
+    async (providerId: string) => {
+      // Set the active provider using the provider ID
+      await settingsModel.setActiveProvider(providerId);
+      closeMenu();
+
+      // Provider selected successfully
+    },
+    [settingsModel, closeMenu]
+  );
+
+  // Update current selection when settings model changes
+  useEffect(() => {
+    const updateCurrentSelection = () => {
+      const activeProvider = settingsModel.getActiveProvider();
+      if (activeProvider) {
+        setCurrentProvider(activeProvider.id);
+        setCurrentModel(activeProvider.model);
+      }
+    };
+
+    updateCurrentSelection();
+    settingsModel.stateChanged.connect(updateCurrentSelection);
+    return () => {
+      settingsModel.stateChanged.disconnect(updateCurrentSelection);
+    };
+  }, [settingsModel]);
+
+  // Get current provider label for display
+  const activeProvider = settingsModel.getActiveProvider();
+  const currentProviderLabel = activeProvider?.name || currentProvider;
+
+  // Use all configured providers (they're already validated when added)
+  const availableProviders = configuredProviders;
+
+  // Get available model combinations from configured providers
+  const availableModels = availableProviders.map(provider => ({
+    provider: provider.id,
+    providerLabel: provider.name,
+    model: provider.model,
+    isSelected:
+      provider.id === currentProvider && provider.model === currentModel
+  }));
+
+  // Show a message if no providers are configured
+  if (availableModels.length === 0) {
+    return (
+      <TooltippedButton
+        onClick={() => {}}
+        tooltip="No providers configured. Please go to AI Settings to add a provider."
+        buttonProps={{
+          size: 'small',
+          variant: 'outlined',
+          color: 'warning',
+          disabled: true,
+          title: 'No Providers Available'
+        }}
+        sx={{
+          minWidth: 'auto',
+          display: 'flex',
+          alignItems: 'center',
+          height: '29px'
+        }}
+      >
+        <Typography
+          variant="caption"
+          sx={{ fontSize: '0.7rem', fontWeight: 500 }}
+        >
+          No Providers
+        </Typography>
+      </TooltippedButton>
+    );
+  }
+
+  return (
+    <>
+      <TooltippedButton
+        onClick={e => {
+          openMenu(e.currentTarget);
+        }}
+        tooltip={`Current Model: ${currentProviderLabel} - ${currentModel}`}
+        buttonProps={{
+          size: 'small',
+          variant: 'contained',
+          color: 'primary',
+          title: 'Select AI Model',
+          onKeyDown: e => {
+            if (e.key !== 'Enter' && e.key !== ' ') {
+              return;
+            }
+            openMenu(e.currentTarget);
+            // Stop propagation to prevent sending message
+            e.stopPropagation();
+          }
+        }}
+        sx={{
+          minWidth: 'auto',
+          display: 'flex',
+          alignItems: 'center',
+          height: '29px'
+        }}
+      >
+        <Typography
+          variant="caption"
+          sx={{ fontSize: '0.7rem', fontWeight: 500, textTransform: 'none' }}
+        >
+          {currentProviderLabel}
+        </Typography>
+      </TooltippedButton>
+
+      <Menu
+        open={menuOpen}
+        onClose={closeMenu}
+        anchorEl={menuAnchorEl}
+        anchorOrigin={{
+          vertical: 'top',
+          horizontal: 'right'
+        }}
+        transformOrigin={{
+          vertical: 'bottom',
+          horizontal: 'right'
+        }}
+        sx={{
+          '& .MuiPaper-root': {
+            maxHeight: '300px',
+            overflowY: 'auto'
+          },
+          '& .MuiMenuItem-root': {
+            padding: '0.5em',
+            paddingRight: '2em',
+            minWidth: '200px'
+          }
+        }}
+      >
+        {availableModels.map(({ provider, providerLabel, isSelected }) => (
+          <MenuItem
+            key={provider}
+            className={SELECT_ITEM_CLASS}
+            onClick={async e => {
+              await selectModel(provider);
+              // Prevent sending message on model selection
+              e.stopPropagation();
+            }}
+            sx={{
+              backgroundColor: isSelected
+                ? 'var(--jp-brand-color3, rgba(33, 150, 243, 0.1))'
+                : 'transparent',
+              '&:hover': {
+                backgroundColor: isSelected
+                  ? 'var(--jp-brand-color3, rgba(33, 150, 243, 0.15))'
+                  : 'var(--jp-layout-color1)'
+              },
+              display: 'flex',
+              alignItems: 'center',
+              gap: '8px'
+            }}
+          >
+            {isSelected ? (
+              <CheckIcon
+                sx={{
+                  color: 'var(--jp-brand-color1, #2196F3)',
+                  fontSize: 16
+                }}
+              />
+            ) : (
+              <div style={{ width: '16px' }} />
+            )}
+            <Typography
+              variant="body2"
+              component="div"
+              sx={{
+                fontWeight: isSelected ? 600 : 400,
+                color: isSelected
+                  ? 'var(--jp-brand-color1, #2196F3)'
+                  : 'inherit'
+              }}
+            >
+              {providerLabel}
+            </Typography>
+          </MenuItem>
+        ))}
+      </Menu>
+    </>
+  );
+}
+
+/**
+ * Factory function returning the toolbar item for model selection.
+ */
+export function createModelSelectItem(
+  settingsModel: AISettingsModel
+): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const modelSelectProps: IModelSelectProps = {
+        ...props,
+        settingsModel
+      };
+      return <ModelSelect {...modelSelectProps} />;
+    },
+    position: 0.5
+  };
+}

package/src/components/stop-button.tsx

@@ -1,12 +1,10 @@
-/*
- * Copyright (c) Jupyter Development Team.
- * Distributed under the terms of the Modified BSD License.
- */
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
 
 import StopIcon from '@mui/icons-material/Stop';
+
 import React from 'react';
 
-import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+import { AIChatModel } from '../chat-model';
 
 /**
  * Properties of the stop button.
@@ -20,7 +18,7 @@ export interface IStopButtonProps
 }
 
 /**
- * The stop button.
+ * The stop button component.
  */
 export function StopButton(props: IStopButtonProps): JSX.Element {
   const tooltip = 'Stop streaming';
@@ -31,6 +29,7 @@ export function StopButton(props: IStopButtonProps): JSX.Element {
       buttonProps={{
         size: 'small',
         variant: 'contained',
+        color: 'error',
         title: tooltip
       }}
     >
@@ -40,17 +39,18 @@ export function StopButton(props: IStopButtonProps): JSX.Element {
 }
 
 /**
- * factory returning the toolbar item.
+ * Factory returning the stop button toolbar item.
  */
-export function stopItem(
-  stopStreaming: () => void
-): InputToolbarRegistry.IToolbarItem {
+export function stopItem(): InputToolbarRegistry.IToolbarItem {
   return {
     element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const { model } = props;
+      const stopStreaming = () =>
+        (model.chatContext as AIChatModel.IAIChatContext).stopStreaming();
       const stopProps: IStopButtonProps = { ...props, stopStreaming };
       return StopButton(stopProps);
     },
     position: 50,
-    hidden: true /* hidden by default */
+    hidden: true // Hidden by default, shown when streaming
   };
 }