Set specific model for prompts and modes (microsoft#252893)

* add model to prompt and mode files

* add diagnostics

* handle model service not initialized

* switch model on mode update

* fix test

* fix value completion range

Author: aeschli

src/vs/workbench/contrib/chat/browser/chatInputPart.ts

@@ -510,6 +510,16 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		}
 	}
 
+	public switchModelByName(modelName: string): boolean {
+		const models = this.getModels();
+		const model = models.find(m => m.metadata.name === modelName);
+		if (model) {
+			this.setCurrentLanguageModel(model);
+			return true;
+		}
+		return false;
+	}
+
 	public switchToNextModel(): void {
 		const models = this.getModels();
 		if (models.length > 0) {
@@ -552,6 +562,11 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		this.chatModeKindKey.set(mode.kind);
 		this._onDidChangeCurrentChatMode.fire();
 
+		const model = mode.model?.get();
+		if (model) {
+			this.switchModelByName(model);
+		}
+
 		if (storeSelection) {
 			this.storageService.store(GlobalLastChatModeKey, mode.kind, StorageScope.APPLICATION, StorageTarget.USER);
 		}
@@ -560,10 +575,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	private modelSupportedForDefaultAgent(model: ILanguageModelChatMetadataAndIdentifier): boolean {
 		// Probably this logic could live in configuration on the agent, or somewhere else, if it gets more complex
 		if (this.currentModeKind === ChatModeKind.Agent || (this.currentModeKind === ChatModeKind.Edit && this.configurationService.getValue(ChatConfiguration.Edits2Enabled))) {
-			const supportsToolsAgent = typeof model.metadata.capabilities?.agentMode === 'undefined' || model.metadata.capabilities.agentMode;
-
-			// Filter out models that don't support tool calling, and models that don't support enough context to have a good experience with the tools agent
-			return supportsToolsAgent && !!model.metadata.capabilities?.toolCalling;
+			return ILanguageModelChatMetadata.suitableForAgentMode(model.metadata);
 		}
 
 		return true;

src/vs/workbench/contrib/chat/browser/chatSelectedTools.ts

@@ -176,7 +176,7 @@ export class ChatSelectedTools extends Disposable {
 			return;
 		}
 		if (mode.kind === ChatModeKind.Agent && mode.customTools && mode.uri) {
-			// apply directly to mode.
+			// apply directly to mode file.
 			this.updateCustomModeTools(mode.uri.get(), enablementMap);
 			return;
 		}

src/vs/workbench/contrib/chat/browser/chatWidget.ts

@@ -1879,7 +1879,7 @@ export class ChatWidget extends Disposable implements IChatWidget {
 
 	private async _applyPromptMetadata(metadata: TPromptMetadata, requestInput: IChatRequestInputOptions): Promise<void> {
 
-		const { mode, tools } = metadata;
+		const { mode, tools, model } = metadata;
 
 		// switch to appropriate chat mode if needed
 		if (mode && mode !== this.input.currentModeKind) {
@@ -1894,17 +1894,21 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		}
 
 		// if not tools to enable are present, we are done
-		if (tools === undefined) {
-			return;
-		}
+		if (tools !== undefined) {
+
 
-		// sanity check on the logic of the `getPromptFilesMetadata` method
-		// and the code above in case this block is moved around somewhere else:
-		// if we have some tools present, the mode must have been equal to `agent`
-		assert(this.input.currentModeKind === ChatModeKind.Agent, `Chat mode must be 'agent' when there are 'tools' defined, got ${this.input.currentModeKind}.`);
+			// sanity check on the logic of the `getPromptFilesMetadata` method
+			// and the code above in case this block is moved around somewhere else:
+			// if we have some tools present, the mode must have been equal to `agent`
+			assert(this.input.currentModeKind === ChatModeKind.Agent, `Chat mode must be 'agent' when there are 'tools' defined, got ${this.input.currentModeKind}.`);
 
-		const enablementMap = this.toolsService.toToolAndToolSetEnablementMap(new Set(tools));
-		this.input.selectedToolsModel.set(enablementMap, true);
+			const enablementMap = this.toolsService.toToolAndToolSetEnablementMap(new Set(tools));
+			this.input.selectedToolsModel.set(enablementMap, true);
+		}
+
+		if (model !== undefined) {
+			this.input.switchModelByName(model);
+		}
 	}
 
 	/**

src/vs/workbench/contrib/chat/common/chatModes.ts

@@ -95,6 +95,7 @@ export class ChatModeService extends Disposable implements IChatModeService {
 						name: cachedMode.name,
 						description: cachedMode.description,
 						tools: cachedMode.customTools,
+						model: cachedMode.model,
 						body: cachedMode.body || ''
 					};
 					const instance = new CustomChatMode(customChatMode);
@@ -190,6 +191,7 @@ export interface IChatModeData {
 	readonly description?: string;
 	readonly kind: ChatModeKind;
 	readonly customTools?: readonly string[];
+	readonly model?: string;
 	readonly body?: string;
 	readonly uri?: URI;
 }
@@ -200,8 +202,10 @@ export interface IChatMode {
 	readonly description: IObservable<string | undefined>;
 	readonly kind: ChatModeKind;
 	readonly customTools?: IObservable<readonly string[] | undefined>;
+	readonly model?: IObservable<string | undefined>;
 	readonly body?: IObservable<string>;
 	readonly uri?: IObservable<URI>;
+
 }
 
 function isCachedChatModeData(data: unknown): data is IChatModeData {
@@ -216,6 +220,7 @@ function isCachedChatModeData(data: unknown): data is IChatModeData {
 		(mode.description === undefined || typeof mode.description === 'string') &&
 		(mode.customTools === undefined || Array.isArray(mode.customTools)) &&
 		(mode.body === undefined || typeof mode.body === 'string') &&
+		(mode.model === undefined || typeof mode.model === 'string') &&
 		(mode.uri === undefined || (typeof mode.uri === 'object' && mode.uri !== null));
 }
 
@@ -224,6 +229,7 @@ export class CustomChatMode implements IChatMode {
 	private readonly _customToolsObservable: ISettableObservable<readonly string[] | undefined>;
 	private readonly _bodyObservable: ISettableObservable<string>;
 	private readonly _uriObservable: ISettableObservable<URI>;
+	private readonly _modelObservable: ISettableObservable<string | undefined>;
 
 	public readonly id: string;
 	public readonly name: string;
@@ -236,6 +242,10 @@ export class CustomChatMode implements IChatMode {
 		return this._customToolsObservable;
 	}
 
+	get model(): IObservable<string | undefined> {
+		return this._modelObservable;
+	}
+
 	get body(): IObservable<string> {
 		return this._bodyObservable;
 	}
@@ -253,6 +263,7 @@ export class CustomChatMode implements IChatMode {
 		this.name = customChatMode.name;
 		this._descriptionObservable = observableValue('description', customChatMode.description);
 		this._customToolsObservable = observableValue('customTools', customChatMode.tools);
+		this._modelObservable = observableValue('model', customChatMode.model);
 		this._bodyObservable = observableValue('body', customChatMode.body);
 		this._uriObservable = observableValue('uri', customChatMode.uri);
 	}
@@ -265,6 +276,7 @@ export class CustomChatMode implements IChatMode {
 			// Note- name is derived from ID, it can't change
 			this._descriptionObservable.set(newData.description, tx);
 			this._customToolsObservable.set(newData.tools, tx);
+			this._modelObservable.set(newData.model, tx);
 			this._bodyObservable.set(newData.body, tx);
 			this._uriObservable.set(newData.uri, tx);
 		});
@@ -277,6 +289,7 @@ export class CustomChatMode implements IChatMode {
 			description: this.description.get(),
 			kind: this.kind,
 			customTools: this.customTools.get(),
+			model: this.model.get(),
 			body: this.body.get(),
 			uri: this.uri.get()
 		};

src/vs/workbench/contrib/chat/common/languageModels.ts

@@ -146,6 +146,13 @@ export interface ILanguageModelChatMetadata {
 	};
 }
 
+export namespace ILanguageModelChatMetadata {
+	export function suitableForAgentMode(metadata: ILanguageModelChatMetadata): boolean {
+		const supportsToolsAgent = typeof metadata.capabilities?.agentMode === 'undefined' || metadata.capabilities.agentMode;
+		return supportsToolsAgent && !!metadata.capabilities?.toolCalling;
+	}
+}
+
 export interface ILanguageModelChatResponse {
 	stream: AsyncIterable<IChatResponseFragment | IChatResponseFragment[]>;
 	result: Promise<any>;

src/vs/workbench/contrib/chat/common/promptSyntax/languageProviders/promptHeaderAutocompletion.ts

@@ -13,6 +13,7 @@ import { CancellationToken } from '../../../../../../base/common/cancellation.js
 import { ILanguageFeaturesService } from '../../../../../../editor/common/services/languageFeatures.js';
 import { CompletionContext, CompletionItem, CompletionItemInsertTextRule, CompletionItemKind, CompletionItemProvider, CompletionList } from '../../../../../../editor/common/languages.js';
 import { Range } from '../../../../../../editor/common/core/range.js';
+import { ILanguageModelChatMetadata, ILanguageModelsService } from '../../languageModels.js';
 
 export class PromptHeaderAutocompletion extends Disposable implements CompletionItemProvider {
 	/**
@@ -28,6 +29,7 @@ export class PromptHeaderAutocompletion extends Disposable implements Completion
 	constructor(
 		@IPromptsService private readonly promptsService: IPromptsService,
 		@ILanguageFeaturesService private readonly languageService: ILanguageFeaturesService,
+		@ILanguageModelsService private readonly languageModelsService: ILanguageModelsService,
 
 	) {
 		super();
@@ -143,13 +145,14 @@ export class PromptHeaderAutocompletion extends Disposable implements Completion
 			return undefined;
 		}
 
+		const whilespaceAfterColon = (lineContent.substring(colonPosition.column).match(/^\s*/)?.[0].length) ?? 0;
 		const values = this.getValueSuggestions(promptType, property);
 		for (const value of values) {
 			const item: CompletionItem = {
 				label: value,
 				kind: CompletionItemKind.Value,
 				insertText: value,
-				range: new Range(position.lineNumber, position.column, position.lineNumber, model.getLineMaxColumn(position.lineNumber)),
+				range: new Range(position.lineNumber, colonPosition.column + whilespaceAfterColon + 1, position.lineNumber, model.getLineMaxColumn(position.lineNumber)),
 			};
 			suggestions.push(item);
 		}
@@ -161,9 +164,9 @@ export class PromptHeaderAutocompletion extends Disposable implements Completion
 			case PromptsType.instructions:
 				return new Set(['applyTo', 'description']);
 			case PromptsType.prompt:
-				return new Set(['mode', 'tools', 'description']);
+				return new Set(['mode', 'tools', 'description', 'model']);
 			default:
-				return new Set(['tools', 'description']);
+				return new Set(['tools', 'description', 'model']);
 		}
 	}
 
