Skip to content

Use better type for ChatParticipantToolToken #248452

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/vs/workbench/api/common/extHostChatAgents2.ts
Original file line number Diff line number Diff line change
@@ -450,13 +450,15 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
const { request, location, history } = await this._createRequest(requestDto, context, detector.extension);

const model = await this.getModelForRequest(request, detector.extension);
const toolInvocationToken = this._tools.getToolInvocationToken(request.sessionId);
const extRequest = typeConvert.ChatAgentRequest.to(
request,
location,
model,
this.getDiagnosticsWhenEnabled(detector.extension),
this.getToolsForRequest(detector.extension, request),
detector.extension,
toolInvocationToken,
this._logService);

return detector.provider.provideParticipantDetection(
@@ -541,13 +543,15 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
stream = new ChatAgentResponseStream(agent.extension, request, this._proxy, this._commands.converter, sessionDisposables);

const model = await this.getModelForRequest(request, agent.extension);
const toolInvocationToken = this._tools.getToolInvocationToken(request.sessionId);
const extRequest = typeConvert.ChatAgentRequest.to(
request,
location,
model,
this.getDiagnosticsWhenEnabled(agent.extension),
this.getToolsForRequest(agent.extension, request),
agent.extension,
toolInvocationToken,
this._logService
);
inFlightRequest = { requestId: requestDto.requestId, extRequest };
@@ -653,6 +657,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
$releaseSession(sessionId: string): void {
this._sessionDisposables.deleteAndDispose(sessionId);
this._onDidDisposeChatSession.fire(sessionId);
this._tools.releaseSession(sessionId);
}

async $provideFollowups(requestDto: Dto<IChatAgentRequest>, handle: number, result: IChatAgentResult, context: { history: IChatAgentHistoryEntryDto[] }, token: CancellationToken): Promise<IChatFollowup[]> {
38 changes: 31 additions & 7 deletions src/vs/workbench/api/common/extHostLanguageModelTools.ts
Original file line number Diff line number Diff line change
@@ -11,15 +11,15 @@ import { IDisposable, toDisposable } from '../../../base/common/lifecycle.js';
import { revive } from '../../../base/common/marshalling.js';
import { generateUuid } from '../../../base/common/uuid.js';
import { IExtensionDescription } from '../../../platform/extensions/common/extensions.js';
import { IPreparedToolInvocation, isToolInvocationContext, IToolInvocation, IToolInvocationContext, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
import { IPreparedToolInvocation, IToolInvocation, IToolInvocationContext, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
import { ExtensionEditToolId, InternalEditToolId } from '../../contrib/chat/common/tools/editFileTool.js';
import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/tools.js';
import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
import { checkProposedApiEnabled, isProposedApiEnabled } from '../../services/extensions/common/extensions.js';
import { Dto, SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js';
import { ExtHostLanguageModelToolsShape, IMainContext, IToolDataDto, MainContext, MainThreadLanguageModelToolsShape } from './extHost.protocol.js';
import { ExtHostLanguageModels } from './extHostLanguageModels.js';
import * as typeConvert from './extHostTypeConverters.js';
import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';

class Tool {

@@ -61,6 +61,8 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
/** A map of all known tools, from other EHs or registered in vscode core */
private readonly _allTools = new Map<string, Tool>();

private readonly _sessionToolInvocationTokens = new Map<string, vscode.ChatParticipantToolToken>();

constructor(
mainContext: IMainContext,
private readonly _languageModels: ExtHostLanguageModels,
@@ -74,6 +76,30 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
});
}

releaseSession(sessionId: string): void {
this._sessionToolInvocationTokens.delete(sessionId);
}

getToolInvocationToken(sessionId: string): vscode.ChatParticipantToolToken {
if (!this._sessionToolInvocationTokens.has(sessionId)) {
const token: vscode.ChatParticipantToolToken = Object.freeze({
value: Symbol('ChatParticipantToolToken') as any
});
this._sessionToolInvocationTokens.set(sessionId, token);
}

return this._sessionToolInvocationTokens.get(sessionId)!;
}

private getContextForToken(token: vscode.ChatParticipantToolToken): IToolInvocationContext {
for (const [sessionId, sessionToken] of this._sessionToolInvocationTokens.entries()) {
if (token === sessionToken) {
return { sessionId };
}
}
throw new Error(`Invalid tool invocation token`);
}

async $countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise<number> {
const fn = this._tokenCountFuncs.get(callId);
if (!fn) {
@@ -90,9 +116,7 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
}

try {
if (options.toolInvocationToken && !isToolInvocationContext(options.toolInvocationToken)) {
throw new Error(`Invalid tool invocation token`);
}
const invocationContext = options.toolInvocationToken && this.getContextForToken(options.toolInvocationToken);

if ((toolId === InternalEditToolId || toolId === ExtensionEditToolId) && !isProposedApiEnabled(extension, 'chatParticipantPrivate')) {
throw new Error(`Invalid tool: ${toolId}`);
@@ -104,7 +128,7 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
callId,
parameters: options.input,
tokenBudget: options.tokenizationOptions?.tokenBudget,
context: options.toolInvocationToken as IToolInvocationContext | undefined,
context: invocationContext,
chatRequestId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.chatRequestId : undefined,
chatInteractionId: isProposedApiEnabled(extension, 'chatParticipantPrivate') ? options.chatInteractionId : undefined,
}, token);
@@ -159,7 +183,7 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape

const options: vscode.LanguageModelToolInvocationOptions<Object> = {
input: dto.parameters,
toolInvocationToken: dto.context as vscode.ChatParticipantToolToken | undefined,
toolInvocationToken: dto.context && this.getToolInvocationToken(dto.context.sessionId),
};
if (isProposedApiEnabled(item.extension, 'chatParticipantPrivate')) {
options.chatRequestId = dto.chatRequestId;
14 changes: 11 additions & 3 deletions src/vs/workbench/api/common/extHostTypeConverters.ts
Original file line number Diff line number Diff line change
@@ -2921,7 +2921,16 @@ export namespace ChatResponsePart {
}

export namespace ChatAgentRequest {
export function to(request: IChatAgentRequest, location2: vscode.ChatRequestEditorData | vscode.ChatRequestNotebookData | undefined, model: vscode.LanguageModelChat, diagnostics: readonly [vscode.Uri, readonly vscode.Diagnostic[]][], tools: Map<string, boolean>, extension: IRelaxedExtensionDescription, logService: ILogService): vscode.ChatRequest {
export function to(
request: IChatAgentRequest,
location2: vscode.ChatRequestEditorData | vscode.ChatRequestNotebookData | undefined,
model: vscode.LanguageModelChat,
diagnostics: readonly [vscode.Uri, readonly vscode.Diagnostic[]][],
tools: Map<string, boolean>,
extension: IRelaxedExtensionDescription,
toolInvocationToken: vscode.ChatParticipantToolToken,
logService: ILogService
): vscode.ChatRequest {
const toolReferences = request.variables.variables.filter(v => v.kind === 'tool');
const variableReferences = request.variables.variables.filter(v => v.kind !== 'tool');
const requestWithAllProps: vscode.ChatRequest = {
@@ -2939,7 +2948,7 @@ export namespace ChatAgentRequest {
acceptedConfirmationData: request.acceptedConfirmationData,
rejectedConfirmationData: request.rejectedConfirmationData,
location2,
toolInvocationToken: Object.freeze({ sessionId: request.sessionId }) as never,
toolInvocationToken,
tools,
model,
editedFileEvents: request.editedFileEvents,
@@ -2961,7 +2970,6 @@ export namespace ChatAgentRequest {
delete (requestWithAllProps as any).tools;
}


return requestWithAllProps;
}
}
Original file line number Diff line number Diff line change
@@ -119,10 +119,6 @@ export interface IToolInvocationContext {
sessionId: string;
}

export function isToolInvocationContext(obj: any): obj is IToolInvocationContext {
return typeof obj === 'object' && typeof obj.sessionId === 'string';
}

export interface IToolResultInputOutputDetails {
readonly input: string;
readonly output: ({ type: 'text'; value: string } | { type: 'data'; mimeType: string; value64: string })[];
7 changes: 6 additions & 1 deletion src/vscode-dts/vscode.d.ts
Original file line number Diff line number Diff line change
@@ -20621,7 +20621,12 @@ declare module 'vscode' {
/**
* A token that can be passed to {@link lm.invokeTool} when invoking a tool inside the context of handling a chat request.
*/
export type ChatParticipantToolToken = never;
export interface ChatParticipantToolToken {
/**
* A unique symbol to prevent accidental construction of this type (see note on {@link LanguageModelToolInvocationOptions.toolInvocationToken}).
*/
readonly value: unique symbol;
}

/**
* Options provided for tool invocation.
Loading
Oops, something went wrong.