Register categories as part of LM provider (#247146)

* Revert "Revert "Support categories in model picker (#247084)" (#247090)"

This reverts commit 29c0ca5.

* Switch model categories to be part of registration

Author: lramos15

src/vs/platform/actionWidget/browser/actionList.ts

@@ -75,7 +75,7 @@ class HeaderRenderer<T> implements IListRenderer<IActionListItem<T>, IHeaderTemp
 	}
 
 	renderElement(element: IActionListItem<T>, _index: number, templateData: IHeaderTemplateData): void {
-		templateData.text.textContent = element.group?.title ?? '';
+		templateData.text.textContent = element.group?.title ?? element.label ?? '';
 	}
 
 	disposeTemplate(_templateData: IHeaderTemplateData): void {

src/vs/workbench/api/common/extHostLanguageModels.ts

@@ -25,6 +25,7 @@ import { IExtHostRpcService } from './extHostRpcService.js';
 import * as typeConvert from './extHostTypeConverters.js';
 import * as extHostTypes from './extHostTypes.js';
 import { SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js';
+import { DEFAULT_MODEL_PICKER_CATEGORY } from '../../contrib/chat/common/modelPicker/modelPickerWidget.js';
 
 export interface IExtHostLanguageModels extends ExtHostLanguageModels { }
 
@@ -181,6 +182,7 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape {
 			targetExtensions: metadata.extensions,
 			isDefault: metadata.isDefault,
 			isUserSelectable: metadata.isUserSelectable,
+			modelPickerCategory: metadata.category ?? DEFAULT_MODEL_PICKER_CATEGORY,
 			capabilities: metadata.capabilities,
 		});
 

src/vs/workbench/contrib/chat/browser/modelPicker/modelPickerWidget.ts

@@ -101,16 +101,53 @@ export class ModelPickerWidget extends Disposable {
 	 * Shows the picker at the specified anchor
 	 */
 	showAt(anchor: HTMLElement | StandardMouseEvent | IAnchor, container?: HTMLElement): void {
-		const items: IActionListItem<ILanguageModelChatMetadataAndIdentifier>[] = this.getActionItems().map(item => ({
-			item: item.model,
-			description: item.model.metadata.description,
-			kind: ActionListItemKind.Action,
-			canPreview: false,
-			group: { title: '', icon: ThemeIcon.fromId(item.isCurrent ? Codicon.check.id : Codicon.blank.id) },
-			disabled: false,
-			hideIcon: false,
-			label: item.model.metadata.name,
-		} satisfies IActionListItem<ILanguageModelChatMetadataAndIdentifier>));
+		const actionItems = this.getActionItems();
+		const items: IActionListItem<ILanguageModelChatMetadataAndIdentifier>[] = [];
+
+		// Group models by categories
+		const modelsByCategory = new Map<string, IModelPickerActionItem[]>();
+
+		// First, group models by their categories
+		for (const item of actionItems) {
+			const category = item.model.metadata.modelPickerCategory;
+			if (!modelsByCategory.has(category.label)) {
+				modelsByCategory.set(category.label, []);
+			}
+			modelsByCategory.get(category.label)!.push(item);
+		}
+
+		for (const [categoryLabel, modelsInCategory] of modelsByCategory.entries()) {
+			// Skip empty categories
+			if (modelsInCategory.length === 0) {
+				continue;
+			}
+
+			// Add category header
+			items.push({
+				label: categoryLabel,
+				kind: ActionListItemKind.Header,
+				canPreview: false,
+				disabled: false,
+				hideIcon: true,
+			} satisfies IActionListItem<ILanguageModelChatMetadataAndIdentifier>);
+
+			// Add models in this category
+			for (const item of modelsInCategory) {
+				items.push({
+					item: item.model,
+					description: item.model.metadata.description,
+					kind: ActionListItemKind.Action,
+					canPreview: false,
+					group: { title: '', icon: ThemeIcon.fromId(item.isCurrent ? Codicon.check.id : Codicon.blank.id) },
+					disabled: false,
+					hideIcon: false,
+					label: item.model.metadata.name,
+				} satisfies IActionListItem<ILanguageModelChatMetadataAndIdentifier>);
+			}
+
+			// Remove this category from the map so we don't process it again
+			modelsByCategory.delete(categoryLabel);
+		}
 
 		const delegate = {
 			onSelect: (item: ILanguageModelChatMetadataAndIdentifier) => {

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

@@ -127,6 +127,7 @@ export interface ILanguageModelChatMetadata {
 
 	readonly isDefault?: boolean;
 	readonly isUserSelectable?: boolean;
+	readonly modelPickerCategory: { label: string };
 	readonly auth?: {
 		readonly providerLabel: string;
 		readonly accountLabel?: string;

src/vs/workbench/contrib/chat/common/modelPicker/modelPickerWidget.ts

@@ -0,0 +1,8 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { localize } from '../../../../../nls.js';
+
+export const DEFAULT_MODEL_PICKER_CATEGORY = { label: localize('chat.modelPicker.other', "Other Models") };

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

@@ -14,6 +14,7 @@ import { ChatMessageRole, IChatResponseFragment, languageModelExtensionPoint, La
 import { IExtensionService, nullExtensionDescription } from '../../../../services/extensions/common/extensions.js';
 import { ExtensionsRegistry } from '../../../../services/extensions/common/extensionsRegistry.js';
 import { MockContextKeyService } from '../../../../../platform/keybinding/test/common/mockKeybindingService.js';
+import { DEFAULT_MODEL_PICKER_CATEGORY } from '../../common/modelPicker/modelPickerWidget.js';
 
 suite('LanguageModels', function () {
 
@@ -50,6 +51,7 @@ suite('LanguageModels', function () {
 				name: 'Pretty Name',
 				vendor: 'test-vendor',
 				family: 'test-family',
+				modelPickerCategory: DEFAULT_MODEL_PICKER_CATEGORY,
 				version: 'test-version',
 				id: 'test-id',
 				maxInputTokens: 100,
@@ -70,6 +72,7 @@ suite('LanguageModels', function () {
 				vendor: 'test-vendor',
 				family: 'test2-family',
 				version: 'test2-version',
+				modelPickerCategory: DEFAULT_MODEL_PICKER_CATEGORY,
 				id: 'test-id',
 				maxInputTokens: 100,
 				maxOutputTokens: 100,
@@ -119,6 +122,7 @@ suite('LanguageModels', function () {
 				id: 'actual-lm',
 				maxInputTokens: 100,
 				maxOutputTokens: 100,
+				modelPickerCategory: DEFAULT_MODEL_PICKER_CATEGORY,
 			},
 			sendChatRequest: async (messages, _from, _options, token) => {
 				// const message = messages.at(-1);

src/vscode-dts/vscode.proposed.chatProvider.d.ts

@@ -70,6 +70,13 @@ declare module 'vscode' {
 			readonly toolCalling?: boolean;
 			readonly agentMode?: boolean;
 		};
+
+		/**
+		 * Optional category to group models by in the model picker.
+		 * Has no effect if `isUserSelectable` is `false`.
+		 * If not specified, the model will appear in the "Other Models" category.
+		 */
+		readonly category?: { label: string };
 	}
 
 	export interface ChatResponseProviderMetadata {