Capture invocation and failure reasons for tools

src/client/chat/baseTool.ts

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import {
+    CancellationToken,
+    LanguageModelTextPart,
+    LanguageModelTool,
+    LanguageModelToolInvocationOptions,
+    LanguageModelToolInvocationPrepareOptions,
+    LanguageModelToolResult,
+    PreparedToolInvocation,
+    Uri,
+    workspace,
+} from 'vscode';
+import { IResourceReference, isCancellationError, resolveFilePath } from './utils';
+import { ErrorWithTelemetrySafeReason } from '../common/errors/errorUtils';
+import { sendTelemetryEvent } from '../telemetry';
+import { EventName } from '../telemetry/constants';
+
+export abstract class BaseTool<T extends IResourceReference> implements LanguageModelTool<T> {
+    constructor(private readonly toolName: string) {}
+
+    async invoke(
+        options: LanguageModelToolInvocationOptions<T>,
+        token: CancellationToken,
+    ): Promise<LanguageModelToolResult> {
+        if (!workspace.isTrusted) {
+            return new LanguageModelToolResult([
+                new LanguageModelTextPart('Cannot use this tool in an untrusted workspace.'),
+            ]);
+        }
+        let error: Error | undefined;
+        const resource = resolveFilePath(options.input.resourcePath);
+        try {
+            return await this.invokeImpl(options, resource, token);
+        } catch (ex) {
+            error = ex as any;
+            throw ex;
+        } finally {
+            const isCancelled = token.isCancellationRequested || (error ? isCancellationError(error) : false);
+            const failed = !!error || isCancelled;
+            const failureCategory = isCancelled
+                ? 'cancelled'
+                : error
+                ? error instanceof ErrorWithTelemetrySafeReason
+                    ? error.telemetrySafeReason
+                    : 'error'
+                : undefined;
+            sendTelemetryEvent(EventName.INVOKE_TOOL, undefined, {
+                toolName: this.toolName,
+                failed,
+                failureCategory,
+            });
+        }
+    }
+    protected abstract invokeImpl(
+        options: LanguageModelToolInvocationOptions<T>,
+        resource: Uri | undefined,
+        token: CancellationToken,
+    ): Promise<LanguageModelToolResult>;
+
+    async prepareInvocation(
+        options: LanguageModelToolInvocationPrepareOptions<T>,
+        token: CancellationToken,
+    ): Promise<PreparedToolInvocation> {
+        const resource = resolveFilePath(options.input.resourcePath);
+        return this.prepareInvocationImpl(options, resource, token);
+    }
+
+    protected abstract prepareInvocationImpl(
+        options: LanguageModelToolInvocationPrepareOptions<T>,
+        resource: Uri | undefined,
+        token: CancellationToken,
+    ): Promise<PreparedToolInvocation>;
+}

src/client/chat/configurePythonEnvTool.ts

@@ -19,18 +19,18 @@ import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/termin
 import {
     getEnvDetailsForResponse,
     getToolResponseIfNotebook,
-    getUntrustedWorkspaceResponse,
     IResourceReference,
     isCancellationError,
     raceCancellationError,
 } from './utils';
-import { resolveFilePath } from './utils';
 import { ITerminalHelper } from '../common/terminal/types';
 import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
 import { CreateVirtualEnvTool } from './createVirtualEnvTool';
 import { ISelectPythonEnvToolArguments, SelectPythonEnvTool } from './selectEnvTool';
+import { BaseTool } from './baseTool';
 
-export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceReference> {
+export class ConfigurePythonEnvTool extends BaseTool<IResourceReference>
+    implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
     private readonly terminalHelper: ITerminalHelper;
     private readonly recommendedEnvService: IRecommendedEnvironmentService;
@@ -40,6 +40,7 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
         private readonly serviceContainer: IServiceContainer,
         private readonly createEnvTool: CreateVirtualEnvTool,
     ) {
+        super(ConfigurePythonEnvTool.toolName);
         this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
             ICodeExecutionService,
             'standard',
@@ -50,14 +51,11 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
         );
     }
 
