diff --git a/src/client/api.ts b/src/client/api.ts index 15fb4d688a89..908da4be7103 100644 --- a/src/client/api.ts +++ b/src/client/api.ts @@ -45,8 +45,10 @@ export function buildApi( TensorboardExtensionIntegration, TensorboardExtensionIntegration, ); - const jupyterIntegration = serviceContainer.get(JupyterExtensionIntegration); const jupyterPythonEnvApi = serviceContainer.get(JupyterExtensionPythonEnvironments); + const environments = buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi); + const jupyterIntegration = serviceContainer.get(JupyterExtensionIntegration); + jupyterIntegration.registerEnvApi(environments); const tensorboardIntegration = serviceContainer.get( TensorboardExtensionIntegration, ); @@ -155,7 +157,7 @@ export function buildApi( stop: (client: BaseLanguageClient): Promise => client.stop(), getTelemetryReporter: () => getTelemetryReporter(), }, - environments: buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi), + environments, }; // In test environment return the DI Container. diff --git a/src/client/chat/installPackagesTool.ts b/src/client/chat/installPackagesTool.ts index a430525e1018..177c63077cbe 100644 --- a/src/client/chat/installPackagesTool.ts +++ b/src/client/chat/installPackagesTool.ts @@ -19,6 +19,7 @@ import { resolveFilePath } from './utils'; import { IModuleInstaller } from '../common/installer/types'; import { ModuleInstallerType } from '../pythonEnvironments/info'; import { IDiscoveryAPI } from '../pythonEnvironments/base/locator'; +import { trackEnvUsedByTool } from './lastUsedEnvs'; export interface IInstallPackageArgs { resourcePath?: string; @@ -66,7 +67,7 @@ export class InstallPackagesTool implements LanguageModelTool= 0; i--) { + if (urisEqual(lastUsedEnvs[i].uri, uri)) { + lastUsedEnvs.splice(i, 1); + } + } + // Add the new entry + lastUsedEnvs.push({ uri, env, dateTime: now }); + // Prune + pruneLastUsedEnvs(); +} + +/** + * Get the last used environment for a given resource (uri), or undefined if not found or expired. + */ +export function getLastEnvUsedByTool( + uri: Uri | undefined, + api: PythonExtension['environments'], +): EnvironmentPath | undefined { + pruneLastUsedEnvs(); + // Find the most recent entry for this uri that is not expired + const item = lastUsedEnvs.find((item) => urisEqual(item.uri, uri)); + if (item) { + return item.env; + } + const envPath = api.getActiveEnvironmentPath(uri); + if (lastUsedEnvs.some((item) => item.env.id === envPath.id)) { + // If this env was already used, return it + return envPath; + } + return undefined; +} + +/** + * Compare two uris (or undefined) for equality. + */ +function urisEqual(a: Uri | undefined, b: Uri | undefined): boolean { + if (a === b) { + return true; + } + if (!a || !b) { + return false; + } + return a.toString() === b.toString(); +} + +/** + * Remove items older than 60 minutes or if the list grows over 100. + */ +function pruneLastUsedEnvs() { + const now = Date.now(); + // Remove items older than 60 minutes + for (let i = lastUsedEnvs.length - 1; i >= 0; i--) { + if (now - lastUsedEnvs[i].dateTime > MAX_TRACKED_AGE) { + lastUsedEnvs.splice(i, 1); + } + } + // If still over 100, remove oldest + if (lastUsedEnvs.length > MAX_TRACKED_URIS) { + lastUsedEnvs.sort((a, b) => b.dateTime - a.dateTime); + lastUsedEnvs.length = MAX_TRACKED_URIS; + } +} diff --git a/src/client/chat/listPackagesTool.ts b/src/client/chat/listPackagesTool.ts index 157bdcf34793..0e410593de6c 100644 --- a/src/client/chat/listPackagesTool.ts +++ b/src/client/chat/listPackagesTool.ts @@ -22,6 +22,7 @@ import { parsePipList } from './pipListUtils'; import { Conda } from '../pythonEnvironments/common/environmentManagers/conda'; import { traceError } from '../logging'; import { IDiscoveryAPI } from '../pythonEnvironments/base/locator'; +import { trackEnvUsedByTool } from './lastUsedEnvs'; export interface IResourceReference { resourcePath?: string; @@ -108,7 +109,7 @@ export async function getPythonPackagesResponse( if (!packages.length) { return 'No packages found'; } - + trackEnvUsedByTool(resourcePath, environment); // Installed Python packages, each in the format or (). 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 or (). The version may be omitted if unknown: ', diff --git a/src/client/chat/utils.ts b/src/client/chat/utils.ts index 6206e01ea655..00a3fbb8393c 100644 --- a/src/client/chat/utils.ts +++ b/src/client/chat/utils.ts @@ -7,6 +7,7 @@ import { PythonExtension, ResolvedEnvironment } from '../api/types'; import { ITerminalHelper, TerminalShellType } from '../common/terminal/types'; import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution'; import { Conda } from '../pythonEnvironments/common/environmentManagers/conda'; +import { trackEnvUsedByTool } from './lastUsedEnvs'; export function resolveFilePath(filepath?: string): Uri | undefined { if (!filepath) { @@ -70,7 +71,7 @@ export async function getEnvironmentDetails( getTerminalCommand(environment, resourcePath, terminalExecutionService, terminalHelper), token, ); - + trackEnvUsedByTool(resourcePath, environment); const message = [ `Following is the information about the Python environment:`, `1. Environment Type: ${environment.environment?.type || 'unknown'}`, diff --git a/src/client/jupyter/jupyterIntegration.ts b/src/client/jupyter/jupyterIntegration.ts index 1136502c1ef2..a80c93916f3f 100644 --- a/src/client/jupyter/jupyterIntegration.ts +++ b/src/client/jupyter/jupyterIntegration.ts @@ -22,8 +22,9 @@ import { import { PylanceApi } from '../activation/node/pylanceApi'; import { ExtensionContextKey } from '../common/application/contextKeys'; import { getDebugpyPath } from '../debugger/pythonDebugger'; -import type { Environment } from '../api/types'; +import type { Environment, EnvironmentPath, PythonExtension } from '../api/types'; import { DisposableBase } from '../common/utils/resourceLifecycle'; +import { getLastEnvUsedByTool } from '../chat/lastUsedEnvs'; type PythonApiForJupyterExtension = { /** @@ -63,6 +64,11 @@ type PythonApiForJupyterExtension = { * @param func : The function that Python should call when requesting the Python path. */ registerJupyterPythonPathFunction(func: (uri: Uri) => Promise): void; + + /** + * Returns the Environment that was last used in a Python tool. + */ + getLastUsedEnvInLmTool(uri: Uri): EnvironmentPath | undefined; }; type JupyterExtensionApi = { @@ -78,6 +84,7 @@ export class JupyterExtensionIntegration { private jupyterExtension: Extension | undefined; private pylanceExtension: Extension | undefined; + private environmentApi: PythonExtension['environments'] | undefined; constructor( @inject(IExtensions) private readonly extensions: IExtensions, @@ -90,6 +97,9 @@ export class JupyterExtensionIntegration { @inject(IContextKeyManager) private readonly contextManager: IContextKeyManager, @inject(IInterpreterService) private interpreterService: IInterpreterService, ) {} + public registerEnvApi(api: PythonExtension['environments']) { + this.environmentApi = api; + } public registerApi(jupyterExtensionApi: JupyterExtensionApi): JupyterExtensionApi | undefined { this.contextManager.setContext(ExtensionContextKey.IsJupyterInstalled, true); @@ -121,6 +131,12 @@ export class JupyterExtensionIntegration { getCondaVersion: () => this.condaService.getCondaVersion(), registerJupyterPythonPathFunction: (func: (uri: Uri) => Promise) => this.registerJupyterPythonPathFunction(func), + getLastUsedEnvInLmTool: (uri) => { + if (!this.environmentApi) { + return undefined; + } + return getLastEnvUsedByTool(uri, this.environmentApi); + }, }); return undefined; } diff --git a/src/test/api.functional.test.ts b/src/test/api.functional.test.ts index 1149dcb7da9d..03016956dbef 100644 --- a/src/test/api.functional.test.ts +++ b/src/test/api.functional.test.ts @@ -19,7 +19,11 @@ import { ServiceManager } from '../client/ioc/serviceManager'; import { IServiceContainer, IServiceManager } from '../client/ioc/types'; import { IDiscoveryAPI } from '../client/pythonEnvironments/base/locator'; import * as pythonDebugger from '../client/debugger/pythonDebugger'; -import { JupyterExtensionPythonEnvironments, JupyterPythonEnvironmentApi } from '../client/jupyter/jupyterIntegration'; +import { + JupyterExtensionIntegration, + JupyterExtensionPythonEnvironments, + JupyterPythonEnvironmentApi, +} from '../client/jupyter/jupyterIntegration'; import { EventEmitter, Uri } from 'vscode'; suite('Extension API', () => { @@ -50,6 +54,9 @@ suite('Extension API', () => { when(serviceContainer.get(IEnvironmentVariablesProvider)).thenReturn( instance(environmentVariablesProvider), ); + when(serviceContainer.get(JupyterExtensionIntegration)).thenReturn( + instance(mock()), + ); when(serviceContainer.get(IInterpreterService)).thenReturn(instance(interpreterService)); const onDidChangePythonEnvironment = new EventEmitter(); const jupyterApi: JupyterPythonEnvironmentApi = {