diff --git a/package.json b/package.json index f552f1fc..df434d44 100644 --- 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 index 004a9832..7315370f 100644 --- 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 outputChannel.show()), commands.registerCommand('python-envs.refreshManager', async (item) => { await refreshManagerCommand(item); diff --git a/src/features/copilotTools.ts b/src/features/copilotTools.ts index f139fece..143858fa 100644 --- 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 { - constructor(private readonly api: PythonProjectEnvironmentApi & PythonPackageGetterApi) {} +export class GetEnvironmentInfoTool implements LanguageModelTool { + 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, + options: LanguageModelToolInvocationOptions, token: CancellationToken, ): Promise { const deferredReturn = createDeferred(); @@ -36,55 +54,192 @@ export class GetPackagesTool implements LanguageModelTool { 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, + _token: CancellationToken, + ): Promise { + 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 or (). 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 { + 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, + token: CancellationToken, + ): Promise { + const deferredReturn = createDeferred(); + 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, + options: LanguageModelToolInvocationPrepareOptions, _token: CancellationToken, ): Promise { - 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 index 2857fa56..03eb936d 100644 --- a/src/test/copilotTools.unit.test.ts +++ b/src/test/copilotTools.unit.test.ts @@ -2,210 +2,421 @@ import * as assert from 'assert'; import * as vscode from 'vscode'; import * as sinon from 'sinon'; import * as typeMoq from 'typemoq'; -import { Package, PythonEnvironment, PythonPackageGetterApi, PythonProjectEnvironmentApi } from '../api'; +import { + Package, + PackageId, + PythonEnvironment, + PythonEnvironmentId, + PythonPackageGetterApi, + PythonPackageManagementApi, + PythonProjectEnvironmentApi, +} from '../api'; import { createDeferred } from '../common/utils/deferred'; -import { GetPackagesTool, IGetActiveFile } from '../features/copilotTools'; - -suite('GetPackagesTool Tests', () => { - let tool: GetPackagesTool; - let mockApi: typeMoq.IMock; +import { + GetEnvironmentInfoTool, + IInstallPackageInput, + InstallPackageTool, + IResourceReference, +} from '../features/copilotTools'; +import { EnvironmentManagers, InternalEnvironmentManager } from '../internal.api'; + +suite('InstallPackageTool Tests', () => { + let installPackageTool: InstallPackageTool; + let mockApi: typeMoq.IMock; let mockEnvironment: typeMoq.IMock; setup(() => { // Create mock functions - mockApi = typeMoq.Mock.ofType(); + mockApi = typeMoq.Mock.ofType< + PythonProjectEnvironmentApi & PythonPackageGetterApi & PythonPackageManagementApi + >(); mockEnvironment = typeMoq.Mock.ofType(); // eslint-disable-next-line @typescript-eslint/no-explicit-any mockEnvironment.setup((x: any) => x.then).returns(() => undefined); - // refresh will always return a resolved promise - mockApi.setup((x) => x.refreshPackages(typeMoq.It.isAny())).returns(() => Promise.resolve()); - - // Create an instance of GetPackagesTool with the mock functions - tool = new GetPackagesTool(mockApi.object); + // Create an instance of InstallPackageTool with the mock functions + installPackageTool = new InstallPackageTool(mockApi.object); }); teardown(() => { sinon.restore(); }); - test('should throw error if filePath is undefined', async () => { - const testFile: IGetActiveFile = { - filePath: '', + test('should throw error if workspacePath is an empty string', async () => { + const testFile: IInstallPackageInput = { + workspacePath: '', + packageList: ['package1', 'package2'], }; const options = { input: testFile, toolInvocationToken: undefined }; const token = new vscode.CancellationTokenSource().token; - await assert.rejects(tool.invoke(options, token), { message: 'Invalid input: filePath is required' }); + await assert.rejects(installPackageTool.invoke(options, token), { + message: 'Invalid input: workspacePath is required', + }); }); test('should throw error for notebook files', async () => { // eslint-disable-next-line @typescript-eslint/no-explicit-any mockEnvironment.setup((x: any) => x.then).returns(() => undefined); - const testFile: IGetActiveFile = { - filePath: 'test.ipynb', + const testFile: IInstallPackageInput = { + workspacePath: 'this/is/a/test/path.ipynb', + packageList: ['package1', 'package2'], }; const options = { input: testFile, toolInvocationToken: undefined }; const token = new vscode.CancellationTokenSource().token; - const result = await tool.invoke(options, token); + const result = await installPackageTool.invoke(options, token); const content = result.content as vscode.LanguageModelTextPart[]; const firstPart = content[0] as vscode.LanguageModelTextPart; - assert.strictEqual( - firstPart.value, - 'An error occurred while fetching packages: Error: Unable to access Jupyter kernels for notebook cells', - ); + assert.strictEqual(firstPart.value.includes('An error occurred while installing packages'), true); }); test('should throw error for notebook cells', async () => { - const testFile: IGetActiveFile = { - filePath: 'test.ipynb#123', + const testFile: IInstallPackageInput = { + workspacePath: 'this/is/a/test/path.ipynb#cell', + packageList: ['package1', 'package2'], }; const options = { input: testFile, toolInvocationToken: undefined }; const token = new vscode.CancellationTokenSource().token; - const result = await tool.invoke(options, token); + const result = await installPackageTool.invoke(options, token); const content = result.content as vscode.LanguageModelTextPart[]; - const firstPart = content[0] as vscode.MarkdownString; + const firstPart = content[0] as vscode.LanguageModelTextPart; - assert.strictEqual( - firstPart.value, - 'An error occurred while fetching packages: Error: Unable to access Jupyter kernels for notebook cells', - ); + assert.strictEqual(firstPart.value.includes('An error occurred while installing packages'), true); }); - test('should return no packages message if no packages are installed', async () => { - const testFile: IGetActiveFile = { - filePath: 'test.py', + test('should throw error if packageList passed in is empty', async () => { + const testFile: IInstallPackageInput = { + workspacePath: 'path/to/workspace', + packageList: [], + }; + + const options = { input: testFile, toolInvocationToken: undefined }; + const token = new vscode.CancellationTokenSource().token; + await assert.rejects(installPackageTool.invoke(options, token), { + message: 'Invalid input: packageList is required and cannot be empty', + }); + }); + + test('should handle cancellation', async () => { + const testFile: IInstallPackageInput = { + workspacePath: 'path/to/workspace', + packageList: ['package1', 'package2'], }; mockApi .setup((x) => x.getEnvironment(typeMoq.It.isAny())) - .returns(() => { + .returns(async () => { return Promise.resolve(mockEnvironment.object); }); - mockApi.setup((x) => x.getPackages(typeMoq.It.isAny())).returns(() => Promise.resolve([])); - const options = { input: testFile, toolInvocationToken: undefined }; - const token = new vscode.CancellationTokenSource().token; - const result = await tool.invoke(options, token); - const content = result.content as vscode.LanguageModelTextPart[]; - const firstPart = content[0] as vscode.MarkdownString; + const tokenSource = new vscode.CancellationTokenSource(); + const token = tokenSource.token; - assert.strictEqual(firstPart.value, 'No packages are installed in the current environment.'); + const deferred = createDeferred(); + installPackageTool.invoke(options, token).then((result) => { + const content = result.content as vscode.LanguageModelTextPart[]; + const firstPart = content[0] as vscode.MarkdownString; + + assert.strictEqual(firstPart.value, 'Operation cancelled by the user.'); + deferred.resolve(); + }); + + tokenSource.cancel(); + await deferred.promise; }); - test('should return just packages if versions do not exist', async () => { - const testFile: IGetActiveFile = { - filePath: 'test.py', + test('should handle packages installation', async () => { + const testFile: IInstallPackageInput = { + workspacePath: 'path/to/workspace', + packageList: ['package1', 'package2'], }; mockApi .setup((x) => x.getEnvironment(typeMoq.It.isAny())) - .returns(() => { + .returns(async () => { return Promise.resolve(mockEnvironment.object); }); - const mockPackages: Package[] = [ - { - pkgId: { id: 'pkg1', managerId: 'pip', environmentId: 'env1' }, - name: 'package1', - displayName: 'package1', - }, - { - pkgId: { id: 'pkg2', managerId: 'pip', environmentId: 'env1' }, - name: 'package2', - displayName: 'package2', - }, - ]; - - mockApi.setup((x) => x.refreshPackages(typeMoq.It.isAny())).returns(() => Promise.resolve()); - mockApi.setup((x) => x.getPackages(typeMoq.It.isAny())).returns(() => Promise.resolve(mockPackages)); + mockApi + .setup((x) => x.managePackages(typeMoq.It.isAny(), typeMoq.It.isAny())) + .returns(() => { + const deferred = createDeferred(); + deferred.resolve(); + return deferred.promise; + }); const options = { input: testFile, toolInvocationToken: undefined }; const token = new vscode.CancellationTokenSource().token; - const result = await tool.invoke(options, token); + + const result = await installPackageTool.invoke(options, token); const content = result.content as vscode.LanguageModelTextPart[]; const firstPart = content[0] as vscode.MarkdownString; - assert.ok( - firstPart.value.includes('The packages installed in the current environment are as follows:') && - firstPart.value.includes('package1') && - firstPart.value.includes('package2'), - ); + assert.strictEqual(firstPart.value.includes('Successfully installed packages'), true); + assert.strictEqual(firstPart.value.includes('package1'), true); + assert.strictEqual(firstPart.value.includes('package2'), true); }); - - test('should return installed packages with versions', async () => { - const testFile: IGetActiveFile = { - filePath: 'test.py', + test('should handle package installation failure', async () => { + const testFile: IInstallPackageInput = { + workspacePath: 'path/to/workspace', + packageList: ['package1', 'package2'], }; mockApi .setup((x) => x.getEnvironment(typeMoq.It.isAny())) - .returns(() => { + .returns(async () => { return Promise.resolve(mockEnvironment.object); }); - const mockPackages: Package[] = [ - { - pkgId: { id: 'pkg1', managerId: 'pip', environmentId: 'env1' }, - name: 'package1', - displayName: 'package1', - version: '1.0.0', - }, - { - pkgId: { id: 'pkg2', managerId: 'pip', environmentId: 'env1' }, - name: 'package2', - displayName: 'package2', - version: '2.0.0', - }, - ]; - - mockApi.setup((x) => x.refreshPackages(typeMoq.It.isAny())).returns(() => Promise.resolve()); - mockApi.setup((x) => x.getPackages(typeMoq.It.isAny())).returns(() => Promise.resolve(mockPackages)); + mockApi + .setup((x) => x.managePackages(typeMoq.It.isAny(), typeMoq.It.isAny())) + .returns(() => { + const deferred = createDeferred(); + deferred.reject(new Error('Installation failed')); + return deferred.promise; + }); const options = { input: testFile, toolInvocationToken: undefined }; const token = new vscode.CancellationTokenSource().token; - const result = await tool.invoke(options, token); + + const result = await installPackageTool.invoke(options, token); const content = result.content as vscode.LanguageModelTextPart[]; const firstPart = content[0] as vscode.MarkdownString; - assert.ok( - firstPart.value.includes('The packages installed in the current environment are as follows:') && - firstPart.value.includes('package1 (1.0.0)') && - firstPart.value.includes('package2 (2.0.0)'), + assert.strictEqual( + firstPart.value.includes('An error occurred while installing packages'), + true, + `error message was ${firstPart.value}`, ); }); - - test('should handle cancellation', async () => { - const testFile: IGetActiveFile = { - filePath: 'test.py', + test('should handle error occurs when getting environment', async () => { + const testFile: IInstallPackageInput = { + workspacePath: 'path/to/workspace', + packageList: ['package1', 'package2'], }; + mockApi + .setup((x) => x.getEnvironment(typeMoq.It.isAny())) + .returns(async () => { + return Promise.reject(new Error('Unable to get environment')); + }); + const options = { input: testFile, toolInvocationToken: undefined }; + const token = new vscode.CancellationTokenSource().token; + const result = await installPackageTool.invoke(options, token); + const content = result.content as vscode.LanguageModelTextPart[]; + const firstPart = content[0] as vscode.MarkdownString; + assert.strictEqual(firstPart.value.includes('An error occurred while installing packages'), true); + }); + test('correct plurality in package installation message', async () => { + const testFile: IInstallPackageInput = { + workspacePath: 'path/to/workspace', + packageList: ['package1'], + }; mockApi .setup((x) => x.getEnvironment(typeMoq.It.isAny())) .returns(async () => { return Promise.resolve(mockEnvironment.object); }); + mockApi + .setup((x) => x.managePackages(typeMoq.It.isAny(), typeMoq.It.isAny())) + .returns(() => { + const deferred = createDeferred(); + deferred.resolve(); + return deferred.promise; + }); + const options = { input: testFile, toolInvocationToken: undefined }; + const token = new vscode.CancellationTokenSource().token; + const result = await installPackageTool.invoke(options, token); + const content = result.content as vscode.LanguageModelTextPart[]; + const firstPart = content[0] as vscode.MarkdownString; + assert.strictEqual(firstPart.value.includes('packages'), false); + assert.strictEqual(firstPart.value.includes('package'), true); + }); +}); +suite('GetEnvironmentInfoTool Tests', () => { + let getEnvironmentInfoTool: GetEnvironmentInfoTool; + let mockApi: typeMoq.IMock; + let mockEnvironment: typeMoq.IMock; + let em: typeMoq.IMock; + let managerSys: typeMoq.IMock; + + setup(() => { + // Create mock functions + mockApi = typeMoq.Mock.ofType< + PythonProjectEnvironmentApi & PythonPackageGetterApi & PythonPackageManagementApi + >(); + mockEnvironment = typeMoq.Mock.ofType(); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + mockEnvironment.setup((x: any) => x.then).returns(() => undefined); + + em = typeMoq.Mock.ofType(); + em.setup((e) => e.managers).returns(() => [managerSys.object]); + em.setup((e) => e.getEnvironmentManager(typeMoq.It.isAnyString())).returns(() => managerSys.object); + + getEnvironmentInfoTool = new GetEnvironmentInfoTool(mockApi.object, em.object); + }); + + teardown(() => { + sinon.restore(); + }); + test('should throw error if resourcePath is an empty string', async () => { + const testFile: IResourceReference = { + resourcePath: '', + }; + const options = { input: testFile, toolInvocationToken: undefined }; + const token = new vscode.CancellationTokenSource().token; + await assert.rejects(getEnvironmentInfoTool.invoke(options, token), { + message: 'Invalid input: resourcePath is required', + }); + }); + test('should throw error if environment is not found', async () => { + const testFile: IResourceReference = { + resourcePath: 'this/is/a/test/path.ipynb', + }; + mockApi + .setup((x) => x.getEnvironment(typeMoq.It.isAny())) + .returns(async () => { + return Promise.reject(new Error('Unable to get environment')); + }); + + const options = { input: testFile, toolInvocationToken: undefined }; + const token = new vscode.CancellationTokenSource().token; + const result = getEnvironmentInfoTool.invoke(options, token); + const content = (await result).content as vscode.LanguageModelTextPart[]; + const firstPart = content[0] as vscode.MarkdownString; + assert.strictEqual(firstPart.value.includes('An error occurred while fetching environment information'), true); + }); + test('should return successful with environment info', async () => { + // Create an instance of GetEnvironmentInfoTool with the mock functions + managerSys = typeMoq.Mock.ofType(); + managerSys.setup((m) => m.id).returns(() => 'ms-python.python:venv'); + managerSys.setup((m) => m.name).returns(() => 'venv'); + managerSys.setup((m) => m.displayName).returns(() => 'Test Manager'); + + em = typeMoq.Mock.ofType(); + em.setup((e) => e.managers).returns(() => [managerSys.object]); + em.setup((e) => e.getEnvironmentManager(typeMoq.It.isAnyString())).returns(() => managerSys.object); + // create mock of PythonEnvironment + const mockEnvironmentSuccess = typeMoq.Mock.ofType(); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + mockEnvironmentSuccess.setup((x: any) => x.then).returns(() => undefined); + mockEnvironmentSuccess.setup((x) => x.version).returns(() => '3.9.1'); + const mockEnvId = typeMoq.Mock.ofType(); + mockEnvId.setup((x) => x.managerId).returns(() => 'ms-python.python:venv'); + mockEnvironmentSuccess.setup((x) => x.envId).returns(() => mockEnvId.object); + mockEnvironmentSuccess + .setup((x) => x.execInfo) + .returns(() => ({ + run: { + executable: 'conda', + args: ['run', '-n', 'env_name', 'python'], + }, + })); + + mockApi + .setup((x) => x.getEnvironment(typeMoq.It.isAny())) + .returns(async () => { + return Promise.resolve(mockEnvironmentSuccess.object); + }); mockApi.setup((x) => x.refreshPackages(typeMoq.It.isAny())).returns(() => Promise.resolve()); - mockApi.setup((x) => x.getPackages(typeMoq.It.isAny())).returns(() => Promise.resolve([])); + const packageAId: PackageId = { + id: 'package1', + managerId: 'ms-python.python:venv', + environmentId: 'env_id', + }; + const packageBId: PackageId = { + id: 'package2', + managerId: 'ms-python.python:venv', + environmentId: 'env_id', + }; + const packageA: Package = { name: 'package1', displayName: 'Package 1', version: '1.0.0', pkgId: packageAId }; + const packageB: Package = { name: 'package2', displayName: 'Package 2', version: '2.0.0', pkgId: packageBId }; + mockApi + .setup((x) => x.getPackages(typeMoq.It.isAny())) + .returns(async () => { + return Promise.resolve([packageA, packageB]); + }); + + const testFile: IResourceReference = { + resourcePath: 'this/is/a/test/path.ipynb', + }; const options = { input: testFile, toolInvocationToken: undefined }; - const tokenSource = new vscode.CancellationTokenSource(); - const token = tokenSource.token; + const token = new vscode.CancellationTokenSource().token; + // run + const result = await getEnvironmentInfoTool.invoke(options, token); + // assert + const content = result.content as vscode.LanguageModelTextPart[]; + const firstPart = content[0] as vscode.MarkdownString; + assert.strictEqual(firstPart.value.includes('3.9.1'), true); + assert.strictEqual(firstPart.value.includes('package1 (1.0.0)'), true); + assert.strictEqual(firstPart.value.includes('package2 (2.0.0)'), true); + assert.strictEqual(firstPart.value.includes(`"conda run -n env_name python"`), true); + assert.strictEqual(firstPart.value.includes('venv'), true); + }); + test('should return successful with weird environment info', async () => { + // create mock of PythonEnvironment + const mockEnvironmentSuccess = typeMoq.Mock.ofType(); - const deferred = createDeferred(); - tool.invoke(options, token).then((result) => { - const content = result.content as vscode.LanguageModelTextPart[]; - const firstPart = content[0] as vscode.MarkdownString; + // Create an instance of GetEnvironmentInfoTool with the mock functions + let managerSys = typeMoq.Mock.ofType(); + managerSys.setup((m) => m.id).returns(() => 'ms-python.python:system'); + managerSys.setup((m) => m.name).returns(() => 'system'); + managerSys.setup((m) => m.displayName).returns(() => 'Test Manager'); - assert.strictEqual(firstPart.value, 'Operation cancelled by the user.'); - deferred.resolve(); - }); + let emSys = typeMoq.Mock.ofType(); + emSys.setup((e) => e.managers).returns(() => [managerSys.object]); + emSys.setup((e) => e.getEnvironmentManager(typeMoq.It.isAnyString())).returns(() => managerSys.object); + getEnvironmentInfoTool = new GetEnvironmentInfoTool(mockApi.object, emSys.object); - tokenSource.cancel(); - await deferred.promise; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + mockEnvironmentSuccess.setup((x: any) => x.then).returns(() => undefined); + mockEnvironmentSuccess.setup((x) => x.version).returns(() => '3.12.1'); + const mockEnvId = typeMoq.Mock.ofType(); + mockEnvId.setup((x) => x.managerId).returns(() => 'ms-python.python:system'); + managerSys.setup((m) => m.name).returns(() => 'system'); + mockEnvironmentSuccess.setup((x) => x.envId).returns(() => mockEnvId.object); + mockEnvironmentSuccess + .setup((x) => x.execInfo) + .returns(() => ({ + run: { + executable: 'path/to/venv/bin/python', + args: [], + }, + })); + + mockApi + .setup((x) => x.getEnvironment(typeMoq.It.isAny())) + .returns(async () => { + return Promise.resolve(mockEnvironmentSuccess.object); + }); + mockApi.setup((x) => x.refreshPackages(typeMoq.It.isAny())).returns(() => Promise.resolve()); + + mockApi + .setup((x) => x.getPackages(typeMoq.It.isAny())) + .returns(async () => { + return Promise.resolve([]); + }); + + const testFile: IResourceReference = { + resourcePath: 'this/is/a/test/path.ipynb', + }; + const options = { input: testFile, toolInvocationToken: undefined }; + const token = new vscode.CancellationTokenSource().token; + // run + const result = await getEnvironmentInfoTool.invoke(options, token); + // assert + const content = result.content as vscode.LanguageModelTextPart[]; + const firstPart = content[0] as vscode.MarkdownString; + assert.strictEqual(firstPart.value.includes('3.12.1'), true); + assert.strictEqual(firstPart.value.includes('"packages": []'), true); + assert.strictEqual(firstPart.value.includes(`"path/to/venv/bin/python"`), true); + assert.strictEqual(firstPart.value.includes('system'), true); }); });