@@ -190,6 +193,22 @@ export class PromptHeaderAutocompletion extends Disposable implements Completion
 		if (property === 'tools' && (promptType === PromptsType.prompt || promptType === PromptsType.mode)) {
 			return ['[]', `['codebase', 'editFiles', 'fetch']`];
 		}
+		if (property === 'model' && (promptType === PromptsType.prompt || promptType === PromptsType.mode)) {
+			return this.getModelNames(promptType === PromptsType.mode);
+		}
 		return [];
 	}
+
+	private getModelNames(agentModeOnly: boolean): string[] {
+		const result = [];
+		for (const model of this.languageModelsService.getLanguageModelIds()) {
+			const metadata = this.languageModelsService.lookupLanguageModel(model);
+			if (metadata && metadata.isUserSelectable !== false) {
+				if (!agentModeOnly || ILanguageModelChatMetadata.suitableForAgentMode(metadata)) {
+					result.push(metadata.name);
+				}
+			}
+		}
+		return result;
+	}
 }

src/vs/workbench/contrib/chat/common/promptSyntax/languageProviders/promptHeaderDiagnosticsProvider.ts

@@ -11,6 +11,14 @@ import { CancellationToken } from '../../../../../../base/common/cancellation.js
 import { ProviderInstanceManagerBase, TProviderClass } from './providerInstanceManagerBase.js';
 import { TDiagnostic, PromptMetadataError, PromptMetadataWarning } from '../parsers/promptHeader/diagnostics.js';
 import { IMarkerData, IMarkerService, MarkerSeverity } from '../../../../../../platform/markers/common/markers.js';
