All Python lm tools in Python extension

DonJayamanne · DonJayamanne · commit a0d7fa774310 · 2025-05-29T15:15:19.000+10:00
diff --git a/package.json b/package.json
@@ -1491,8 +1491,7 @@
                         }
                     },
                     "required": []
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "get_python_executable_details",
@@ -1516,8 +1515,7 @@
                         }
                     },
                     "required": []
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "install_python_packages",
@@ -1551,8 +1549,7 @@
                     "required": [
                         "packageList"
                     ]
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "configure_python_environment",
@@ -1575,8 +1572,7 @@
                         }
                     },
                     "required": []
-                },
-                "when": "!pythonEnvExtensionInstalled"
+                }
             },
             {
                 "name": "create_virtual_environment",
diff --git a/src/client/api.ts b/src/client/api.ts
@@ -27,6 +27,7 @@ import { ApiForPylance } from './pylanceApi';
 import { getTelemetryReporter } from './telemetry';
 import { TensorboardExtensionIntegration } from './tensorBoard/tensorboardIntegration';
 import { getDebugpyPath } from './debugger/pythonDebugger';
+import { registerApi } from './envExt/api.internal';
 
 export function buildApi(
     ready: Promise<void>,
@@ -61,6 +62,12 @@ export function buildApi(
         jupyter: {
             registerHooks(): void;
         };
+        /**
+         * Internal API just for Python Env Extension untill the API is stabilized.
+         */
+        pythonEnvironment: {
+            registerApi: typeof registerApi;
+        };
         /**
          * Internal API just for Tensorboard, hence don't include in the official types.
          */
@@ -117,6 +124,9 @@ export function buildApi(
         jupyter: {
             registerHooks: () => jupyterIntegration.integrateWithJupyterExtension(),
         },
+        pythonEnvironment: {
+            registerApi: registerApi,
+        },
         tensorboard: {
             registerHooks: () => tensorboardIntegration.integrateWithTensorboardExtension(),
         },
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,8 @@ import { getEnvironmentDetails, getToolResponseIfNotebook, IResourceReference, r
 import { resolveFilePath } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
+import { isPrivateApiRegistered, listPackages, useEnvExtension } from '../envExt/api.internal';
+import { traceError } from '../logging';
 
 export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
@@ -56,14 +58,35 @@ 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() && isPrivateApiRegistered()) {
+            try {
+                const pkgs = await listPackages(resourcePath, token);
+                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');
+                }
+            } catch (ex) {
+                traceError(`Error invoking list_install_python_package_ex tool: ${ex}`);
+            }
+        }
+        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 { installPackages, isPrivateApiRegistered, useEnvExtension } from '../envExt/api.internal';
 
 export interface IInstallPackageArgs extends IResourceReference {
     packageList: string[];
@@ -50,6 +51,12 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
             return notebookResponse;
         }
 
+        if (useEnvExtension() && isPrivateApiRegistered()) {
+            await installPackages(resourcePath, options.input.packageList, token);
+            const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(', ')}`;
+            return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);
+        }
+
         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

Original file line number	Diff line number	Diff line change
`@@ -1491,8 +1491,7 @@`
`1491`	`1491`	`}`
`1492`	`1492`	`},`
`1493`	`1493`	`"required": []`
`1494`		`- },`
`1495`		`- "when": "!pythonEnvExtensionInstalled"`
	`1494`	`+ }`
`1496`	`1495`	`},`
`1497`	`1496`	`{`
`1498`	`1497`	`"name": "get_python_executable_details",`
`@@ -1516,8 +1515,7 @@`
`1516`	`1515`	`}`
`1517`	`1516`	`},`
`1518`	`1517`	`"required": []`
`1519`		`- },`
`1520`		`- "when": "!pythonEnvExtensionInstalled"`
	`1518`	`+ }`
`1521`	`1519`	`},`
`1522`	`1520`	`{`
`1523`	`1521`	`"name": "install_python_packages",`
`@@ -1551,8 +1549,7 @@`
`1551`	`1549`	`"required": [`
`1552`	`1550`	`"packageList"`
`1553`	`1551`	`]`
`1554`		`- },`
`1555`		`- "when": "!pythonEnvExtensionInstalled"`
	`1552`	`+ }`
`1556`	`1553`	`},`
`1557`	`1554`	`{`
`1558`	`1555`	`"name": "configure_python_environment",`
`@@ -1575,8 +1572,7 @@`
`1575`	`1572`	`}`
`1576`	`1573`	`},`
`1577`	`1574`	`"required": []`
`1578`		`- },`
`1579`		`- "when": "!pythonEnvExtensionInstalled"`
	`1575`	`+ }`
`1580`	`1576`	`},`
`1581`	`1577`	`{`
`1582`	`1578`	`"name": "create_virtual_environment",`