Skip to content

Commit 28f8a9c

Browse files
eleanorjboydeleanorjboyd
committed
iteration 1
1 parent d6fc846 commit 28f8a9c

File tree

4 files changed

+408
-142
lines changed

4 files changed

+408
-142
lines changed

‎package.json‎

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -494,22 +494,52 @@
494494
],
495495
"languageModelTools": [
496496
{
497-
"name": "python_get_packages",
498-
"displayName": "Get Python Packages",
499-
"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.",
500-
"toolReferenceName": "pythonGetPackages",
497+
"name": "python_environment_tool",
498+
"displayName": "Get Python Environment Information",
499+
"modelDescription": "Returns the information about the Python environment for the given file or workspace. Information includes environment type, python version, run command and installed packages with versions.",
500+
"toolReferenceName": "pythonGetEnvironmentInfo",
501501
"tags": [],
502502
"icon": "$(files)",
503+
"canBeReferencedInPrompt": true,
503504
"inputSchema": {
504505
"type": "object",
505506
"properties": {
506-
"filePath": {
507+
"resourcePath": {
507508
"type": "string"
508509
}
509510
},
510-
"description": "The path to the Python file or workspace to get the installed packages for.",
511+
"description": "The path to the Python file or workspace to get the environment information for.",
511512
"required": [
512-
"filePath"
513+
"resourcePath"
514+
]
515+
}
516+
},
517+
{
518+
"name": "python_install_package_tool",
519+
"displayName": "Install Python Package",
520+
"modelDescription": "Installs Python packages in a workspace. You should call this when you want to install packages in the user's environment.",
521+
"toolReferenceName": "pythonInstallPackage",
522+
"tags": [],
523+
"icon": "$(package)",
524+
"canBeReferencedInPrompt": true,
525+
"inputSchema": {
526+
"type": "object",
527+
"properties": {
528+
"packageList": {
529+
"type": "array",
530+
"items": {
531+
"type": "string"
532+
},
533+
"description": "The list of packages to install."
534+
},
535+
"workspacePath": {
536+
"type": "string",
537+
"description": "The path to the Python workspace to identify which environment to install packages in."
538+
}
539+
},
540+
"required": [
541+
"packageList",
542+
"workspacePath"
513543
]
514544
}
515545
}

‎src/extension.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ import { ensureCorrectVersion } from './common/extVersion';
5454
import { ExistingProjects } from './features/creators/existingProjects';
5555
import { AutoFindProjects } from './features/creators/autoFindProjects';
5656
import { registerTools } from './common/lm.apis';
57-
import { GetPackagesTool } from './features/copilotTools';
57+
import { GetEnvironmentInfoTool, InstallPackageTool } from './features/copilotTools';
5858
import { TerminalActivationImpl } from './features/terminal/terminalActivationState';
5959
import { getEnvironmentForTerminal } from './features/terminal/utils';
6060

@@ -108,7 +108,8 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
108108

109109
context.subscriptions.push(
110110
registerCompletionProvider(envManagers),
111-
registerTools('python_get_packages', new GetPackagesTool(api)),
111+
registerTools('python_environment_tool', new GetEnvironmentInfoTool(api)),
112+
registerTools('python_install_package_tool', new InstallPackageTool(api)),
112113
commands.registerCommand('python-envs.viewLogs', () => outputChannel.show()),
113114
commands.registerCommand('python-envs.refreshManager', async (item) => {
114115
await refreshManagerCommand(item);

‎src/features/copilotTools.ts‎

Lines changed: 167 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -8,26 +8,40 @@ import {
88
PreparedToolInvocation,
99
Uri,
1010
} from 'vscode';
11-
import { PythonPackageGetterApi, PythonProjectEnvironmentApi } from '../api';
11+
import {
12+
PythonCommandRunConfiguration,
13+
PythonEnvironment,
14+
PythonEnvironmentExecutionInfo,
15+
PythonPackageGetterApi,
16+
PythonPackageManagementApi,
17+
PythonProjectEnvironmentApi,
18+
} from '../api';
1219
import { createDeferred } from '../common/utils/deferred';
1320

14-
export interface IGetActiveFile {
15-
filePath?: string;
21+
export interface IResourceReference {
22+
resourcePath?: string;
23+
}
24+
25+
interface EnvironmentInfo {
26+
type: string; // e.g. conda, venv, virtualenv, sys
27+
version: string;
28+
runCommand: string;
29+
packages: string[] | string; //include versions too
1630
}
1731

1832
/**
19-
* A tool to get the list of installed Python packages in the active environment.
33+
* A tool to get the information about the Python environment.
2034
*/
21-
export class GetPackagesTool implements LanguageModelTool<IGetActiveFile> {
35+
export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
2236
constructor(private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi) {}
2337
/**
24-
* Invokes the tool to get the list of installed packages.
38+
* Invokes the tool to get the information about the Python environment.
2539
* @param options - The invocation options containing the file path.
2640
* @param token - The cancellation token.
27-
* @returns The result containing the list of installed packages or an error message.
41+
* @returns The result containing the information about the Python environment or an error message.
2842
*/
2943
async invoke(
30-
options: LanguageModelToolInvocationOptions<IGetActiveFile>,
44+
options: LanguageModelToolInvocationOptions<IResourceReference>,
3145
token: CancellationToken,
3246
): Promise<LanguageModelToolResult> {
3347
const deferredReturn = createDeferred<LanguageModelToolResult>();
@@ -36,55 +50,181 @@ export class GetPackagesTool implements LanguageModelTool<IGetActiveFile> {
3650
deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
3751
});
3852

39-
const parameters: IGetActiveFile = options.input;
53+
const parameters: IResourceReference = options.input;
4054

41-
if (parameters.filePath === undefined || parameters.filePath === '') {
42-
throw new Error('Invalid input: filePath is required');
55+
if (parameters.resourcePath === undefined || parameters.resourcePath === '') {
56+
throw new Error('Invalid input: resourcePath is required');
4357
}
44-
const fileUri = Uri.file(parameters.filePath);
58+
const resourcePath: Uri = Uri.file(parameters.resourcePath);
4559

4660
try {
47-
const environment = await this.api.getEnvironment(fileUri);
61+
// environment info set to default values
62+
const envInfo: EnvironmentInfo = {
63+
type: 'no type found',
64+
version: 'no version found',
65+
packages: 'no packages found',
66+
runCommand: 'no run command found',
67+
};
68+
69+
// environment
70+
const environment: PythonEnvironment | undefined = await this.api.getEnvironment(resourcePath);
4871
if (!environment) {
4972
// Check if the file is a notebook or a notebook cell to throw specific error messages.
50-
if (fileUri.fsPath.endsWith('.ipynb') || fileUri.fsPath.includes('.ipynb#')) {
73+
if (resourcePath.fsPath.endsWith('.ipynb') || resourcePath.fsPath.includes('.ipynb#')) {
5174
throw new Error('Unable to access Jupyter kernels for notebook cells');
5275
}
53-
throw new Error('No environment found');
76+
throw new Error('No environment found for the provided resource path: ' + resourcePath.fsPath);
5477
}
78+
79+
const execInfo: PythonEnvironmentExecutionInfo = environment.execInfo;
80+
const run: PythonCommandRunConfiguration = execInfo.run;
81+
envInfo.runCommand = run.executable + (run.args ? ` ${run.args.join(' ')}` : '');
82+
// TODO: check if this is the right way to get type
83+
envInfo.type = environment.envId.managerId.split(':')[1];
84+
envInfo.version = environment.version;
85+
86+
// does this need to be refreshed prior to returning to get any new packages?
5587
await this.api.refreshPackages(environment);
5688
const installedPackages = await this.api.getPackages(environment);
57-
58-
let resultMessage: string;
5989
if (!installedPackages || installedPackages.length === 0) {
60-
resultMessage = 'No packages are installed in the current environment.';
90+
envInfo.packages = [];
6191
} else {
62-
const packageNames = installedPackages
63-
.map((pkg) => (pkg.version ? `${pkg.name} (${pkg.version})` : pkg.name))
64-
.join(', ');
65-
resultMessage = 'The packages installed in the current environment are as follows:\n' + packageNames;
92+
envInfo.packages = installedPackages.map((pkg) =>
93+
pkg.version ? `${pkg.name} (${pkg.version})` : pkg.name,
94+
);
6695
}
6796

68-
const textPart = new LanguageModelTextPart(resultMessage || '');
97+
// format and return
98+
const textPart = BuildEnvironmentInfoContent(envInfo);
6999
deferredReturn.resolve({ content: [textPart] });
70100
} catch (error) {
71-
const errorMessage: string = `An error occurred while fetching packages: ${error}`;
101+
const errorMessage: string = `An error occurred while fetching environment information: ${error}`;
72102
deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
73103
}
74104
return deferredReturn.promise;
75105
}
76-
77106
/**
78107
* Prepares the invocation of the tool.
79108
* @param _options - The preparation options.
80109
* @param _token - The cancellation token.
81110
* @returns The prepared tool invocation.
82111
*/
83112
async prepareInvocation?(
84-
_options: LanguageModelToolInvocationPrepareOptions<IGetActiveFile>,
113+
_options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
85114
_token: CancellationToken,
86115
): Promise<PreparedToolInvocation> {
87-
const message = 'Preparing to fetch the list of installed Python packages...';
116+
const message = 'Preparing to fetch Python environment information...';
117+
return {
118+
invocationMessage: message,
119+
};
120+
}
121+
}
122+
123+
function BuildEnvironmentInfoContent(envInfo: EnvironmentInfo): LanguageModelTextPart {
124+
// Create a formatted string that looks like JSON but preserves comments
125+
const content = `{
126+
// type of python environment; sys means it is the system python
127+
"environmentType": ${JSON.stringify(envInfo.type)},
128+
// python version of the environment
129+
"pythonVersion": ${JSON.stringify(envInfo.version)},
130+
// command to run python in this environment, will include command with active environment if applicable
131+
"runCommand": ${JSON.stringify(envInfo.runCommand)},
132+
// installed python packages and their versions if know in the format <name> (<version>), empty array is returned if no packages are installed.
133+
"packages": ${JSON.stringify(Array.isArray(envInfo.packages) ? envInfo.packages : envInfo.packages, null, 2)}
134+
}`;
135+
136+
return new LanguageModelTextPart(content);
137+
}
138+
139+
/**
140+
* The input interface for the Install Package Tool.
141+
*/
142+
export interface IInstallPackageInput {
143+
packageList: string[];
144+
workspacePath?: string;
145+
}
146+
147+
/**
148+
* A tool to install Python packages in the active environment.
149+
*/
150+
export class InstallPackageTool implements LanguageModelTool<IInstallPackageInput> {
151+
constructor(
152+
private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi & PythonPackageManagementApi,
153+
) {}
154+
155+
/**
156+
* Invokes the tool to install Python packages in the active environment.
157+
* @param options - The invocation options containing the package list.
158+
* @param token - The cancellation token.
159+
* @returns The result containing the installation status or an error message.
160+
*/
161+
async invoke(
162+
options: LanguageModelToolInvocationOptions<IInstallPackageInput>,
163+
token: CancellationToken,
164+
): Promise<LanguageModelToolResult> {
165+
const deferredReturn = createDeferred<LanguageModelToolResult>();
166+
token.onCancellationRequested(() => {
167+
const errorMessage: string = `Operation cancelled by the user.`;
168+
deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
169+
});
170+
171+
const parameters: IInstallPackageInput = options.input;
172+
const workspacePath = parameters.workspacePath ? Uri.file(parameters.workspacePath) : undefined;
173+
if (!workspacePath) {
174+
throw new Error('Invalid input: workspacePath is required');
175+
}
176+
177+
if (!parameters.packageList || parameters.packageList.length === 0) {
178+
throw new Error('Invalid input: packageList is required and cannot be empty');
179+
}
180+
const packageCount = parameters.packageList.length;
181+
const packagePlurality = packageCount === 1 ? 'package' : 'packages';
182+
183+
try {
184+
const environment = await this.api.getEnvironment(workspacePath);
185+
if (!environment) {
186+
// Check if the file is a notebook or a notebook cell to throw specific error messages.
187+
if (workspacePath.fsPath.endsWith('.ipynb') || workspacePath.fsPath.includes('.ipynb#')) {
188+
throw new Error('Unable to access Jupyter kernels for notebook cells');
189+
}
190+
throw new Error('No environment found');
191+
}
192+
193+
// Install the packages
194+
await this.api.installPackages(environment, parameters.packageList);
195+
const resultMessage = `Successfully installed ${packagePlurality}: ${parameters.packageList.join(', ')}`;
196+
197+
// Refresh packages after installation to update the package view
198+
//TODO: do I want the await?
199+
await this.api.refreshPackages(environment);
200+
201+
deferredReturn.resolve({
202+
content: [new LanguageModelTextPart(resultMessage)],
203+
});
204+
} catch (error) {
205+
const errorMessage = `An error occurred while installing ${packagePlurality}: ${error}`;
206+
207+
deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult);
208+
}
209+
210+
return deferredReturn.promise;
211+
}
212+
213+
/**
214+
* Prepares the invocation of the tool.
215+
* @param options - The preparation options.
216+
* @param _token - The cancellation token.
217+
* @returns The prepared tool invocation.
218+
*/
219+
async prepareInvocation?(
220+
options: LanguageModelToolInvocationPrepareOptions<IInstallPackageInput>,
221+
_token: CancellationToken,
222+
): Promise<PreparedToolInvocation> {
223+
const packageList = options.input.packageList || [];
224+
const packageCount = packageList.length;
225+
const packageText = packageCount === 1 ? 'package' : 'packages';
226+
const message = `Preparing to install Python ${packageText}: ${packageList.join(', ')}...`;
227+
88228
return {
89229
invocationMessage: message,
90230
};

0 commit comments

Comments
 (0)