Copilot tools (#280)

eleanorjboyd · web-flow · commit 9dd0fd9dd6d9 · 2025-04-11T18:02:52.000Z
adds support for two copilot tools:
- get environment information
- install packages
diff --git a/package.json b/package.json
@@ -491,22 +491,52 @@
         ],
         "languageModelTools": [
             {
-                "name": "python_get_packages",
-                "displayName": "Get Python Packages",
-                "modelDescription": "Returns the packages installed in the given Python file's environment. You should call this when you want to generate Python code to determine the users preferred packages. Also call this to determine if you need to provide installation instructions in a response.",
-                "toolReferenceName": "pythonGetPackages",
+                "name": "python_environment_tool",
+                "displayName": "Get Python Environment Information",
+                "modelDescription": "Provides details about the Python environment for a specified file or workspace, including environment type, Python version, run command, and installed packages with their versions.",
+                "toolReferenceName": "pythonGetEnvironmentInfo",
                 "tags": [],
                 "icon": "$(files)",
+                "canBeReferencedInPrompt": true,
                 "inputSchema": {
                     "type": "object",
                     "properties": {
-                        "filePath": {
+                        "resourcePath": {
                             "type": "string"
                         }
                     },
-                    "description": "The path to the Python file or workspace to get the installed packages for.",
+                    "description": "The path to the Python file or workspace to get the environment information for.",
                     "required": [
-                        "filePath"
+                        "resourcePath"
+                    ]
+                }
+            },
+            {
+                "name": "python_install_package_tool",
+                "displayName": "Install Python Package",
+                "modelDescription": "Installs Python packages in the given workspace. Use this tool to install packages in the user's chosen environment.",
+                "toolReferenceName": "pythonInstallPackage",
+                "tags": [],
+                "icon": "$(package)",
+                "canBeReferencedInPrompt": true,
+                "inputSchema": {
+                    "type": "object",
+                    "properties": {
+                        "packageList": {
+                            "type": "array",
+                            "items": {
+                                "type": "string"
+                            },
+                            "description": "The list of packages to install."
+                        },
+                        "workspacePath": {
+                            "type": "string",
+                            "description": "Path to Python workspace that determines the environment for package installation."
+                        }
+                    },
+                    "required": [
+                        "packageList",
+                        "workspacePath"
                     ]
                 }
             }
diff --git a/src/extension.ts b/src/extension.ts
@@ -53,7 +53,7 @@ import { ensureCorrectVersion } from './common/extVersion';
 import { ExistingProjects } from './features/creators/existingProjects';
 import { AutoFindProjects } from './features/creators/autoFindProjects';
 import { registerTools } from './common/lm.apis';
-import { GetPackagesTool } from './features/copilotTools';
+import { GetEnvironmentInfoTool, InstallPackageTool } from './features/copilotTools';
 import { TerminalActivationImpl } from './features/terminal/terminalActivationState';
 import { getEnvironmentForTerminal } from './features/terminal/utils';
 
@@ -107,7 +107,8 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
 
     context.subscriptions.push(
         registerCompletionProvider(envManagers),
-        registerTools('python_get_packages', new GetPackagesTool(api)),
+        registerTools('python_environment_tool', new GetEnvironmentInfoTool(api, envManagers)),
+        registerTools('python_install_package_tool', new InstallPackageTool(api)),
         commands.registerCommand('python-envs.viewLogs', () => outputChannel.show()),
         commands.registerCommand('python-envs.refreshManager', async (item) => {
             await refreshManagerCommand(item);
diff --git a/src/features/copilotTools.ts b/src/features/copilotTools.ts
@@ -8,26 +8,44 @@ import {
     PreparedToolInvocation,
     Uri,
 } from 'vscode';
-import { PythonPackageGetterApi, PythonProjectEnvironmentApi } from '../api';
+import {
+    PackageManagementOptions,
+    PythonEnvironment,
+    PythonEnvironmentExecutionInfo,
+    PythonPackageGetterApi,
+    PythonPackageManagementApi,
+    PythonProjectEnvironmentApi,
+} from '../api';
 import { createDeferred } from '../common/utils/deferred';
+import { EnvironmentManagers } from '../internal.api';
+
+export interface IResourceReference {
+    resourcePath?: string;
+}
 
-export interface IGetActiveFile {
-    filePath?: string;
+interface EnvironmentInfo {
+    type: string; // e.g. conda, venv, virtualenv, sys
+    version: string;
+    runCommand: string;
+    packages: string[] | string; //include versions too
 }
 
 /**
- * A tool to get the list of installed Python packages in the active environment.
+ * A tool to get the information about the Python environment.
  */
-export class GetPackagesTool implements LanguageModelTool<IGetActiveFile> {
-    constructor(private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi) {}
+export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
+    constructor(
+        private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi,
+        private readonly envManagers: EnvironmentManagers,
+    ) {}
     /**
-     * Invokes the tool to get the list of installed packages.
+     * 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 list of installed packages or an error message.
+     * @returns The result containing the information about the Python environment or an error message.
      */
     async invoke(
-        options: LanguageModelToolInvocationOptions<IGetActiveFile>,
+        options: LanguageModelToolInvocationOptions<IResourceReference>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
         const deferredReturn = createDeferred<LanguageModelToolResult>();
@@ -36,55 +54,192 @@ export class GetPackagesTool implements LanguageModelTool<IGetActiveFile> {
             deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
         });
 
-        const parameters: IGetActiveFile = options.input;
+        const parameters: IResourceReference = options.input;
 
-        if (parameters.filePath === undefined || parameters.filePath === '') {
-            throw new Error('Invalid input: filePath is required');
+        if (parameters.resourcePath === undefined || parameters.resourcePath === '') {
+            throw new Error('Invalid input: resourcePath is required');
         }
-        const fileUri = Uri.file(parameters.filePath);
+        const resourcePath: Uri = Uri.parse(parameters.resourcePath);
+
+        // environment info set to default values
+        const envInfo: EnvironmentInfo = {
+            type: 'no type found',
+            version: 'no version found',
+            packages: 'no packages found',
+            runCommand: 'no run command found',
+        };
 
         try {
-            const environment = await this.api.getEnvironment(fileUri);
+            // environment
+            const environment: PythonEnvironment | undefined = await this.api.getEnvironment(resourcePath);
             if (!environment) {
-                // Check if the file is a notebook or a notebook cell to throw specific error messages.
-                if (fileUri.fsPath.endsWith('.ipynb') || fileUri.fsPath.includes('.ipynb#')) {
-                    throw new Error('Unable to access Jupyter kernels for notebook cells');
-                }
-                throw new Error('No environment found');
+                throw new Error('No environment found for the provided resource path: ' + resourcePath.fsPath);
             }
+
+            const execInfo: PythonEnvironmentExecutionInfo = environment.execInfo;
+            const executable = execInfo?.activatedRun?.executable ?? execInfo?.run.executable ?? 'python';
+            const args = execInfo?.activatedRun?.args ?? execInfo?.run.args ?? [];
+            envInfo.runCommand = args.length > 0 ? `${executable} ${args.join(' ')}` : executable;
+            envInfo.version = environment.version;
+
+            // get the environment type or manager if type is not available
+            try {
+                const managerId = environment.envId.managerId;
+                const manager = this.envManagers.getEnvironmentManager(managerId);
+                envInfo.type = manager?.name || 'cannot be determined';
+            } catch {
+                envInfo.type = environment.envId.managerId || 'cannot be determined';
+            }
+
+            // TODO: remove refreshPackages here eventually once terminal isn't being used as a fallback
             await this.api.refreshPackages(environment);
             const installedPackages = await this.api.getPackages(environment);
-
-            let resultMessage: string;
             if (!installedPackages || installedPackages.length === 0) {
-                resultMessage = 'No packages are installed in the current environment.';
+                envInfo.packages = [];
             } else {
-                const packageNames = installedPackages
-                    .map((pkg) => (pkg.version ? `${pkg.name} (${pkg.version})` : pkg.name))
-                    .join(', ');
-                resultMessage = 'The packages installed in the current environment are as follows:\n' + packageNames;
+                envInfo.packages = installedPackages.map((pkg) =>
+                    pkg.version ? `${pkg.name} (${pkg.version})` : pkg.name,
+                );
             }
 
-            const textPart = new LanguageModelTextPart(resultMessage || '');
+            // format and return
+            const textPart = BuildEnvironmentInfoContent(envInfo);
             deferredReturn.resolve({ content: [textPart] });
         } catch (error) {
-            const errorMessage: string = `An error occurred while fetching packages: ${error}`;
+            const errorMessage: string = `An error occurred while fetching environment information: ${error}`;
+            const partialContent = BuildEnvironmentInfoContent(envInfo);
+            const combinedContent = new LanguageModelTextPart(`${errorMessage}\n\n${partialContent.value}`);
+            deferredReturn.resolve({ content: [combinedContent] } as LanguageModelToolResult);
+        }
+        return deferredReturn.promise;
+    }
+    /**
+     * Prepares the invocation of the tool.
+     * @param _options - The preparation options.
+     * @param _token - The cancellation token.
+     * @returns The prepared tool invocation.
+     */
+    async prepareInvocation?(
+        _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        _token: CancellationToken,
+    ): Promise<PreparedToolInvocation> {
+        const message = 'Preparing to fetch Python environment information...';
+        return {
+            invocationMessage: message,
+        };
+    }
+}
+
+function BuildEnvironmentInfoContent(envInfo: EnvironmentInfo): LanguageModelTextPart {
+    // Create a formatted string that looks like JSON but preserves comments
+    let envTypeDescriptor: string = `This environment is managed by ${envInfo.type} environment manager. Use the install tool to install packages into this environment.`;
+
+    if (envInfo.type === 'system') {
+        envTypeDescriptor =
+            'System pythons are pythons that ship with the OS or are installed globally. These python installs may be used by the OS for running services and core functionality. Confirm with the user before installing packages into this environment, as it can lead to issues with any services on the OS.';
+    }
+    const content = `{
+    // ${JSON.stringify(envTypeDescriptor)}
+  "environmentType": ${JSON.stringify(envInfo.type)},
+  // Python version of the environment
+  "pythonVersion": ${JSON.stringify(envInfo.version)},
+  // Use this command to run Python script or code in the terminal.
+  "runCommand": ${JSON.stringify(envInfo.runCommand)},
+  // 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.
+  "packages": ${JSON.stringify(Array.isArray(envInfo.packages) ? envInfo.packages : envInfo.packages, null, 2)}
+}`;
+
+    return new LanguageModelTextPart(content);
+}
+
+/**
+ * The input interface for the Install Package Tool.
+ */
+export interface IInstallPackageInput {
+    packageList: string[];
+    workspacePath?: string;
+}
+
+/**
+ * A tool to install Python packages in the active environment.
+ */
+export class InstallPackageTool implements LanguageModelTool<IInstallPackageInput> {
+    constructor(
+        private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi & PythonPackageManagementApi,
+    ) {}
+
+    /**
+     * Invokes the tool to install Python packages in the active environment.
+     * @param options - The invocation options containing the package list.
+     * @param token - The cancellation token.
+     * @returns The result containing the installation status or an error message.
+     */
+    async invoke(
+        options: LanguageModelToolInvocationOptions<IInstallPackageInput>,
+        token: CancellationToken,
+    ): Promise<LanguageModelToolResult> {
+        const deferredReturn = createDeferred<LanguageModelToolResult>();
+        token.onCancellationRequested(() => {
+            const errorMessage: string = `Operation cancelled by the user.`;
+            deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
+        });
+
+        const parameters: IInstallPackageInput = options.input;
+        const workspacePath = parameters.workspacePath ? Uri.file(parameters.workspacePath) : undefined;
+        if (!workspacePath) {
+            throw new Error('Invalid input: workspacePath is required');
+        }
+
+        if (!parameters.packageList || parameters.packageList.length === 0) {
+            throw new Error('Invalid input: packageList is required and cannot be empty');
+        }
+        const packageCount = parameters.packageList.length;
+        const packagePlurality = packageCount === 1 ? 'package' : 'packages';
+
+        try {
+            const environment = await this.api.getEnvironment(workspacePath);
+            if (!environment) {
+                // Check if the file is a notebook or a notebook cell to throw specific error messages.
+                if (workspacePath.fsPath.endsWith('.ipynb') || workspacePath.fsPath.includes('.ipynb#')) {
+                    throw new Error('Unable to access Jupyter kernels for notebook cells');
+                }
+                throw new Error('No environment found');
+            }
+
+            // Install the packages
+            const pkgManagementOptions: PackageManagementOptions = {
+                install: parameters.packageList,
+            };
+            await this.api.managePackages(environment, pkgManagementOptions);
+            const resultMessage = `Successfully installed ${packagePlurality}: ${parameters.packageList.join(', ')}`;
+
+            deferredReturn.resolve({
+                content: [new LanguageModelTextPart(resultMessage)],
+            });
+        } catch (error) {
+            const errorMessage = `An error occurred while installing ${packagePlurality}: ${error}`;
+
             deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
         }
+
         return deferredReturn.promise;
     }
 
     /**
      * Prepares the invocation of the tool.
-     * @param _options - The preparation options.
+     * @param options - The preparation options.
      * @param _token - The cancellation token.
      * @returns The prepared tool invocation.
      */
     async prepareInvocation?(
-        _options: LanguageModelToolInvocationPrepareOptions<IGetActiveFile>,
+        options: LanguageModelToolInvocationPrepareOptions<IInstallPackageInput>,
         _token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
-        const message = 'Preparing to fetch the list of installed Python packages...';
+        const packageList = options.input.packageList || [];
+        const packageCount = packageList.length;
+        const packageText = packageCount === 1 ? 'package' : 'packages';
+        const message = `Preparing to install Python ${packageText}: ${packageList.join(', ')}...`;
+
         return {
             invocationMessage: message,
         };
diff --git a/src/test/copilotTools.unit.test.ts b/src/test/copilotTools.unit.test.ts