+import { PromptHeader } from '../parsers/promptHeader/promptHeader.js';
+import { PromptToolsMetadata } from '../parsers/promptHeader/metadata/tools.js';
+import { PromptModelMetadata } from '../parsers/promptHeader/metadata/model.js';
+import { ModeHeader } from '../parsers/promptHeader/modeHeader.js';
+import { ILanguageModelChatMetadata, ILanguageModelsService } from '../../languageModels.js';
+import { ILanguageModelToolsService } from '../../languageModelToolsService.js';
+import { localize } from '../../../../../../nls.js';
+import { ChatModeKind } from '../../constants.js';
 
 /**
  * Unique ID of the markers provider class.
@@ -26,8 +34,16 @@ class PromptHeaderDiagnosticsProvider extends ProviderInstanceBase {
 		model: ITextModel,
 		@IPromptsService promptsService: IPromptsService,
 		@IMarkerService private readonly markerService: IMarkerService,
+		@ILanguageModelsService private readonly languageModelsService: ILanguageModelsService,
+		@ILanguageModelToolsService private readonly languageModelToolsService: ILanguageModelToolsService,
 	) {
 		super(model, promptsService);
+		this._register(languageModelsService.onDidChangeLanguageModels(() => {
+			this.onPromptSettled(undefined, CancellationToken.None);
+		}));
+		this._register(languageModelToolsService.onDidChangeTools(() => {
+			this.onPromptSettled(undefined, CancellationToken.None);
+		}));
 	}
 
 	/**
@@ -47,27 +63,95 @@ class PromptHeaderDiagnosticsProvider extends ProviderInstanceBase {
 
 		// header parsing process is separate from the prompt parsing one, hence
 		// apply markers only after the header is settled and so has diagnostics
-		header.settled.then(() => {
-			// by the time the promise finishes, the token might have been cancelled
-			// already due to a new 'onSettle' event, hence don't apply outdated markers
-			if (token.isCancellationRequested) {
-				return;
-			}
+		await header.settled;
+		// by the time the promise finishes, the token might have been cancelled
+		// already due to a new 'onSettle' event, hence don't apply outdated markers
+		if (token.isCancellationRequested) {
+			return;
+		}
 
-			const markers: IMarkerData[] = [];
-			for (const diagnostic of header.diagnostics) {
-				markers.push(toMarker(diagnostic));
-			}
+		const markers: IMarkerData[] = [];
+		for (const diagnostic of header.diagnostics) {
+			markers.push(toMarker(diagnostic));
+		}
+
+		if (header instanceof PromptHeader) {
+			this.validateTools(header.metadataUtility.tools, header.metadata.mode, markers);
+			this.validateModel(header.metadataUtility.model, header.metadata.mode, markers);
+		} else if (header instanceof ModeHeader) {
+			this.validateTools(header.metadataUtility.tools, ChatModeKind.Agent, markers);
+			this.validateModel(header.metadataUtility.model, ChatModeKind.Agent, markers);
 
-			this.markerService.changeOne(
-				MARKERS_OWNER_ID,
-				this.model.uri,
-				markers,
-			);
-		});
+		}
 
+		this.markerService.changeOne(
+			MARKERS_OWNER_ID,
+			this.model.uri,
+			markers,
+		);
 		return;
 	}
+	validateModel(modelNode: PromptModelMetadata | undefined, modeKind: ChatModeKind | undefined, markers: IMarkerData[]) {
+		if (!modelNode || modelNode.value === undefined) {
+			return;
+		}
+		const languageModes = this.languageModelsService.getLanguageModelIds();
+		if (languageModes.length === 0) {
+			// likely the service is not initialized yet
+			return;
+		}
+		const modelMetadata = this.findModelByName(languageModes, modelNode.value);
+		if (!modelMetadata) {
+			markers.push({
+				message: localize('promptHeaderDiagnosticsProvider.modelNotFound', "Unknown model '{0}'", modelNode.value),
+				severity: MarkerSeverity.Warning,
+				...modelNode.range,
+			});
+		} else if (modeKind === ChatModeKind.Agent && !ILanguageModelChatMetadata.suitableForAgentMode(modelMetadata)) {
+			markers.push({
+				message: localize('promptHeaderDiagnosticsProvider.modelNotSuited', "Model '{0}' is not suited for agent mode", modelNode.value),
+				severity: MarkerSeverity.Warning,
+				...modelNode.range,
+			});
+		}
+
+	}
+	findModelByName(languageModes: string[], modelName: string): ILanguageModelChatMetadata | undefined {
+		for (const model of languageModes) {
+			const metadata = this.languageModelsService.lookupLanguageModel(model);
+			if (metadata && metadata.isUserSelectable !== false && metadata.name === modelName) {
+				return metadata;
+			}
+		}
+		return undefined;
+	}
+
+	validateTools(tools: PromptToolsMetadata | undefined, modeKind: ChatModeKind | undefined, markers: IMarkerData[]) {
+		if (!tools || tools.value === undefined || modeKind === ChatModeKind.Ask || modeKind === ChatModeKind.Edit) {
+			return;
+		}
+		const toolNames = new Set(tools.value);
+		if (toolNames.size === 0) {
+			return;
+		}
+		for (const tool of this.languageModelToolsService.getTools()) {
+			toolNames.delete(tool.toolReferenceName ?? tool.displayName);
+		}
+		for (const toolSet of this.languageModelToolsService.toolSets.get()) {
+			toolNames.delete(toolSet.referenceName);
+		}
+
+		for (const toolName of toolNames) {
+			const range = tools.getToolRange(toolName);
+			if (range) {
+				markers.push({
+					message: localize('promptHeaderDiagnosticsProvider.toolNotFound', "Unknown tool '{0}'", toolName),
+					severity: MarkerSeverity.Warning,
+					...range,
+				});
+			}
+		}
+	}
 
 	/**
 	 * Returns a string representation of this object.