All Python lm tools in Python extension (#25116)

DonJayamanne · web-flow · commit 9b17438067c6 · 2025-05-29T21:49:17.000Z
diff --git a/package.json b/package.json
@@ -1474,9 +1474,10 @@
                 "displayName": "Get Python Environment Info",
                 "userDescription": "%python.languageModelTools.get_python_environment_details.userDescription%",
                 "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",
+                "toolReferenceName": "getPythonEnvironmentInfo",
                 "tags": [
                     "python",
+                    "python environment",
                     "extension_installed_by_tool",
                     "enable_other_tool_configure_python_environment"
                 ],
@@ -1491,17 +1492,17 @@
                         }
                     },
                     "required": []
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "get_python_executable_details",
                 "displayName": "Get Python Executable",
                 "userDescription": "%python.languageModelTools.get_python_executable_details.userDescription%",
                 "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",
+                "toolReferenceName": "getPythonExecutableCommand",
                 "tags": [
                     "python",
+                    "python environment",
                     "extension_installed_by_tool",
                     "enable_other_tool_configure_python_environment"
                 ],
@@ -1516,17 +1517,17 @@
                         }
                     },
                     "required": []
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "install_python_packages",
                 "displayName": "Install Python Package",
                 "userDescription": "%python.languageModelTools.install_python_packages.userDescription%",
                 "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",
+                "toolReferenceName": "installPythonPackage",
                 "tags": [
                     "python",
+                    "python environment",
                     "install python package",
                     "extension_installed_by_tool",
                     "enable_other_tool_configure_python_environment"
@@ -1551,8 +1552,7 @@
                     "required": [
                         "packageList"
                     ]
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "configure_python_environment",
@@ -1562,6 +1562,7 @@
                 "toolReferenceName": "configurePythonEnvironment",
                 "tags": [
                     "python",
+                    "python environment",
                     "extension_installed_by_tool"
                 ],
                 "icon": "$(gear)",
@@ -1575,8 +1576,7 @@
                         }
                     },
                     "required": []
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "create_virtual_environment",
diff --git a/src/client/chat/configurePythonEnvTool.ts b/src/client/chat/configurePythonEnvTool.ts
@@ -28,7 +28,6 @@ import { ITerminalHelper } from '../common/terminal/types';
 import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
 import { CreateVirtualEnvTool } from './createVirtualEnvTool';
 import { ISelectPythonEnvToolArguments, SelectPythonEnvTool } from './selectEnvTool';
-import { useEnvExtension } from '../envExt/api.internal';
 
 export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
