Ensure notebook tools are recommended when attempting to use Python tools for Notebooks

Author: DonJayamanne

src/client/chat/configurePythonEnvTool.ts

@@ -19,7 +19,7 @@ import { PythonExtension, ResolvedEnvironment } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
-import { getEnvironmentDetails, raceCancellationError } from './utils';
+import { getEnvironmentDetails, raceCancellationError, throwIfNotebookUri } from './utils';
 import { resolveFilePath } from './utils';
 import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -135,6 +135,7 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
             return {};
         }
         const resource = resolveFilePath(options.input.resourcePath);
+        throwIfNotebookUri(resource);
         const recommededEnv = await this.recommendedEnvService.getRecommededEnvironment(resource);
         // Already selected workspace env, hence nothing to do.
         if (recommededEnv?.reason === 'workspaceUserSelected' && workspace.workspaceFolders?.length) {

src/client/chat/getExecutableTool.ts

@@ -16,7 +16,7 @@ import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
-import { getEnvDisplayName, getEnvironmentDetails, raceCancellationError } from './utils';
+import { getEnvDisplayName, getEnvironmentDetails, raceCancellationError, throwIfNotebookUri } from './utils';
 import { resolveFilePath } from './utils';
 import { traceError } from '../logging';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -72,6 +72,7 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
         token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
         const resourcePath = resolveFilePath(options.input.resourcePath);
+        throwIfNotebookUri(resourcePath);
         const envName = await raceCancellationError(getEnvDisplayName(this.discovery, resourcePath, this.api), token);
         return {
             invocationMessage: envName

src/client/chat/getPythonEnvTool.ts

@@ -17,7 +17,7 @@ import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
 import { IProcessServiceFactory, IPythonExecutionFactory } from '../common/process/types';
-import { getEnvironmentDetails, raceCancellationError } from './utils';
+import { getEnvironmentDetails, raceCancellationError, throwIfNotebookUri } from './utils';
 import { resolveFilePath } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -91,9 +91,11 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
     }
 
     async prepareInvocation?(
-        _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
         _token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
+        const resourcePath = resolveFilePath(options.input.resourcePath);
+        throwIfNotebookUri(resourcePath);
         return {
             invocationMessage: l10n.t('Fetching Python environment information'),
         };

src/client/chat/installPackagesTool.ts

@@ -14,7 +14,7 @@ import {
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
-import { getEnvDisplayName, raceCancellationError } from './utils';
+import { getEnvDisplayName, raceCancellationError, throwIfNotebookUri } from './utils';
 import { resolveFilePath } from './utils';
 import { IModuleInstaller } from '../common/installer/types';
 import { ModuleInstallerType } from '../pythonEnvironments/info';
@@ -84,6 +84,7 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
     ): Promise<PreparedToolInvocation> {
         const resourcePath = resolveFilePath(options.input.resourcePath);
         const packageCount = options.input.packageList.length;
+        throwIfNotebookUri(resourcePath);
 
         const envName = await raceCancellationError(getEnvDisplayName(this.discovery, resourcePath, this.api), token);
         let title = '';

src/client/chat/utils.ts

@@ -1,12 +1,13 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { CancellationError, CancellationToken, Uri, workspace } from 'vscode';
+import { CancellationError, CancellationToken, extensions, NotebookDocument, Uri, workspace } from 'vscode';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
 import { PythonExtension, ResolvedEnvironment } from '../api/types';
 import { ITerminalHelper, TerminalShellType } from '../common/terminal/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
 import { Conda } from '../pythonEnvironments/common/environmentManagers/conda';
+import { JUPYTER_EXTENSION_ID } from '../common/constants';
 
 export function resolveFilePath(filepath?: string): Uri | undefined {
     if (!filepath) {
@@ -115,3 +116,30 @@ async function getCondaRunCommand(environment: ResolvedEnvironment) {
     }
     return { command: cmd[0], args: cmd.slice(1) };
 }
+
+export function throwIfNotebookUri(resource: Uri | undefined) {
+    if (!resource) {
+        return;
+    }
+    const notebook = workspace.notebookDocuments.find(
+        (doc) => doc.uri.toString() === resource.toString() || doc.uri.path === resource.path,
+    );
+    if ((notebook && isJupyterNotebook(notebook)) || resource.path.toLowerCase().endsWith('.ipynb')) {
+        const isJupyterExtensionAvailable = extensions.getExtension(JUPYTER_EXTENSION_ID);
+        if (isJupyterExtensionAvailable) {
+            throw new Error(
+                'This tool cannot be used for Jupyter Notebooks, try using notebook specific tools instead.',
+            );
+        }
+        throw new Error(
+            `This tool cannot be used for Jupyter Notebooks. Install the Jupyter Extension (${JUPYTER_EXTENSION_ID}) & try using notebook specific tools instead.`,
+        );
+    }
+    if (notebook) {
+        throw new Error('This tool cannot be used for Notebooks, try using notebook specific tools instead.');
+    }
+}
+
+function isJupyterNotebook(notebook: NotebookDocument) {
+    return notebook.notebookType === 'jupyter-notebook';
+}