Skip to content

Commit b9c5c2f

Browse files
committed
API to get last used env in a LM tool
1 parent a653a6e commit b9c5c2f

File tree

6 files changed

+108
-6
lines changed

6 files changed

+108
-6
lines changed

src/client/api.ts

+4-2
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,10 @@ export function buildApi(
4545
TensorboardExtensionIntegration,
4646
TensorboardExtensionIntegration,
4747
);
48-
const jupyterIntegration = serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration);
4948
const jupyterPythonEnvApi = serviceContainer.get<JupyterPythonEnvironmentApi>(JupyterExtensionPythonEnvironments);
49+
const environments = buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi);
50+
const jupyterIntegration = serviceContainer.get<JupyterExtensionIntegration>(JupyterExtensionIntegration);
51+
jupyterIntegration.registerEnvApi(environments);
5052
const tensorboardIntegration = serviceContainer.get<TensorboardExtensionIntegration>(
5153
TensorboardExtensionIntegration,
5254
);
@@ -155,7 +157,7 @@ export function buildApi(
155157
stop: (client: BaseLanguageClient): Promise<void> => client.stop(),
156158
getTelemetryReporter: () => getTelemetryReporter(),
157159
},
158-
environments: buildEnvironmentApi(discoveryApi, serviceContainer, jupyterPythonEnvApi),
160+
environments,
159161
};
160162

161163
// In test environment return the DI Container.

src/client/chat/installPackagesTool.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import { resolveFilePath } from './utils';
1919
import { IModuleInstaller } from '../common/installer/types';
2020
import { ModuleInstallerType } from '../pythonEnvironments/info';
2121
import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
22+
import { trackEnvUsedByTool } from './lastUsedEnvs';
2223