@@ -78,10 +77,7 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
 
         if (await this.createEnvTool.shouldCreateNewVirtualEnv(resource, token)) {
             try {
-                // If the Python Env extension is available, then use that.
-                // create_quick_virtual_environment
-                const toolName = useEnvExtension() ? 'create_quick_virtual_environment' : CreateVirtualEnvTool.toolName;
-                return await lm.invokeTool(toolName, options, token);
+                return await lm.invokeTool(CreateVirtualEnvTool.toolName, options, token);
             } catch (ex) {
                 if (isCancellationError(ex)) {
                     const input: ISelectPythonEnvToolArguments = {
diff --git a/src/client/chat/createVirtualEnvTool.ts b/src/client/chat/createVirtualEnvTool.ts
@@ -4,6 +4,7 @@
 import {
     CancellationError,
     CancellationToken,
+    commands,
     l10n,
     LanguageModelTool,
     LanguageModelToolInvocationOptions,
@@ -39,6 +40,8 @@ import { isStableVersion } from '../pythonEnvironments/info/pythonVersion';
 import { createVirtualEnvironment } from '../pythonEnvironments/creation/createEnvApi';
 import { traceError, traceVerbose, traceWarn } from '../logging';
 import { StopWatch } from '../common/utils/stopWatch';
+import { useEnvExtension } from '../envExt/api.internal';
+import { PythonEnvironment } from '../envExt/types';
 
 interface ICreateVirtualEnvToolParams extends IResourceReference {
     packageList?: string[]; // Added only becausewe have ability to create a virtual env with list of packages same tool within the in Python Env extension.
@@ -66,7 +69,7 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
     }
 
     async invoke(
-        options: LanguageModelToolInvocationOptions<IResourceReference>,
+        options: LanguageModelToolInvocationOptions<ICreateVirtualEnvToolParams>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
         const resource = resolveFilePath(options.input.resourcePath);
@@ -83,14 +86,26 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
                 disposables.add(interpreterPathService.onDidChange(() => resolve()));
             });
 
-            const created = await raceCancellationError(
-                createVirtualEnvironment({
-                    interpreter: preferredGlobalPythonEnv.id,
-                    workspaceFolder,
-                }),
-                token,
-            );
-            if (!created?.path) {
+            let createdEnvPath: string | undefined = undefined;
+            if (useEnvExtension()) {
+                const result: PythonEnvironment | undefined = await commands.executeCommand('python-envs.createAny', {
+                    quickCreate: true,
+                    additionalPackages: options.input.packageList || [],
+                    uri: workspaceFolder.uri,
+                    selectEnvironment: true,
+                });
+                createdEnvPath = result?.environmentPath.fsPath;
+            } else {
+                const created = await raceCancellationError(
+                    createVirtualEnvironment({
+                        interpreter: preferredGlobalPythonEnv.id,
+                        workspaceFolder,
+                    }),
+                    token,
+                );
+                createdEnvPath = created?.path;
+            }
+            if (!createdEnvPath) {
                 traceWarn(`${CreateVirtualEnvTool.toolName} tool not invoked, virtual env not created.`);
                 throw new CancellationError();
             }
@@ -102,7 +117,7 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
             const stopWatch = new StopWatch();
             let env: ResolvedEnvironment | undefined;
             while (stopWatch.elapsedTime < 5_000 || !env) {
-                env = await this.api.resolveEnvironment(created.path);
+                env = await this.api.resolveEnvironment(createdEnvPath);
                 if (env) {
                     break;
                 } else {
@@ -150,7 +165,7 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
     }
 
     async prepareInvocation?(
-        options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        options: LanguageModelToolInvocationPrepareOptions<ICreateVirtualEnvToolParams>,
         token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
         const resource = resolveFilePath(options.input.resourcePath);
diff --git a/src/client/chat/getPythonEnvTool.ts b/src/client/chat/getPythonEnvTool.ts
@@ -20,6 +20,7 @@ import { getEnvironmentDetails, getToolResponseIfNotebook, IResourceReference, r
 import { resolveFilePath } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
+import { getEnvExtApi, useEnvExtension } from '../envExt/api.internal';
 
 export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
@@ -56,14 +57,33 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         if (!environment || !environment.version) {
             throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
         }
-        const packages = await getPythonPackagesResponse(
-            environment,
-            this.pythonExecFactory,
-            this.processServiceFactory,
-            resourcePath,
-            token,
-        );
 
+        let packages = '';
+        if (useEnvExtension()) {
+            const api = await getEnvExtApi();
+            const env = await api.getEnvironment(resourcePath);
+            const pkgs = env ? await api.getPackages(env) : [];
+            if (pkgs && pkgs.length > 0) {
+                // Installed Python packages, each in the format <name> or <name> (<version>). The version may be omitted if unknown. Returns an empty array if no packages are installed.
+                const response = [
+                    'Below is a list of the Python packages, each in the format <name> or <name> (<version>). The version may be omitted if unknown: ',
+                ];
+                pkgs.forEach((pkg) => {
+                    const version = pkg.version;
+                    response.push(version ? `- ${pkg.name} (${version})` : `- ${pkg.name}`);
+                });
+                packages = response.join('\n');
+            }
+        }
+        if (!packages) {
+            packages = await getPythonPackagesResponse(
+                environment,
+                this.pythonExecFactory,
+                this.processServiceFactory,
+                resourcePath,
+                token,
+            );
+        }
         const message = await getEnvironmentDetails(
             resourcePath,
             this.api,
diff --git a/src/client/chat/index.ts b/src/client/chat/index.ts
@@ -1,13 +1,12 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { commands, extensions, lm } from 'vscode';
+import { lm } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { InstallPackagesTool } from './installPackagesTool';
 import { IExtensionContext } from '../common/types';
 import { DisposableStore } from '../common/utils/resourceLifecycle';
-import { ENVS_EXTENSION_ID } from '../envExt/api.internal';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
 import { GetExecutableTool } from './getExecutableTool';
 import { GetEnvironmentInfoTool } from './getPythonEnvTool';
@@ -21,11 +20,6 @@ export function registerTools(
     environmentsApi: PythonExtension['environments'],
     serviceContainer: IServiceContainer,
 ) {
-    if (extensions.getExtension(ENVS_EXTENSION_ID)) {
-        return;
-    }
-    const contextKey = 'pythonEnvExtensionInstalled';
-    commands.executeCommand('setContext', contextKey, false);
     const ourTools = new DisposableStore();
     context.subscriptions.push(ourTools);
 
@@ -55,14 +49,4 @@ export function registerTools(
             new ConfigurePythonEnvTool(environmentsApi, serviceContainer, createVirtualEnvTool),
         ),
     );
-    ourTools.add(
-        extensions.onDidChange(() => {
-            const envExtension = extensions.getExtension(ENVS_EXTENSION_ID);
-            if (envExtension) {
-                envExtension.activate();
-                commands.executeCommand('setContext', contextKey, true);
-                ourTools.dispose();
-            }
-        }),
-    );
 }
diff --git a/src/client/chat/installPackagesTool.ts b/src/client/chat/installPackagesTool.ts
@@ -25,6 +25,7 @@ import { resolveFilePath } from './utils';
 import { IModuleInstaller } from '../common/installer/types';
 import { ModuleInstallerType } from '../pythonEnvironments/info';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
+import { getEnvExtApi, useEnvExtension } from '../envExt/api.internal';
 
 export interface IInstallPackageArgs extends IResourceReference {
     packageList: string[];
@@ -50,6 +51,24 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
             return notebookResponse;
         }
 
+        if (useEnvExtension()) {
+            const api = await getEnvExtApi();
+            const env = await api.getEnvironment(resourcePath);
+            if (env) {
+                await raceCancellationError(api.managePackages(env, { install: options.input.packageList }), token);
+                const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(
+                    ', ',
+                )}`;
+                return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);
+            } else {
+                return new LanguageModelToolResult([
+                    new LanguageModelTextPart(
+                        `Packages not installed. No environment found for: ${resourcePath?.fsPath}`,
+                    ),
+                ]);
+            }
+        }
+
         try {
             // environment
             const envPath = this.api.getActiveEnvironmentPath(resourcePath);
diff --git a/src/client/chat/utils.ts b/src/client/chat/utils.ts
@@ -17,6 +17,7 @@ import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/termin
 import { Conda } from '../pythonEnvironments/common/environmentManagers/conda';
 import { JUPYTER_EXTENSION_ID, NotebookCellScheme } from '../common/constants';
 import { dirname, join } from 'path';
+import { resolveEnvironment, useEnvExtension } from '../envExt/api.internal';
 
 export interface IResourceReference {
     resourcePath?: string;
@@ -76,18 +77,46 @@ export async function getEnvironmentDetails(
 ): Promise<string> {
     // environment
     const envPath = api.getActiveEnvironmentPath(resourcePath);
-    const environment = await raceCancellationError(api.resolveEnvironment(envPath), token);
-    if (!environment || !environment.version) {
-        throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
+    let envType = '';
+    let envVersion = '';
+    let runCommand = '';
+    if (useEnvExtension()) {
+        const environment =
+            (await raceCancellationError(resolveEnvironment(envPath.id), token)) ||
+            (await raceCancellationError(resolveEnvironment(envPath.path), token));
+        if (!environment || !environment.version) {
+            throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
+        }
+        envVersion = environment.version;
+        try {
+            const managerId = environment.envId.managerId;
+            envType =
+                (!managerId.endsWith(':') && managerId.includes(':') ? managerId.split(':').reverse()[0] : '') ||
+                'unknown';
+        } catch {
+            envType = 'unknown';
+        }
+
+        const execInfo = environment.execInfo;
+        const executable = execInfo?.activatedRun?.executable ?? execInfo?.run.executable ?? 'python';
+        const args = execInfo?.activatedRun?.args ?? execInfo?.run.args ?? [];
+        runCommand = terminalHelper.buildCommandForTerminal(TerminalShellType.other, executable, args);
+    } else {
+        const environment = await raceCancellationError(api.resolveEnvironment(envPath), token);
+        if (!environment || !environment.version) {
+            throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
+        }
+        envType = environment.environment?.type || 'unknown';
+        envVersion = environment.version.sysVersion || 'unknown';
+        runCommand = await raceCancellationError(
+            getTerminalCommand(environment, resourcePath, terminalExecutionService, terminalHelper),
+            token,
+        );
     }
-    const runCommand = await raceCancellationError(
-        getTerminalCommand(environment, resourcePath, terminalExecutionService, terminalHelper),
-        token,
-    );
     const message = [
         `Following is the information about the Python environment:`,
-        `1. Environment Type: ${environment.environment?.type || 'unknown'}`,
-        `2. Version: ${environment.version.sysVersion || 'unknown'}`,
+        `1. Environment Type: ${envType}`,
+        `2. Version: ${envVersion}`,
         '',
         `3. Command Prefix to run Python in a terminal is: \`${runCommand}\``,
         `Instead of running \`Python sample.py\` in the terminal, you will now run: \`${runCommand} sample.py\``,
@@ -183,7 +212,8 @@ export function doesWorkspaceHaveVenvOrCondaEnv(resource: Uri | undefined, api:
             env.environment?.folderUri &&
             env.executable.sysPrefix &&
             dirname(env.executable.sysPrefix) === workspaceFolder.uri.fsPath &&
-            env.environment.name === '.venv' &&
+            ((env.environment.name || '').startsWith('.venv') ||
+                env.executable.sysPrefix === join(workspaceFolder.uri.fsPath, '.venv')) &&
             env.environment.type === 'VirtualEnvironment'
         );
     };
@@ -192,7 +222,8 @@ export function doesWorkspaceHaveVenvOrCondaEnv(resource: Uri | undefined, api:
             env.environment?.folderUri &&
             env.executable.sysPrefix &&
             dirname(env.executable.sysPrefix) === workspaceFolder.uri.fsPath &&
-            env.environment.folderUri.fsPath === join(workspaceFolder.uri.fsPath, '.conda') &&
+            (env.environment.folderUri.fsPath === join(workspaceFolder.uri.fsPath, '.conda') ||
+                env.executable.sysPrefix === join(workspaceFolder.uri.fsPath, '.conda')) &&
             env.environment.type === 'Conda'
         );
     };
diff --git a/src/client/envExt/api.internal.ts b/src/client/envExt/api.internal.ts
diff --git a/src/client/envExt/api.legacy.ts b/src/client/envExt/api.legacy.ts
diff --git a/src/client/envExt/types.ts b/src/client/envExt/types.ts
