|
| 1 | +/* |
| 2 | + * Copyright (c) Jupyter Development Team. |
| 3 | + * Distributed under the terms of the Modified BSD License. |
| 4 | + */ |
| 5 | + |
| 6 | +import { JupyterFrontEndPlugin } from '@jupyterlab/application'; |
| 7 | +import type { Contents } from '@jupyterlab/services'; |
| 8 | +import type { DocumentRegistry } from '@jupyterlab/docregistry'; |
| 9 | +import { |
| 10 | + IChatCommandProvider, |
| 11 | + IChatCommandRegistry, |
| 12 | + IInputModel, |
| 13 | + ChatCommand |
| 14 | +} from '@jupyter/chat'; |
| 15 | + |
| 16 | +const CONTEXT_COMMANDS_PROVIDER_ID = |
| 17 | + '@jupyter-ai/core:context-commands-provider'; |
| 18 | + |
| 19 | +/** |
| 20 | + * A command provider that provides completions for context commands like `@file`. |
| 21 | + */ |
| 22 | +export class ContextCommandsProvider implements IChatCommandProvider { |
| 23 | + public id: string = CONTEXT_COMMANDS_PROVIDER_ID; |
| 24 | + private _context_commands: ChatCommand[] = [ |
| 25 | + // TODO: add an icon! |
| 26 | + // import FindInPage from '@mui/icons-material/FindInPage'; |
| 27 | + // may need to change the API to allow JSX els as icons |
| 28 | + { |
| 29 | + name: '@file', |
| 30 | + providerId: this.id, |
| 31 | + replaceWith: '@file:', |
| 32 | + description: 'Include a file with your prompt' |
| 33 | + } |
| 34 | + ]; |
| 35 | + |
| 36 | + constructor( |
| 37 | + contentsManager: Contents.IManager, |
| 38 | + docRegistry: DocumentRegistry |
| 39 | + ) { |
| 40 | + this._contentsManager = contentsManager; |
| 41 | + this._docRegistry = docRegistry; |
| 42 | + } |
| 43 | + |
| 44 | + async getChatCommands(inputModel: IInputModel) { |
| 45 | + // do nothing if the current word does not start with '@'. |
| 46 | + const currentWord = inputModel.currentWord; |
| 47 | + if (!currentWord || !currentWord.startsWith('@')) { |
| 48 | + return []; |
| 49 | + } |
| 50 | + |
| 51 | + // if the current word starts with `@file:`, return a list of valid file |
| 52 | + // paths. |
| 53 | + if (currentWord.startsWith('@file:')) { |
| 54 | + const searchPath = currentWord.split('@file:')[1]; |
| 55 | + const commands = await getPathCompletions( |
| 56 | + this._contentsManager, |
| 57 | + this._docRegistry, |
| 58 | + searchPath |
| 59 | + ); |
| 60 | + return commands; |
| 61 | + } |
| 62 | + |
| 63 | + // otherwise, a context command has not yet been specified. return a list of |
| 64 | + // valid context commands. |
| 65 | + const commands = this._context_commands.filter(cmd => |
| 66 | + cmd.name.startsWith(currentWord) |
| 67 | + ); |
| 68 | + return commands; |
| 69 | + } |
| 70 | + |
| 71 | + async handleChatCommand( |
| 72 | + command: ChatCommand, |
| 73 | + inputModel: IInputModel |
| 74 | + ): Promise<void> { |
| 75 | + // no handling needed because `replaceWith` is set in each command. |
| 76 | + return; |
| 77 | + } |
| 78 | + |
| 79 | + private _contentsManager: Contents.IManager; |
| 80 | + private _docRegistry: DocumentRegistry; |
| 81 | +} |
| 82 | + |
| 83 | +/** |
| 84 | + * Returns the parent path and base name given a path. The parent path will |
| 85 | + * always include a trailing "/" if non-empty. |
| 86 | + * |
| 87 | + * Examples: |
| 88 | + * - "package.json" => ["", "package.json"] |
| 89 | + * - "foo/bar" => ["foo/", "bar"] |
| 90 | + * - "a/b/c/d.txt" => ["a/b/c/", "d.txt"] |
| 91 | + * |
| 92 | + */ |
| 93 | +function getParentAndBase(path: string): [string, string] { |
| 94 | + const components = path.split('/'); |
| 95 | + let parentPath: string; |
| 96 | + let basename: string; |
| 97 | + if (components.length === 1) { |
| 98 | + parentPath = ''; |
| 99 | + basename = components[0]; |
| 100 | + } else { |
| 101 | + parentPath = components.slice(0, -1).join('/') + '/'; |
| 102 | + basename = components[components.length - 1] ?? ''; |
| 103 | + } |
| 104 | + |
| 105 | + return [parentPath, basename]; |
| 106 | +} |
| 107 | + |
| 108 | +async function getPathCompletions( |
| 109 | + contentsManager: Contents.IManager, |
| 110 | + docRegistry: DocumentRegistry, |
| 111 | + searchPath: string |
| 112 | +) { |
| 113 | + const [parentPath, basename] = getParentAndBase(searchPath); |
| 114 | + const parentDir = await contentsManager.get(parentPath); |
| 115 | + const commands: ChatCommand[] = []; |
| 116 | + |
| 117 | + if (!Array.isArray(parentDir.content)) { |
| 118 | + // return nothing if parentDir is invalid / points to a non-directory file |
| 119 | + return []; |
| 120 | + } |
| 121 | + |
| 122 | + const children = parentDir.content |
| 123 | + // filter the children of the parent directory to only include file names that |
| 124 | + // start with the specified base name (case-insensitive). |
| 125 | + .filter((a: Contents.IModel) => { |
| 126 | + return a.name.toLowerCase().startsWith(basename.toLowerCase()); |
| 127 | + }) |
| 128 | + // sort the list, showing directories first while ensuring entries are shown |
| 129 | + // in alphabetic (lexicographically ascending) order. |
| 130 | + .sort((a: Contents.IModel, b: Contents.IModel) => { |
| 131 | + const aPrimaryKey = a.type === 'directory' ? -1 : 1; |
| 132 | + const bPrimaryKey = b.type === 'directory' ? -1 : 1; |
| 133 | + const primaryKey = aPrimaryKey - bPrimaryKey; |
| 134 | + const secondaryKey = a.name < b.name ? -1 : 1; |
| 135 | + |
| 136 | + return primaryKey || secondaryKey; |
| 137 | + }); |
| 138 | + |
| 139 | + for (const child of children) { |
| 140 | + // get icon |
| 141 | + const { icon } = docRegistry.getFileTypeForModel(child); |
| 142 | + |
| 143 | + // compute list of results, while handling directories and non-directories |
| 144 | + // appropriately. |
| 145 | + const isDirectory = child.type === 'directory'; |
| 146 | + let newCommand: ChatCommand; |
| 147 | + if (isDirectory) { |
| 148 | + newCommand = { |
| 149 | + name: child.name + '/', |
| 150 | + providerId: CONTEXT_COMMANDS_PROVIDER_ID, |
| 151 | + icon, |
| 152 | + description: 'Search this directory', |
| 153 | + replaceWith: '@file:' + parentPath + child.name + '/' |
| 154 | + }; |
| 155 | + } else { |
| 156 | + newCommand = { |
| 157 | + name: child.name, |
| 158 | + providerId: CONTEXT_COMMANDS_PROVIDER_ID, |
| 159 | + icon, |
| 160 | + description: 'Attach this file', |
| 161 | + replaceWith: '@file:' + parentPath + child.name + ' ' |
| 162 | + }; |
| 163 | + } |
| 164 | + commands.push(newCommand); |
| 165 | + } |
| 166 | + |
| 167 | + return commands; |
| 168 | +} |
| 169 | + |
| 170 | +export const contextCommandsPlugin: JupyterFrontEndPlugin<void> = { |
| 171 | + id: '@jupyter-ai/core:context-commands-plugin', |
| 172 | + description: 'Adds Jupyter AI context commands to the chat commands menu.', |
| 173 | + autoStart: true, |
| 174 | + requires: [IChatCommandRegistry], |
| 175 | + activate: (app, registry: IChatCommandRegistry) => { |
| 176 | + const { serviceManager, docRegistry } = app; |
| 177 | + registry.addProvider( |
| 178 | + new ContextCommandsProvider(serviceManager.contents, docRegistry) |
| 179 | + ); |
| 180 | + } |
| 181 | +}; |
0 commit comments