Add the chat model to the IChatProvider method's arguments

Author: brichet

packages/jupyter-chat/src/chat-commands/registry.ts

@@ -6,6 +6,7 @@
 import { Token } from '@lumino/coreutils';
 import { ChatCommand, IChatCommandProvider } from './types';
 import { IInputModel } from '../input-model';
+import { IChatModel } from '../model';
 
 /**
  * Interface of a chat command registry, which tracks a list of chat command
@@ -20,7 +21,11 @@ export interface IChatCommandRegistry {
    * Handles a chat command by calling `handleChatCommand()` on the provider
    * corresponding to this chat command.
    */
-  handleChatCommand(command: ChatCommand, inputModel: IInputModel): void;
+  handleChatCommand(
+    command: ChatCommand,
+    inputModel: IInputModel,
+    chatModel: IChatModel
+  ): void;
 }
 
 /**
@@ -39,7 +44,11 @@ export class ChatCommandRegistry implements IChatCommandRegistry {
     return Array.from(this._providers.values());
   }
 
-  handleChatCommand(command: ChatCommand, inputModel: IInputModel) {
+  handleChatCommand(
+    command: ChatCommand,
+    inputModel: IInputModel,
+    chatModel: IChatModel
+  ) {
     const provider = this._providers.get(command.providerId);
     if (!provider) {
       console.error(
@@ -49,7 +58,7 @@ export class ChatCommandRegistry implements IChatCommandRegistry {
       return;
     }
 
-    provider.handleChatCommand(command, inputModel);
+    provider.handleChatCommand(command, inputModel, chatModel);
   }
 
   private _providers: Map<string, IChatCommandProvider>;

packages/jupyter-chat/src/chat-commands/types.ts

@@ -5,6 +5,7 @@
 
 import { LabIcon } from '@jupyterlab/ui-components';
 import { IInputModel } from '../input-model';
+import { IChatModel } from '../model';
 
 export type ChatCommand = {
   /**
@@ -54,14 +55,18 @@ export interface IChatCommandProvider {
    * valid chat commands that match the current word. The current word is
    * space-separated word at the user's cursor.
    */
-  getChatCommands(inputModel: IInputModel): Promise<ChatCommand[]>;
+  getChatCommands(
+    inputModel: IInputModel,
+    chatModel: IChatModel
+  ): Promise<ChatCommand[]>;
 
   /**
    * Function called when a chat command is run by the user through the chat
    * commands menu.
    */
   handleChatCommand(
     command: ChatCommand,
-    inputModel: IInputModel
+    inputModel: IInputModel,
+    chatModel: IChatModel
   ): Promise<void>;
 }

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

@@ -18,19 +18,24 @@ import React, { useEffect, useRef, useState } from 'react';
 
 import { AttachmentPreviewList } from './attachments';
 import { AttachButton, CancelButton, SendButton } from './input';
-import { IInputModel, InputModel } from '../input-model';
-import { IAttachment, Selection } from '../types';
 import { useChatCommands } from './input/use-chat-commands';
+import { IInputModel, InputModel } from '../input-model';
 import { IChatCommandRegistry } from '../chat-commands';
+import { IChatModel } from '../model';
+import { IAttachment, Selection } from '../types';
 
 const INPUT_BOX_CLASS = 'jp-chat-input-container';
 
 export function ChatInput(props: ChatInput.IProps): JSX.Element {
-  const { documentManager, model } = props;
+  const { chatModel, documentManager, model } = props;
   const [input, setInput] = useState<string>(model.value);
   const inputRef = useRef<HTMLInputElement>();
 
-  const chatCommands = useChatCommands(model, props.chatCommandRegistry);
+  const chatCommands = useChatCommands(
+    model,
+    chatModel,
+    props.chatCommandRegistry
+  );
 
   const [sendWithShiftEnter, setSendWithShiftEnter] = useState<boolean>(
     model.config.sendWithShiftEnter ?? false
@@ -273,9 +278,13 @@ export namespace ChatInput {
    */
   export interface IProps {
     /**
-     * The chat model.
+     * The input model.
      */
     model: IInputModel;
+    /**
+     * The chat model.
+     */
+    chatModel: IChatModel;
     /**
      * The function to be called to send the message.
      */

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

@@ -417,6 +417,7 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
             hideIncludeSelection={true}
             chatCommandRegistry={props.chatCommandRegistry}
             documentManager={props.documentManager}
+            chatModel={model}
           />
         ) : (
           <MarkdownRenderer

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

@@ -47,6 +47,7 @@ export function ChatBody(props: Chat.IChatBodyProps): JSX.Element {
         model={model.input}
         documentManager={props.documentManager}
         chatCommandRegistry={props.chatCommandRegistry}
+        chatModel={model}
       />
     </AttachmentOpenerContext.Provider>
   );

packages/jupyter-chat/src/components/input/use-chat-commands.tsx

@@ -13,6 +13,7 @@ import { Box } from '@mui/material';
 
 import { ChatCommand, IChatCommandRegistry } from '../../chat-commands';
 import { IInputModel } from '../../input-model';
+import { IChatModel } from '../../model';
 
 type AutocompleteProps = GenericAutocompleteProps<any, any, any, any>;
 
@@ -32,6 +33,7 @@ type UseChatCommandsReturn = {
  */
 export function useChatCommands(
   inputModel: IInputModel,
+  chatModel: IChatModel,
   chatCommandRegistry?: IChatCommandRegistry
 ): UseChatCommandsReturn {
   // whether an option is highlighted in the chat commands menu
@@ -63,7 +65,7 @@ export function useChatCommands(
         // TODO: optimize performance when this method is truly async
         try {
           newCommands = newCommands.concat(
-            await provider.getChatCommands(inputModel)
+            await provider.getChatCommands(inputModel, chatModel)
           );
         } catch (e) {
           console.error(
@@ -116,7 +118,7 @@ export function useChatCommands(
     }
 
     // otherwise, defer handling to the command provider
-    chatCommandRegistry.handleChatCommand(command, inputModel);
+    chatCommandRegistry.handleChatCommand(command, inputModel, chatModel);
   };
 
   return {

python/jupyterlab-chat/src/chat-commands/providers/emoji.ts

@@ -8,7 +8,8 @@ import {
   IChatCommandProvider,
   IChatCommandRegistry,
   ChatCommand,
-  IInputModel
+  IInputModel,
+  IChatModel
 } from '@jupyter/chat';
 
 export class EmojiCommandProvider implements IChatCommandProvider {
@@ -47,7 +48,7 @@ export class EmojiCommandProvider implements IChatCommandProvider {
   // regex used to test the current word
   private _regex: RegExp = /^:\w*:?/;
 
-  async getChatCommands(inputModel: IInputModel) {
+  async getChatCommands(inputModel: IInputModel, chatModel: IChatModel) {
     const match = inputModel.currentWord?.match(this._regex)?.[0];
     if (!match) {
       return [];
@@ -61,7 +62,8 @@ export class EmojiCommandProvider implements IChatCommandProvider {
 
   async handleChatCommand(
     command: ChatCommand,
-    inputModel: IInputModel
+    inputModel: IInputModel,
+    chatModel: IChatModel
   ): Promise<void> {
     // no handling needed because `replaceWith` is set in each command.
     return;