Updates

Author: DonJayamanne

src/client/api.ts

@@ -27,7 +27,6 @@ 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>,
@@ -62,12 +61,6 @@ 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.
          */
@@ -124,9 +117,6 @@ export function buildApi(
         jupyter: {
             registerHooks: () => jupyterIntegration.integrateWithJupyterExtension(),
         },
-        pythonEnvironment: {
-            registerApi: registerApi,
-        },
         tensorboard: {
             registerHooks: () => tensorboardIntegration.integrateWithTensorboardExtension(),
         },

src/client/chat/getPythonEnvTool.ts

@@ -20,8 +20,7 @@ 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';
+import { getEnvExtApi, useEnvExtension } from '../envExt/api.internal';
 
 export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
@@ -60,22 +59,20 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         }
 
         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 (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) {

src/client/chat/installPackagesTool.ts

@@ -25,7 +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';
+import { getEnvExtApi, useEnvExtension } from '../envExt/api.internal';
 
 export interface IInstallPackageArgs extends IResourceReference {
     packageList: string[];
@@ -51,8 +51,12 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
             return notebookResponse;
         }
 
-        if (useEnvExtension() && isPrivateApiRegistered()) {
-            await installPackages(resourcePath, options.input.packageList, token);
+        if (useEnvExtension()) {
+            const api = await getEnvExtApi();
+            const env = await api.getEnvironment(resourcePath);
+            if (env) {
+                await raceCancellationError(api.installPackages(env, options.input.packageList), token);
+            }
             const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(', ')}`;
             return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);
         }

src/client/envExt/api.internal.ts

@@ -155,37 +155,3 @@ export function registerEnvExtFeatures(
         );
     }
 }
-
-type PrivateApi = {
-    listPackages: (resource: Uri | undefined, token: CancellationToken) => Promise<Package[] | undefined>;
-    installPackages: (resource: Uri | undefined, packages: string[], token: CancellationToken) => Promise<void>;
-};
-
-let privateApi: PrivateApi | undefined;
-export function registerApi(api: {
-    listPackages: (resource: Uri | undefined, token: CancellationToken) => Promise<Package[] | undefined>;
-    installPackages: (resource: Uri | undefined, packages: string[], token: CancellationToken) => Promise<void>;
-}): void {
-    privateApi = api;
-}
-
-export function isPrivateApiRegistered(): boolean {
-    return !!privateApi;
-}
-
-export function listPackages(resource: Uri | undefined, token: CancellationToken): Promise<Package[] | undefined> {
-    if (privateApi) {
-        return privateApi.listPackages(resource, token);
-    }
-    return Promise.resolve(undefined);
-}
-export function installPackages(
-    resource: Uri | undefined,
-    packages: string[],
-    token: CancellationToken,
-): Promise<void> {
-    if (privateApi) {
-        return privateApi.installPackages(resource, packages, token);
-    }
-    return Promise.resolve();
-}