Configuration tool with improved workflow and UX

Author: DonJayamanne

package.json

@@ -1475,6 +1475,7 @@
                 "modelDescription": "This tool will retrieve the details of the Python Environment for the specified file or workspace. The details returned include the 1. Type of Environment (conda, venv, etec), 2. Version of Python, 3. List of all installed packages with their versions. ALWAYS call configure_python_environment before using this tool.",
                 "toolReferenceName": "pythonGetEnvironmentInfo",
                 "tags": [
+                    "extension_installed_by_tool",
                     "enable_other_tool_configure_python_environment"
                 ],
                 "icon": "$(snake)",
@@ -1498,6 +1499,7 @@
                 "modelDescription": "This tool will retrieve the details of the Python Environment for the specified file or workspace. ALWAYS use this tool before executing any Python command in the terminal. This tool returns the details of how to construct the fully qualified path and or command including details such as arguments required to run Python in a terminal. Note: Instead of executing `python --version` or `python -c 'import sys; print(sys.executable)'`, use this tool to get the Python executable path to replace the `python` command. E.g. instead of using `python -c 'import sys; print(sys.executable)'`, use this tool to build the command `conda run -n <env_name> -c 'import sys; print(sys.executable)'`. ALWAYS call configure_python_environment before using this tool.",
                 "toolReferenceName": "pythonExecutableCommand",
                 "tags": [
+                    "extension_installed_by_tool",
                     "enable_other_tool_configure_python_environment"
                 ],
                 "icon": "$(terminal)",
@@ -1521,6 +1523,7 @@
                 "modelDescription": "Installs Python packages in the given workspace. Use this tool to install packages in the user's chosen environment. ALWAYS call configure_python_environment before using this tool.",
                 "toolReferenceName": "pythonInstallPackage",
                 "tags": [
+                    "extension_installed_by_tool",
                     "enable_other_tool_configure_python_environment"
                 ],
                 "icon": "$(package)",
@@ -1552,7 +1555,9 @@
                 "modelDescription": "This tool configures a Python environment in the given workspace. ALWAYS Use this tool to set up the user's chosen environment and ALWAYS call this tool before using any other Python related tools.",
                 "userDescription": "%python.languageModelTools.configure_python_environment.userDescription%",
                 "toolReferenceName": "configurePythonEnvironment",
-                "tags": [],
+                "tags": [
+                    "extension_installed_by_tool"
+                ],
                 "icon": "$(gear)",
                 "canBeReferencedInPrompt": true,
                 "inputSchema": {
@@ -1566,6 +1571,46 @@
                     "required": []
                 },
                 "when": "!pythonEnvExtensionInstalled"
+            },
+            {
+                "name": "create_virtual_environment",
+                "displayName": "Create a Virtual Environment",
+                "modelDescription": "This tool will create a Virual Environment",
+                "tags": [
+                    "extension_installed_by_tool"
+                ],
+                "canBeReferencedInPrompt": false,
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {
+                        "resourcePath": {
+                            "type": "string",
+                            "description": "The path to the Python file or workspace for which a Python Environment needs to be configured."
+                        }
+                    },
+                    "required": []
+                },
+                "when": "false"
+            },
+            {
+                "name": "selectEnvironment",
+                "displayName": "Select a Python Environment",
+                "modelDescription": "This tool will prompt the user to select an existing Python Environment",
+                "tags": [
+                    "extension_installed_by_tool"
+                ],
+                "canBeReferencedInPrompt": false,
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {
+                        "resourcePath": {
+                            "type": "string",
+                            "description": "The path to the Python file or workspace for which a Python Environment needs to be configured."
+                        }
+                    },
+                    "required": []
+                },
+                "when": "false"
             }
         ]
     },

src/client/chat/configurePythonEnvTool.ts

@@ -3,48 +3,40 @@
 
 import {
     CancellationToken,
-    l10n,
-    LanguageModelTextPart,
     LanguageModelTool,
     LanguageModelToolInvocationOptions,
     LanguageModelToolInvocationPrepareOptions,
     LanguageModelToolResult,
     PreparedToolInvocation,
     Uri,
     workspace,
-    commands,
-    QuickPickItem,
+    lm,
 } from 'vscode';
-import { PythonExtension, ResolvedEnvironment } from '../api/types';
+import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
-import { getEnvironmentDetails, getToolResponseIfNotebook, raceCancellationError } from './utils';
+import {
+    getEnvDetailsForResponse,
+    getToolResponseIfNotebook,
+    IResourceReference,
+    isCancellationError,
+    raceCancellationError,
+} from './utils';
 import { resolveFilePath } from './utils';
-import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
 import { ITerminalHelper } from '../common/terminal/types';
