|
| 1 | +import { |
| 2 | + CancellationToken, |
| 3 | + LanguageModelTextPart, |
| 4 | + LanguageModelTool, |
| 5 | + LanguageModelToolInvocationOptions, |
| 6 | + LanguageModelToolInvocationPrepareOptions, |
| 7 | + LanguageModelToolResult, |
| 8 | + PreparedToolInvocation, |
| 9 | + Uri, |
| 10 | +} from 'vscode'; |
| 11 | +import { PythonPackageGetterApi, PythonProjectEnvironmentApi } from './api'; |
| 12 | +import { createDeferred } from './common/utils/deferred'; |
| 13 | + |
| 14 | +export interface IGetActiveFile { |
| 15 | + filePath?: string; |
| 16 | +} |
| 17 | + |
| 18 | +/** |
| 19 | + * A tool to get the list of installed Python packages in the active environment. |
| 20 | + */ |
| 21 | +export class GetPackagesTool implements LanguageModelTool<IGetActiveFile> { |
| 22 | + constructor(private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi) {} |
| 23 | + /** |
| 24 | + * Invokes the tool to get the list of installed packages. |
| 25 | + * @param options - The invocation options containing the file path. |
| 26 | + * @param token - The cancellation token. |
| 27 | + * @returns The result containing the list of installed packages or an error message. |
| 28 | + */ |
| 29 | + async invoke( |
| 30 | + options: LanguageModelToolInvocationOptions<IGetActiveFile>, |
| 31 | + token: CancellationToken, |
| 32 | + ): Promise<LanguageModelToolResult> { |
| 33 | + const deferredReturn = createDeferred<LanguageModelToolResult>(); |
| 34 | + token.onCancellationRequested(() => { |
| 35 | + const errorMessage: string = `Operation cancelled by the user.`; |
| 36 | + deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult); |
| 37 | + }); |
| 38 | + |
| 39 | + const parameters: IGetActiveFile = options.input; |
| 40 | + |
| 41 | + if (parameters.filePath === undefined || parameters.filePath === '') { |
| 42 | + throw new Error('Invalid input: filePath is required'); |
| 43 | + } |
| 44 | + const fileUri = Uri.file(parameters.filePath); |
| 45 | + |
| 46 | + try { |
| 47 | + const environment = await this.api.getEnvironment(fileUri); |
| 48 | + if (!environment) { |
| 49 | + // 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#')) { |
| 51 | + throw new Error('Unable to access Jupyter kernels for notebook cells'); |
| 52 | + } |
| 53 | + throw new Error('No environment found'); |
| 54 | + } |
| 55 | + await this.api.refreshPackages(environment); |
| 56 | + const installedPackages = await this.api.getPackages(environment); |
| 57 | + |
| 58 | + let resultMessage: string; |
| 59 | + if (!installedPackages || installedPackages.length === 0) { |
| 60 | + resultMessage = 'No packages are installed in the current environment.'; |
| 61 | + } 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; |
| 66 | + } |
| 67 | + |
| 68 | + const textPart = new LanguageModelTextPart(resultMessage || ''); |
| 69 | + deferredReturn.resolve({ content: [textPart] }); |
| 70 | + } catch (error) { |
| 71 | + const errorMessage: string = `An error occurred while fetching packages: ${error}`; |
| 72 | + deferredReturn.resolve({ content: [new LanguageModelTextPart(errorMessage)] } as LanguageModelToolResult); |
| 73 | + } |
| 74 | + return deferredReturn.promise; |
| 75 | + } |
| 76 | + |
| 77 | + /** |
| 78 | + * Prepares the invocation of the tool. |
| 79 | + * @param _options - The preparation options. |
| 80 | + * @param _token - The cancellation token. |
| 81 | + * @returns The prepared tool invocation. |
| 82 | + */ |
| 83 | + async prepareInvocation?( |
| 84 | + _options: LanguageModelToolInvocationPrepareOptions<IGetActiveFile>, |
| 85 | + _token: CancellationToken, |
| 86 | + ): Promise<PreparedToolInvocation> { |
| 87 | + const message = 'Preparing to fetch the list of installed Python packages...'; |
| 88 | + return { |
| 89 | + invocationMessage: message, |
| 90 | + }; |
| 91 | + } |
| 92 | +} |
0 commit comments