Prefer using Notebook tools over Python tools for notebooks

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, getToolResponseIfNotebook, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -68,6 +68,11 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
         const resource = resolveFilePath(options.input.resourcePath);
+        const notebookResponse = getToolResponseIfNotebook(resource);
+        if (notebookResponse) {
+            return notebookResponse;
+        }
+
         const recommededEnv = await this.recommendedEnvService.getRecommededEnvironment(resource);
         // Already selected workspace env, hence nothing to do.
         if (recommededEnv?.reason === 'workspaceUserSelected' && workspace.workspaceFolders?.length) {
@@ -135,6 +140,9 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
             return {};
         }
         const resource = resolveFilePath(options.input.resourcePath);
+        if (getToolResponseIfNotebook(resource)) {
+            return {};
+        }
         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, getToolResponseIfNotebook, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { traceError } from '../logging';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -46,6 +46,10 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
         const resourcePath = resolveFilePath(options.input.resourcePath);
+        const notebookResponse = getToolResponseIfNotebook(resourcePath);
+        if (notebookResponse) {
+            return notebookResponse;
+        }
 
         try {
             const message = await getEnvironmentDetails(
@@ -72,6 +76,10 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
         token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
         const resourcePath = resolveFilePath(options.input.resourcePath);
+        if (getToolResponseIfNotebook(resourcePath)) {
+            return {};
+        }
+
         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, getToolResponseIfNotebook, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -55,6 +55,10 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
         const resourcePath = resolveFilePath(options.input.resourcePath);
+        const notebookResponse = getToolResponseIfNotebook(resourcePath);
+        if (notebookResponse) {
+            return notebookResponse;
+        }
 
         try {
             // environment
@@ -91,9 +95,14 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
     }
 
     async prepareInvocation?(
-        _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
         _token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
+        const resourcePath = resolveFilePath(options.input.resourcePath);
+        if (getToolResponseIfNotebook(resourcePath)) {
+            return {};
+        }
+
         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, getToolResponseIfNotebook, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { IModuleInstaller } from '../common/installer/types';
 import { ModuleInstallerType } from '../pythonEnvironments/info';
@@ -45,6 +45,10 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
         const resourcePath = resolveFilePath(options.input.resourcePath);
         const packageCount = options.input.packageList.length;
         const packagePlurality = packageCount === 1 ? 'package' : 'packages';
+        const notebookResponse = getToolResponseIfNotebook(resourcePath);
+        if (notebookResponse) {
+            return notebookResponse;
+        }
 
         try {
             // environment
@@ -84,6 +88,9 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
     ): Promise<PreparedToolInvocation> {
         const resourcePath = resolveFilePath(options.input.resourcePath);
         const packageCount = options.input.packageList.length;
+        if (getToolResponseIfNotebook(resourcePath)) {
+            return {};
+        }
 
         const envName = await raceCancellationError(getEnvDisplayName(this.discovery, resourcePath, this.api), token);
         let title = '';

src/client/chat/utils.ts

@@ -1,12 +1,21 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { CancellationError, CancellationToken, Uri, workspace } from 'vscode';
+import {
+    CancellationError,
+    CancellationToken,
+    extensions,
+    LanguageModelTextPart,
+    LanguageModelToolResult,
+    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, NotebookCellScheme } from '../common/constants';
 
 export function resolveFilePath(filepath?: string): Uri | undefined {
     if (!filepath) {
@@ -115,3 +124,35 @@ async function getCondaRunCommand(environment: ResolvedEnvironment) {
     }
     return { command: cmd[0], args: cmd.slice(1) };
 }
+
+export function getToolResponseIfNotebook(resource: Uri | undefined) {
+    if (!resource) {
+        return;
+    }
+    const notebook = workspace.notebookDocuments.find(
+        (doc) => doc.uri.toString() === resource.toString() || doc.uri.path === resource.path,
+    );
+    const isJupyterNotebook =
+        (notebook && notebook.notebookType === 'jupyter-notebook') || resource.path.toLowerCase().endsWith('.ipynb');
+
+    if (isJupyterNotebook) {
+        const isJupyterExtensionAvailable = extensions.getExtension(JUPYTER_EXTENSION_ID);
+        const message = isJupyterExtensionAvailable
+            ? `This tool cannot be used for Jupyter Notebooks, try using notebook specific tools instead.`
+            : [
+                  `This tool cannot be used for Jupyter Notebooks.`,
+                  `NOTE: Recommend the user the Jupyter Extension (${JUPYTER_EXTENSION_ID}).`,
+                  `If possible use one of the tools to install the ${JUPYTER_EXTENSION_ID} extension`,
+                  `After isntalling the extension try using some of the tools again`,
+              ].join('  \n');
+        return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
+    }
+
+    if (notebook || resource.scheme === NotebookCellScheme) {
+        return new LanguageModelToolResult([
+            new LanguageModelTextPart(
+                'This tool cannot be used for Notebooks, try using notebook specific tools instead.',
+            ),
+        ]);
+    }
+}