-import { raceTimeout } from '../common/utils/async';
-import { Commands, Octicons } from '../common/constants';
-import { CreateEnvironmentResult } from '../pythonEnvironments/creation/proposed.createEnvApis';
-import { IInterpreterPathService } from '../common/types';
-import { DisposableStore } from '../common/utils/resourceLifecycle';
-import { Common, InterpreterQuickPickList } from '../common/utils/localize';
-import { QuickPickItemKind } from '../../test/mocks/vsc';
-import { showQuickPick } from '../common/vscodeApis/windowApis';
-import { SelectEnvironmentResult } from '../interpreter/configuration/interpreterSelector/commands/setInterpreter';
-
-export interface IResourceReference {
-    resourcePath?: string;
-}
-
-let _environmentConfigured = false;
+import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
+import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
+import { CreateVirtualEnvTool } from './createVirtualEnvTool';
+import { SelectPythonEnvTool } from './selectEnvTool';
 
 export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
     private readonly terminalHelper: ITerminalHelper;
     private readonly recommendedEnvService: IRecommendedEnvironmentService;
     public static readonly toolName = 'configure_python_environment';
     constructor(
+        private readonly discoveryApi: IDiscoveryAPI,
         private readonly api: PythonExtension['environments'],
         private readonly serviceContainer: IServiceContainer,
     ) {
@@ -57,12 +49,7 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
             IRecommendedEnvironmentService,
         );
     }
-    /**
-     * Invokes the tool to get the information about the Python environment.
-     * @param options - The invocation options containing the file path.
-     * @param token - The cancellation token.
-     * @returns The result containing the information about the Python environment or an error message.
-     */
+
     async invoke(
         options: LanguageModelToolInvocationOptions<IResourceReference>,
         token: CancellationToken,
@@ -73,22 +60,14 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
             return notebookResponse;
         }
 
