API to get last used env in a LM tool

src/client/api.ts

@@ -45,8 +45,10 @@ export function buildApi(
         TensorboardExtensionIntegration,
         TensorboardExtensionIntegration,
     );
-    const jupyterIntegration = serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration);
     const jupyterPythonEnvApi = serviceContainer.get<JupyterPythonEnvironmentApi>(JupyterExtensionPythonEnvironments);
+    const environments = buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi);
+    const jupyterIntegration = serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration);
+    jupyterIntegration.registerEnvApi(environments);
     const tensorboardIntegration = serviceContainer.get<TensorboardExtensionIntegration>(
         TensorboardExtensionIntegration,
     );
@@ -155,7 +157,7 @@ export function buildApi(
             stop: (client: BaseLanguageClient): Promise<void> => client.stop(),
             getTelemetryReporter: () => getTelemetryReporter(),
         },
-        environments: buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi),
+        environments,
     };
 
     // In test environment return the DI Container.

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<IInstallPackageArg
             for (const packageName of options.input.packageList) {
                 await installer.installModule(packageName, resourcePath, token, undefined, { installAsProcess: true });
             }
-
+            trackEnvUsedByTool(resourcePath, environment);
             // format and return
             const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(', ')}`;
             return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);

src/client/chat/lastUsedEnvs.ts

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { EnvironmentPath, PythonExtension } from '../api/types';
+import { Uri } from 'vscode';
+
+const MAX_TRACKED_URIS = 100; // Maximum number of environments to track
+const MAX_TRACKED_AGE = 15 * 60 * 1000; // Maximum age of tracked environments in milliseconds (15 minutes)
+
+type LastUsedEnvEntry = { uri: Uri | undefined; env: EnvironmentPath; dateTime: number };
+const lastUsedEnvs: LastUsedEnvEntry[] = [];
+
+/**
+ * Track the use of an environment for a given resource (uri).
+ * Prunes items older than 60 minutes or if the list grows over 100.
+ */
+export function trackEnvUsedByTool(uri: Uri | undefined, env: EnvironmentPath) {
+    const now = Date.now();
+    // Remove any previous entry for this uri
+    for (let i = lastUsedEnvs.length - 1; i >= 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;
+    }
+}

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 <name> or <name> (<version>). 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 <name> or <name> (<version>). The version may be omitted if unknown: ',

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'}`,

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<string | undefined>): 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<JupyterExtensionApi> | undefined;
 
     private pylanceExtension: Extension<PylanceApi> | 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<string | undefined>) =>
                 this.registerJupyterPythonPathFunction(func),
+            getLastUsedEnvInLmTool: (uri) => {
+                if (!this.environmentApi) {
+                    return undefined;
+                }
+                return getLastEnvUsedByTool(uri, this.environmentApi);
+            },
         });
         return undefined;
     }

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>(IEnvironmentVariablesProvider)).thenReturn(
             instance(environmentVariablesProvider),
         );
+        when(serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration)).thenReturn(
+            instance(mock<JupyterExtensionIntegration>()),
+        );
         when(serviceContainer.get<IInterpreterService>(IInterpreterService)).thenReturn(instance(interpreterService));
         const onDidChangePythonEnvironment = new EventEmitter<Uri>();
         const jupyterApi: JupyterPythonEnvironmentApi = {