-    async invoke(
+    async invokeImpl(
         options: LanguageModelToolInvocationOptions<IResourceReference>,
+        resource: Uri | undefined,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
-        if (!workspace.isTrusted) {
-            return getUntrustedWorkspaceResponse();
-        }
-        const resource = resolveFilePath(options.input.resourcePath);
         const notebookResponse = getToolResponseIfNotebook(resource);
         if (notebookResponse) {
             return notebookResponse;
@@ -101,8 +99,9 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
         }
     }
 
-    async prepareInvocation?(
+    async prepareInvocationImpl(
         _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        _resource: Uri | undefined,
         _token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
         return {

src/client/chat/createVirtualEnvTool.ts

@@ -22,12 +22,10 @@ import {
     doesWorkspaceHaveVenvOrCondaEnv,
     getDisplayVersion,
     getEnvDetailsForResponse,
-    getUntrustedWorkspaceResponse,
     IResourceReference,
     isCancellationError,
     raceCancellationError,
 } from './utils';
-import { resolveFilePath } from './utils';
 import { ITerminalHelper } from '../common/terminal/types';
 import { raceTimeout, sleep } from '../common/utils/async';
 import { IInterpreterPathService } from '../common/types';
@@ -44,12 +42,14 @@ import { StopWatch } from '../common/utils/stopWatch';
 import { useEnvExtension } from '../envExt/api.internal';
 import { PythonEnvironment } from '../envExt/types';
 import { hideEnvCreation } from '../pythonEnvironments/creation/provider/hideEnvCreation';
+import { BaseTool } from './baseTool';
 
 interface ICreateVirtualEnvToolParams extends IResourceReference {
     packageList?: string[]; // Added only becausewe have ability to create a virtual env with list of packages same tool within the in Python Env extension.
 }
 
-export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnvToolParams> {
+export class CreateVirtualEnvTool extends BaseTool<ICreateVirtualEnvToolParams>
+    implements LanguageModelTool<ICreateVirtualEnvToolParams> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
     private readonly terminalHelper: ITerminalHelper;
     private readonly recommendedEnvService: IRecommendedEnvironmentService;
@@ -60,6 +60,7 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
         private readonly api: PythonExtension['environments'],
         private readonly serviceContainer: IServiceContainer,
     ) {
+        super(CreateVirtualEnvTool.toolName);
         this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
             ICodeExecutionService,
             'standard',
@@ -70,14 +71,11 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
         );
     }
 
-    async invoke(
+    async invokeImpl(
         options: LanguageModelToolInvocationOptions<ICreateVirtualEnvToolParams>,
+        resource: Uri | undefined,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
-        if (!workspace.isTrusted) {
-            return getUntrustedWorkspaceResponse();
-        }
-        const resource = resolveFilePath(options.input.resourcePath);
         let info = await this.getPreferredEnvForCreation(resource);
         if (!info) {
             traceWarn(`Called ${CreateVirtualEnvTool.toolName} tool not invoked, no preferred environment found.`);
@@ -170,11 +168,11 @@ export class CreateVirtualEnvTool implements LanguageModelTool<ICreateVirtualEnv
         return info ? true : false;
     }
 
-    async prepareInvocation?(
-        options: LanguageModelToolInvocationPrepareOptions<ICreateVirtualEnvToolParams>,
+    async prepareInvocationImpl(
+        _options: LanguageModelToolInvocationPrepareOptions<ICreateVirtualEnvToolParams>,
+        resource: Uri | undefined,
         token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
-        const resource = resolveFilePath(options.input.resourcePath);
         const info = await raceCancellationError(this.getPreferredEnvForCreation(resource), token);
         if (!info) {
             return {};

src/client/chat/getExecutableTool.ts

@@ -10,7 +10,7 @@ import {
     LanguageModelToolInvocationPrepareOptions,
     LanguageModelToolResult,
     PreparedToolInvocation,
-    workspace,
+    Uri,
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
@@ -20,15 +20,14 @@ import {
     getEnvDisplayName,
     getEnvironmentDetails,
     getToolResponseIfNotebook,
-    getUntrustedWorkspaceResponse,
     IResourceReference,
     raceCancellationError,
 } from './utils';
-import { resolveFilePath } from './utils';
 import { ITerminalHelper } from '../common/terminal/types';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
+import { BaseTool } from './baseTool';
 
-export class GetExecutableTool implements LanguageModelTool<IResourceReference> {
+export class GetExecutableTool extends BaseTool<IResourceReference> implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
     private readonly terminalHelper: ITerminalHelper;
     public static readonly toolName = 'get_python_executable_details';
@@ -37,21 +36,18 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
         private readonly serviceContainer: IServiceContainer,
         private readonly discovery: IDiscoveryAPI,
     ) {
+        super(GetExecutableTool.toolName);
         this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
             ICodeExecutionService,
             'standard',
         );
         this.terminalHelper = this.serviceContainer.get<ITerminalHelper>(ITerminalHelper);
     }
-    async invoke(
-        options: LanguageModelToolInvocationOptions<IResourceReference>,
+    async invokeImpl(
+        _options: LanguageModelToolInvocationOptions<IResourceReference>,
+        resourcePath: Uri | undefined,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
-        if (!workspace.isTrusted) {
-            return getUntrustedWorkspaceResponse();
-        }
-
-        const resourcePath = resolveFilePath(options.input.resourcePath);
         const notebookResponse = getToolResponseIfNotebook(resourcePath);
         if (notebookResponse) {
             return notebookResponse;
@@ -68,11 +64,11 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
         return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
     }
 
-    async prepareInvocation?(
-        options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+    async prepareInvocationImpl(
+        _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        resourcePath: Uri | undefined,
         token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
-        const resourcePath = resolveFilePath(options.input.resourcePath);
         if (getToolResponseIfNotebook(resourcePath)) {
             return {};
         }

src/client/chat/getPythonEnvTool.ts

@@ -10,26 +10,22 @@ import {
     LanguageModelToolInvocationPrepareOptions,
     LanguageModelToolResult,
     PreparedToolInvocation,
-    workspace,
+    Uri,
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
 import { IProcessServiceFactory, IPythonExecutionFactory } from '../common/process/types';
-import {
-    getEnvironmentDetails,
-    getToolResponseIfNotebook,
-    getUntrustedWorkspaceResponse,
-    IResourceReference,
-    raceCancellationError,
-} from './utils';
-import { resolveFilePath } from './utils';
+import { getEnvironmentDetails, getToolResponseIfNotebook, IResourceReference, raceCancellationError } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
 import { getEnvExtApi, useEnvExtension } from '../envExt/api.internal';
+import { ErrorWithTelemetrySafeReason } from '../common/errors/errorUtils';
+import { BaseTool } from './baseTool';
 
-export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceReference> {
+export class GetEnvironmentInfoTool extends BaseTool<IResourceReference>
+    implements LanguageModelTool<IResourceReference> {
     private readonly terminalExecutionService: TerminalCodeExecutionProvider;
     private readonly pythonExecFactory: IPythonExecutionFactory;
     private readonly processServiceFactory: IProcessServiceFactory;
@@ -39,6 +35,7 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         private readonly api: PythonExtension['environments'],
         private readonly serviceContainer: IServiceContainer,
     ) {
+        super(GetEnvironmentInfoTool.toolName);
         this.terminalExecutionService = this.serviceContainer.get<TerminalCodeExecutionProvider>(
             ICodeExecutionService,
             'standard',
@@ -48,15 +45,11 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         this.terminalHelper = this.serviceContainer.get<ITerminalHelper>(ITerminalHelper);
     }
 
-    async invoke(
-        options: LanguageModelToolInvocationOptions<IResourceReference>,
+    async invokeImpl(
+        _options: LanguageModelToolInvocationOptions<IResourceReference>,
+        resourcePath: Uri | undefined,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
-        if (!workspace.isTrusted) {
-            return getUntrustedWorkspaceResponse();
-        }
-
-        const resourcePath = resolveFilePath(options.input.resourcePath);
         const notebookResponse = getToolResponseIfNotebook(resourcePath);
         if (notebookResponse) {
             return notebookResponse;
@@ -66,7 +59,10 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         const envPath = this.api.getActiveEnvironmentPath(resourcePath);
         const environment = await raceCancellationError(this.api.resolveEnvironment(envPath), token);
         if (!environment || !environment.version) {
-            throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
+            throw new ErrorWithTelemetrySafeReason(
+                'No environment found for the provided resource path: ' + resourcePath?.fsPath,
+                'noEnvFound',
+            );
         }
 
         let packages = '';
@@ -107,11 +103,11 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
         return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
     }
 
-    async prepareInvocation?(
-        options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+    async prepareInvocationImpl(
+        _options: LanguageModelToolInvocationPrepareOptions<IResourceReference>,
+        resourcePath: Uri | undefined,
         _token: CancellationToken,
     ): Promise<PreparedToolInvocation> {
-        const resourcePath = resolveFilePath(options.input.resourcePath);
         if (getToolResponseIfNotebook(resourcePath)) {
             return {};
         }