-        const recommededEnv = await this.recommendedEnvService.getRecommededEnvironment(resource);
-        // Already selected workspace env, hence nothing to do.
-        if (recommededEnv?.reason === 'workspaceUserSelected' && workspace.workspaceFolders?.length) {
-            return await getEnvDetailsForResponse(
-                recommededEnv.environment,
-                this.api,
-                this.terminalExecutionService,
-                this.terminalHelper,
-                resource,
-                token,
-            );
-        }
-        // No workspace folders, and the user selected a global environment.
-        if (recommededEnv?.reason === 'globalUserSelected' && !workspace.workspaceFolders?.length) {
-            return await getEnvDetailsForResponse(
-                recommededEnv.environment,
+        const workspaceSpecificEnv = await raceCancellationError(
+            this.hasAlreadyGotAWorkspaceSpecificEnvironment(resource),
+            token,
+        );
+
+        if (workspaceSpecificEnv) {
+            return getEnvDetailsForResponse(
+                workspaceSpecificEnv,
                 this.api,
                 this.terminalExecutionService,
                 this.terminalHelper,
@@ -97,174 +76,46 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
             );
         }
 
-        if (!workspace.workspaceFolders?.length) {
-            const selected = await Promise.resolve(commands.executeCommand(Commands.Set_Interpreter));
-            const env = await this.api.resolveEnvironment(this.api.getActiveEnvironmentPath(resource));
-            if (selected && env) {
-                return await getEnvDetailsForResponse(
-                    env,
-                    this.api,
-                    this.terminalExecutionService,
-                    this.terminalHelper,
-                    resource,
-                    token,
-                );
+        let reason: 'cancelled' | undefined;
+        if (
+            // eslint-disable-next-line @typescript-eslint/no-use-before-define
+            await new CreateVirtualEnvTool(this.discoveryApi, this.api, this.serviceContainer).canCreateNewVirtualEnv(
+                resolveFilePath(options.input.resourcePath),
+                token,
+            )
+        ) {
+            reason = 'cancelled';
+            try {
+                return await lm.invokeTool(CreateVirtualEnvTool.toolName, options, token);
+            } catch (ex) {
+                // If the user cancelled the tool, then we should not invoke the select env tool.
+                if (!isCancellationError(ex)) {
+                    throw ex;
+                }
             }
-            return new LanguageModelToolResult([
-                new LanguageModelTextPart('User did not select a Python environment.'),
-            ]);
         }
 
-        const selected = await showCreateAndSelectEnvironmentQuickPick(resource, this.serviceContainer);
-        const env = await this.api.resolveEnvironment(this.api.getActiveEnvironmentPath(resource));
-        if (selected && env) {
-            return await getEnvDetailsForResponse(
-                env,
-                this.api,
-                this.terminalExecutionService,
-                this.terminalHelper,
-                resource,
-                token,
-            );
-        }
-        return new LanguageModelToolResult([
-            new LanguageModelTextPart('User did not create nor select a Python environment.'),
-        ]);
+        return lm.invokeTool(SelectPythonEnvTool.toolName, { ...options, input: { ...options.input, reason } }, token);
     }
 
     async prepareInvocation?(
-        options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
         _token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
-        if (_environmentConfigured) {
-            return {};
-        }
-        const resource = resolveFilePath(options.input.resourcePath);
-        if (getToolResponseIfNotebook(resource)) {
-            return {};
-        }
+        return {
+            invocationMessage: 'Configuring a Python Environment',
+        };
+    }
+
+    async hasAlreadyGotAWorkspaceSpecificEnvironment(resource: Uri | undefined) {
         const recommededEnv = await this.recommendedEnvService.getRecommededEnvironment(resource);
         // Already selected workspace env, hence nothing to do.
         if (recommededEnv?.reason === 'workspaceUserSelected' && workspace.workspaceFolders?.length) {
-            return {};
+            return recommededEnv.environment;
         }
         // No workspace folders, and the user selected a global environment.
         if (recommededEnv?.reason === 'globalUserSelected' && !workspace.workspaceFolders?.length) {
-            return {};
-        }
-
-        if (!workspace.workspaceFolders?.length) {
-            return {
-                confirmationMessages: {
-                    title: l10n.t('Configure a Python Environment?'),
-                    message: l10n.t('You will be prompted to select a Python Environment.'),
-                },
-            };
-        }
-        return {
-            confirmationMessages: {
-                title: l10n.t('Configure a Python Environment?'),
-                message: l10n.t(
-                    [
-                        'The recommended option is to create a new Python Environment, providing the benefit of isolating packages from other environments.  ',
-                        'Optionally you could select an existing Python Environment.',
-                    ].join('\n'),
-                ),
-            },
-        };
-    }
-}
-
-async function getEnvDetailsForResponse(
-    environment: ResolvedEnvironment | undefined,
-    api: PythonExtension['environments'],
-    terminalExecutionService: TerminalCodeExecutionProvider,
-    terminalHelper: ITerminalHelper,
-    resource: Uri | undefined,
-    token: CancellationToken,
-): Promise<LanguageModelToolResult> {
-    const envPath = api.getActiveEnvironmentPath(resource);
-    environment = environment || (await raceCancellationError(api.resolveEnvironment(envPath), token));
-    if (!environment || !environment.version) {
-        throw new Error('No environment found for the provided resource path: ' + resource?.fsPath);
-    }
-    const message = await getEnvironmentDetails(
-        resource,
-        api,
-        terminalExecutionService,
-        terminalHelper,
-        undefined,
-        token,
-    );
-    return new LanguageModelToolResult([
-        new LanguageModelTextPart(`A Python Environment has been configured.  \n` + message),
-    ]);
-}
-
-async function showCreateAndSelectEnvironmentQuickPick(
-    uri: Uri | undefined,
-    serviceContainer: IServiceContainer,
-): Promise<boolean | undefined> {
-    const createLabel = `${Octicons.Add} ${InterpreterQuickPickList.create.label}`;
-    const selectLabel = l10n.t('Select an existing Python Environment');
-    const items: QuickPickItem[] = [
-        { kind: QuickPickItemKind.Separator, label: Common.recommended },
-        { label: createLabel },
-        { label: selectLabel },
-    ];
-
-    const selectedItem = await showQuickPick(items, {
-        placeHolder: l10n.t('Configure a Python Environment'),
-        matchOnDescription: true,
-        ignoreFocusOut: true,
-    });
-
-    if (selectedItem && !Array.isArray(selectedItem) && selectedItem.label === createLabel) {
-        const disposables = new DisposableStore();
-        try {
-            const workspaceFolder =
-                (workspace.workspaceFolders?.length && uri ? workspace.getWorkspaceFolder(uri) : undefined) ||
-                (workspace.workspaceFolders?.length === 1 ? workspace.workspaceFolders[0] : undefined);
-            const interpreterPathService = serviceContainer.get<IInterpreterPathService>(IInterpreterPathService);
-            const interpreterChanged = new Promise<void>((resolve) => {
-                disposables.add(interpreterPathService.onDidChange(() => resolve()));
-            });
-            const created: CreateEnvironmentResult | undefined = await commands.executeCommand(
-                Commands.Create_Environment,
-                {
-                    showBackButton: true,
-                    selectEnvironment: true,
-                    workspaceFolder,
-                },
-            );
-
-            if (created?.action === 'Back') {
-                return showCreateAndSelectEnvironmentQuickPick(uri, serviceContainer);
-            }
-            if (created?.action === 'Cancel') {
-                return undefined;
-            }
-            if (created?.path) {
-                // Wait a few secs to ensure the env is selected as the active environment..
-                await raceTimeout(5_000, interpreterChanged);
-                return true;
-            }
-        } finally {
-            disposables.dispose();
-        }
-    }
-    if (selectedItem && !Array.isArray(selectedItem) && selectedItem.label === selectLabel) {
-        const result = (await Promise.resolve(
-            commands.executeCommand(Commands.Set_Interpreter, { hideCreateVenv: true, showBackButton: true }),
-        )) as SelectEnvironmentResult | undefined;
-        if (result?.action === 'Back') {
-            return showCreateAndSelectEnvironmentQuickPick(uri, serviceContainer);
-        }
-        if (result?.action === 'Cancel') {
-            return undefined;
-        }
-        if (result?.path) {
-            return true;
+            return recommededEnv.environment;
         }
     }
 }