Skip to content

Commit 180ba38

Browse files
DonJayamanneDonJayamanne
authored
Capture invocation and failure reasons for tools (#25220)
1 parent 01d872b commit 180ba38

File tree

11 files changed

+195
-98
lines changed

11 files changed

+195
-98
lines changed

‎src/client/chat/baseTool.ts

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import {
5+
CancellationToken,
6+
LanguageModelTextPart,
7+
LanguageModelTool,
8+
LanguageModelToolInvocationOptions,
9+
LanguageModelToolInvocationPrepareOptions,
10+
LanguageModelToolResult,
11+
PreparedToolInvocation,
12+
Uri,
13+
workspace,
14+
} from 'vscode';
15+
import { IResourceReference, isCancellationError, resolveFilePath } from './utils';
16+
import { ErrorWithTelemetrySafeReason } from '../common/errors/errorUtils';
17+
import { sendTelemetryEvent } from '../telemetry';
18+
import { EventName } from '../telemetry/constants';
19+
20+
export abstract class BaseTool<T extends IResourceReference> implements LanguageModelTool<T> {
21+
constructor(private readonly toolName: string) {}
22+
23+
async invoke(
24+
options: LanguageModelToolInvocationOptions<T>,
25+
token: CancellationToken,
26+
): Promise<LanguageModelToolResult> {
27+
if (!workspace.isTrusted) {
28+
return new LanguageModelToolResult([
29+
new LanguageModelTextPart('Cannot use this tool in an untrusted workspace.'),
30+
]);
31+
}
32+
let error: Error | undefined;
33+
const resource = resolveFilePath(options.input.resourcePath);
34+
try {
35+
return await this.invokeImpl(options, resource, token);
36+
} catch (ex) {
37+
error = ex as any;
38+
throw ex;
39+
} finally {
40+
const isCancelled = token.isCancellationRequested || (error ? isCancellationError(error) : false);
41+
const failed = !!error || isCancelled;
42+
const failureCategory = isCancelled
43+
? 'cancelled'
44+
: error
45+
? error instanceof ErrorWithTelemetrySafeReason
46+
? error.telemetrySafeReason
47+
: 'error'
48+
: undefined;
49+
sendTelemetryEvent(EventName.INVOKE_TOOL, undefined, {
50+
toolName: this.toolName,
51+
failed,
52+
failureCategory,
53+
});
54+
}
55+
}
56+
protected abstract invokeImpl(
57+
options: LanguageModelToolInvocationOptions<T>,
58+
resource: Uri | undefined,
59+
token: CancellationToken,
60+
): Promise<LanguageModelToolResult>;
61+
62+
async prepareInvocation(
63+
options: LanguageModelToolInvocationPrepareOptions<T>,
64+
token: CancellationToken,
65+
): Promise<PreparedToolInvocation> {
66+
const resource = resolveFilePath(options.input.resourcePath);
67+
return this.prepareInvocationImpl(options, resource, token);
68+
}
69+
70+
protected abstract prepareInvocationImpl(
71+
options: LanguageModelToolInvocationPrepareOptions<T>,
72+
resource: Uri | undefined,
73+
token: CancellationToken,
74+
): Promise<PreparedToolInvocation>;
75+
}

‎src/client/chat/configurePythonEnvTool.ts

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,18 @@ import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/termin
1919
import {
2020
getEnvDetailsForResponse,
2121
getToolResponseIfNotebook,
22-
getUntrustedWorkspaceResponse,
2322
IResourceReference,
2423
isCancellationError,
2524
raceCancellationError,
2625
} from './utils';
27-
import { resolveFilePath } from './utils';
2826
import { ITerminalHelper } from '../common/terminal/types';
2927
import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
3028
import { CreateVirtualEnvTool } from './createVirtualEnvTool';
3129
import { ISelectPythonEnvToolArguments, SelectPythonEnvTool } from './selectEnvTool';
30+
import { BaseTool } from './baseTool';
3231

33-
export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceReference> {
32+
export class ConfigurePythonEnvTool extends BaseTool<IResourceReference>
33+
implements LanguageModelTool<IResourceReference> {
3434
private readonly terminalExecutionService: TerminalCodeExecutionProvider;
3535
private readonly terminalHelper: ITerminalHelper;
3636
private readonly recommendedEnvService: IRecommendedEnvironmentService;
@@ -40,6 +40,7 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
4040
private readonly serviceContainer: IServiceContainer,
4141
private readonly createEnvTool: CreateVirtualEnvTool,
4242
) {
43+
super(ConfigurePythonEnvTool.toolName);
4344
this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
4445
ICodeExecutionService,
4546
'standard',
@@ -50,14 +51,11 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
5051
);
5152
}
5253

53-
async invoke(
54+
async invokeImpl(
5455
options: LanguageModelToolInvocationOptions<IResourceReference>,
56+
resource: Uri | undefined,
5557
token: CancellationToken,
5658
): Promise<LanguageModelToolResult> {
57-
if (!workspace.isTrusted) {
58-
return getUntrustedWorkspaceResponse();
59-
}
60-
const resource = resolveFilePath(options.input.resourcePath);
6159
const notebookResponse = getToolResponseIfNotebook(resource);
6260
if (notebookResponse) {
6361
return notebookResponse;
@@ -101,8 +99,9 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
10199
}
102100
}
103101

104-
async prepareInvocation?(
102+
async prepareInvocationImpl(
105103
_options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
104+
_resource: Uri | undefined,
106105
_token: CancellationToken,
107106
): Promise<PreparedToolInvocation> {
108107
return {

‎src/client/chat/createVirtualEnvTool.ts

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,10 @@ import {
2222
doesWorkspaceHaveVenvOrCondaEnv,
2323
getDisplayVersion,
2424
getEnvDetailsForResponse,
25-
getUntrustedWorkspaceResponse,
2625
IResourceReference,
2726
isCancellationError,
2827
raceCancellationError,
2928
} from './utils';
30-
import { resolveFilePath } from './utils';
3129
import { ITerminalHelper } from '../common/terminal/types';
3230
import { raceTimeout, sleep } from '../common/utils/async';
3331
import { IInterpreterPathService } from '../common/types';
@@ -44,12 +42,14 @@ import { StopWatch } from '../common/utils/stopWatch';
4442
import { useEnvExtension } from '../envExt/api.internal';
4543
import { PythonEnvironment } from '../envExt/types';
4644
import { hideEnvCreation } from '../pythonEnvironments/creation/provider/hideEnvCreation';
45+
import { BaseTool } from './baseTool';
4746

4847
interface ICreateVirtualEnvToolParams extends IResourceReference {
4948
packageList?: string[]; // Added only becausewe have ability to create a virtual env with list of packages same tool within the in Python Env extension.
5049
}
5150

52-
export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnvToolParams> {
51+
export class CreateVirtualEnvTool extends BaseTool<ICreateVirtualEnvToolParams>
52+
implements LanguageModelTool<ICreateVirtualEnvToolParams> {
5353
private readonly terminalExecutionService: TerminalCodeExecutionProvider;
5454
private readonly terminalHelper: ITerminalHelper;
5555
private readonly recommendedEnvService: IRecommendedEnvironmentService;
@@ -60,6 +60,7 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
6060
private readonly api: PythonExtension['environments'],
6161
private readonly serviceContainer: IServiceContainer,
6262
) {
63+
super(CreateVirtualEnvTool.toolName);
6364
this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
6465
ICodeExecutionService,
6566
'standard',
@@ -70,14 +71,11 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
7071
);
7172
}
7273

73-
async invoke(
74+
async invokeImpl(
7475
options: LanguageModelToolInvocationOptions<ICreateVirtualEnvToolParams>,
76+
resource: Uri | undefined,
7577
token: CancellationToken,
7678
): Promise<LanguageModelToolResult> {
77-
if (!workspace.isTrusted) {
78-
return getUntrustedWorkspaceResponse();
79-
}
80-
const resource = resolveFilePath(options.input.resourcePath);
8179
let info = await this.getPreferredEnvForCreation(resource);
8280
if (!info) {
8381
traceWarn(`Called ${CreateVirtualEnvTool.toolName} tool not invoked, no preferred environment found.`);
@@ -170,11 +168,11 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
170168
return info ? true : false;
171169
}
172170

173-
async prepareInvocation?(
174-
options: LanguageModelToolInvocationPrepareOptions<ICreateVirtualEnvToolParams>,
171+
async prepareInvocationImpl(
172+
_options: LanguageModelToolInvocationPrepareOptions<ICreateVirtualEnvToolParams>,
173+
resource: Uri | undefined,
175174
token: CancellationToken,
176175
): Promise<PreparedToolInvocation> {
177-
const resource = resolveFilePath(options.input.resourcePath);
178176
const info = await raceCancellationError(this.getPreferredEnvForCreation(resource), token);
179177
if (!info) {
180178
return {};

‎src/client/chat/getExecutableTool.ts

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {
1010
LanguageModelToolInvocationPrepareOptions,
1111
LanguageModelToolResult,
1212
PreparedToolInvocation,
13-
workspace,
13+
Uri,
1414
} from 'vscode';
1515
import { PythonExtension } from '../api/types';
1616
import { IServiceContainer } from '../ioc/types';
@@ -20,15 +20,14 @@ import {
2020
getEnvDisplayName,
2121
getEnvironmentDetails,
2222
getToolResponseIfNotebook,
23-
getUntrustedWorkspaceResponse,
2423
IResourceReference,
2524
raceCancellationError,
2625
} from './utils';
27-
import { resolveFilePath } from './utils';
2826
import { ITerminalHelper } from '../common/terminal/types';
2927
import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
28+
import { BaseTool } from './baseTool';
3029

31-
export class GetExecutableTool implements LanguageModelTool<IResourceReference> {
30+
export class GetExecutableTool extends BaseTool<IResourceReference> implements LanguageModelTool<IResourceReference> {
3231
private readonly terminalExecutionService: TerminalCodeExecutionProvider;
3332
private readonly terminalHelper: ITerminalHelper;
3433
public static readonly toolName = 'get_python_executable_details';
@@ -37,21 +36,18 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
3736
private readonly serviceContainer: IServiceContainer,
3837
private readonly discovery: IDiscoveryAPI,
3938
) {
39+
super(GetExecutableTool.toolName);
4040
this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
4141
ICodeExecutionService,
4242
'standard',
4343
);
4444
this.terminalHelper = this.serviceContainer.get<ITerminalHelper>(ITerminalHelper);
4545
}
46-
async invoke(
47-
options: LanguageModelToolInvocationOptions<IResourceReference>,
46+
async invokeImpl(
47+
_options: LanguageModelToolInvocationOptions<IResourceReference>,
48+
resourcePath: Uri | undefined,
4849
token: CancellationToken,
4950
): Promise<LanguageModelToolResult> {
50-
if (!workspace.isTrusted) {
51-
return getUntrustedWorkspaceResponse();
52-
}
53-
54-
const resourcePath = resolveFilePath(options.input.resourcePath);
5551
const notebookResponse = getToolResponseIfNotebook(resourcePath);
5652
if (notebookResponse) {
5753
return notebookResponse;
@@ -68,11 +64,11 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
6864
return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
6965
}
7066

71-
async prepareInvocation?(
72-
options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
67+
async prepareInvocationImpl(
68+
_options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
69+
resourcePath: Uri | undefined,
7370
token: CancellationToken,
7471
): Promise<PreparedToolInvocation> {
75-
const resourcePath = resolveFilePath(options.input.resourcePath);
7672
if (getToolResponseIfNotebook(resourcePath)) {
7773
return {};
7874
}

‎src/client/chat/getPythonEnvTool.ts

Lines changed: 17 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -10,26 +10,22 @@ import {
1010
LanguageModelToolInvocationPrepareOptions,
1111
LanguageModelToolResult,
1212
PreparedToolInvocation,
13-
workspace,
13+
Uri,
1414
} from 'vscode';
1515
import { PythonExtension } from '../api/types';
1616
import { IServiceContainer } from '../ioc/types';
1717
import { ICodeExecutionService } from '../terminals/types';
1818
import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
1919
import { IProcessServiceFactory, IPythonExecutionFactory } from '../common/process/types';
20-
import {
21-
getEnvironmentDetails,
22-
getToolResponseIfNotebook,
23-
getUntrustedWorkspaceResponse,
24-
IResourceReference,
25-
raceCancellationError,
26-
} from './utils';
27-
import { resolveFilePath } from './utils';
20+
import { getEnvironmentDetails, getToolResponseIfNotebook, IResourceReference, raceCancellationError } from './utils';
2821
import { getPythonPackagesResponse } from './listPackagesTool';
2922
import { ITerminalHelper } from '../common/terminal/types';
3023
import { getEnvExtApi, useEnvExtension } from '../envExt/api.internal';
24+
import { ErrorWithTelemetrySafeReason } from '../common/errors/errorUtils';
25+
import { BaseTool } from './baseTool';
3126

32-
export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
27+
export class GetEnvironmentInfoTool extends BaseTool<IResourceReference>
28+
implements LanguageModelTool<IResourceReference> {
3329
private readonly terminalExecutionService: TerminalCodeExecutionProvider;
3430
private readonly pythonExecFactory: IPythonExecutionFactory;
3531
private readonly processServiceFactory: IProcessServiceFactory;
@@ -39,6 +35,7 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
3935
private readonly api: PythonExtension['environments'],
4036
private readonly serviceContainer: IServiceContainer,
4137
) {
38+
super(GetEnvironmentInfoTool.toolName);
4239
this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
4340
ICodeExecutionService,
4441
'standard',
@@ -48,15 +45,11 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
4845
this.terminalHelper = this.serviceContainer.get<ITerminalHelper>(ITerminalHelper);
4946
}
5047

51-
async invoke(
52-
options: LanguageModelToolInvocationOptions<IResourceReference>,
48+
async invokeImpl(
49+
_options: LanguageModelToolInvocationOptions<IResourceReference>,
50+
resourcePath: Uri | undefined,
5351
token: CancellationToken,
5452
): Promise<LanguageModelToolResult> {
55-
if (!workspace.isTrusted) {
56-
return getUntrustedWorkspaceResponse();
57-
}
58-
59-
const resourcePath = resolveFilePath(options.input.resourcePath);
6053
const notebookResponse = getToolResponseIfNotebook(resourcePath);
6154
if (notebookResponse) {
6255
return notebookResponse;
@@ -66,7 +59,10 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
6659
const envPath = this.api.getActiveEnvironmentPath(resourcePath);
6760
const environment = await raceCancellationError(this.api.resolveEnvironment(envPath), token);
6861
if (!environment || !environment.version) {
69-
throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
62+
throw new ErrorWithTelemetrySafeReason(
63+
'No environment found for the provided resource path: ' + resourcePath?.fsPath,
64+
'noEnvFound',
65+
);
7066
}
7167

7268
let packages = '';
@@ -107,11 +103,11 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
107103
return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
108104
}
109105

110-
async prepareInvocation?(
111-
options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
106+
async prepareInvocationImpl(
107+
_options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
108+
resourcePath: Uri | undefined,
112109
_token: CancellationToken,
113110
): Promise<PreparedToolInvocation> {
114-
const resourcePath = resolveFilePath(options.input.resourcePath);
115111
if (getToolResponseIfNotebook(resourcePath)) {
116112
return {};
117113
}

0 commit comments

Comments
 (0)