2324
export interface IInstallPackageArgs {
2425
resourcePath?: string;
@@ -66,7 +67,7 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
6667
for (const packageName of options.input.packageList) {
6768
await installer.installModule(packageName, resourcePath, token, undefined, { installAsProcess: true });
6869
}
69-
70+
trackEnvUsedByTool(resourcePath, environment);
7071
// format and return
7172
const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(', ')}`;
7273
return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);

src/client/chat/lastUsedEnvs.ts

+81
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import { EnvironmentPath, PythonExtension } from '../api/types';
5+
import { Uri } from 'vscode';
6+
7+
const MAX_TRACKED_URIS = 100; // Maximum number of environments to track
8+
const MAX_TRACKED_AGE = 15 * 60 * 1000; // Maximum age of tracked environments in milliseconds (15 minutes)
9+
10+
type LastUsedEnvEntry = { uri: Uri | undefined; env: EnvironmentPath; dateTime: number };
11+
const lastUsedEnvs: LastUsedEnvEntry[] = [];
12+
13+
/**
14+
* Track the use of an environment for a given resource (uri).
15+
* Prunes items older than 60 minutes or if the list grows over 100.
16+
*/
17+
export function trackEnvUsedByTool(uri: Uri | undefined, env: EnvironmentPath) {
18+
const now = Date.now();
19+
// Remove any previous entry for this uri
20+
for (let i = lastUsedEnvs.length - 1; i >= 0; i--) {
21+
if (urisEqual(lastUsedEnvs[i].uri, uri)) {
22+
lastUsedEnvs.splice(i, 1);
23+
}
24+
}
25+
// Add the new entry
26+
lastUsedEnvs.push({ uri, env, dateTime: now });
27+
// Prune
28+
pruneLastUsedEnvs();
29+
}
30+
31+
/**
32+
* Get the last used environment for a given resource (uri), or undefined if not found or expired.
33+
*/
34+
export function getLastEnvUsedByTool(
35+
uri: Uri | undefined,
36+
api: PythonExtension['environments'],
37+
): EnvironmentPath | undefined {
38+
pruneLastUsedEnvs();
39+
// Find the most recent entry for this uri that is not expired
40+
const item = lastUsedEnvs.find((item) => urisEqual(item.uri, uri));
41+
if (item) {
42+
return item.env;
43+
}
44+
const envPath = api.getActiveEnvironmentPath(uri);
45+
if (lastUsedEnvs.some((item) => item.env.id === envPath.id)) {
46+
// If this env was already used, return it
47+
return envPath;
48+
}
49+
return undefined;
50+
}
51+
52+
/**
53+
* Compare two uris (or undefined) for equality.
54+
*/
55+
function urisEqual(a: Uri | undefined, b: Uri | undefined): boolean {
56+
if (a === b) {
57+
return true;
58+
}
59+
if (!a || !b) {
60+
return false;
61+
}
62+
return a.toString() === b.toString();
63+
}
64+
65+
/**
66+
* Remove items older than 60 minutes or if the list grows over 100.
67+
*/
68+
function pruneLastUsedEnvs() {
69+
const now = Date.now();
70+
// Remove items older than 60 minutes
71+
for (let i = lastUsedEnvs.length - 1; i >= 0; i--) {
72+
if (now - lastUsedEnvs[i].dateTime > MAX_TRACKED_AGE) {
73+
lastUsedEnvs.splice(i, 1);
74+
}
75+
}
76+
// If still over 100, remove oldest
77+
if (lastUsedEnvs.length > MAX_TRACKED_URIS) {
78+
lastUsedEnvs.sort((a, b) => b.dateTime - a.dateTime);
79+
lastUsedEnvs.length = MAX_TRACKED_URIS;
80+
}
81+
}

src/client/chat/listPackagesTool.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import { parsePipList } from './pipListUtils';
2222
import { Conda } from '../pythonEnvironments/common/environmentManagers/conda';
2323
import { traceError } from '../logging';
2424
import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
25+
import { trackEnvUsedByTool } from './lastUsedEnvs';
2526

2627
export interface IResourceReference {
2728
resourcePath?: string;
@@ -108,7 +109,7 @@ export async function getPythonPackagesResponse(
108109
if (!packages.length) {
109110
return 'No packages found';
110111
}
111-
112+
trackEnvUsedByTool(resourcePath, environment);
112113
// 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.
113114
const response = [
114115
'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

+2-1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { PythonExtension, ResolvedEnvironment } from '../api/types';
77
import { ITerminalHelper, TerminalShellType } from '../common/terminal/types';
88
import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
99
import { Conda } from '../pythonEnvironments/common/environmentManagers/conda';
10+
import { trackEnvUsedByTool } from './lastUsedEnvs';
1011

1112
export function resolveFilePath(filepath?: string): Uri | undefined {
1213
if (!filepath) {
@@ -70,7 +71,7 @@ export async function getEnvironmentDetails(
7071
getTerminalCommand(environment, resourcePath, terminalExecutionService, terminalHelper),
7172
token,
7273
);
73-
74+
trackEnvUsedByTool(resourcePath, environment);
7475
const message = [
7576
`Following is the information about the Python environment:`,
7677
`1. Environment Type: ${environment.environment?.type || 'unknown'}`,

src/client/jupyter/jupyterIntegration.ts

+17-1
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,9 @@ import {
2222
import { PylanceApi } from '../activation/node/pylanceApi';
2323
import { ExtensionContextKey } from '../common/application/contextKeys';
2424
import { getDebugpyPath } from '../debugger/pythonDebugger';
25-
import type { Environment } from '../api/types';
25+
import type { Environment, EnvironmentPath, PythonExtension } from '../api/types';
2626
import { DisposableBase } from '../common/utils/resourceLifecycle';
27+
import { getLastEnvUsedByTool } from '../chat/lastUsedEnvs';
2728

2829
type PythonApiForJupyterExtension = {
2930
/**
@@ -63,6 +64,11 @@ type PythonApiForJupyterExtension = {
6364
* @param func : The function that Python should call when requesting the Python path.
6465
*/
6566
registerJupyterPythonPathFunction(func: (uri: Uri) => Promise<string | undefined>): void;
67+
68+
/**
69+
* Returns the Environment that was last used in a Python tool.
70+
*/
71+
getLastUsedEnvInLmTool(uri: Uri): EnvironmentPath | undefined;
6672
};
6773

6874
type JupyterExtensionApi = {
@@ -78,6 +84,7 @@ export class JupyterExtensionIntegration {
7884
private jupyterExtension: Extension<JupyterExtensionApi> | undefined;
7985

8086
private pylanceExtension: Extension<PylanceApi> | undefined;
87+
private environmentApi: PythonExtension['environments'] | undefined;
8188

8289
constructor(
8390
@inject(IExtensions) private readonly extensions: IExtensions,
@@ -90,6 +97,9 @@ export class JupyterExtensionIntegration {
9097
@inject(IContextKeyManager) private readonly contextManager: IContextKeyManager,
9198
@inject(IInterpreterService) private interpreterService: IInterpreterService,
9299
) {}
100+
public registerEnvApi(api: PythonExtension['environments']) {
101+
this.environmentApi = api;
102+
}
93103

94104
public registerApi(jupyterExtensionApi: JupyterExtensionApi): JupyterExtensionApi | undefined {
95105
this.contextManager.setContext(ExtensionContextKey.IsJupyterInstalled, true);
@@ -121,6 +131,12 @@ export class JupyterExtensionIntegration {
121131
getCondaVersion: () => this.condaService.getCondaVersion(),
122132
registerJupyterPythonPathFunction: (func: (uri: Uri) => Promise<string | undefined>) =>
123133
this.registerJupyterPythonPathFunction(func),
134+
getLastUsedEnvInLmTool: (uri) => {
135+
if (!this.environmentApi) {
136+
return undefined;
137+
}
138+
return getLastEnvUsedByTool(uri, this.environmentApi);
139+
},
124140
});
125141
return undefined;
126142
}

0 commit comments

Comments
 (0)