Allow experiment model to be default (#268)

* Use exp model when available as the default model

* only override default for known models

Author: karthiknadig

src/extension/prompt/vscode-node/endpointProviderImpl.ts

@@ -13,7 +13,7 @@ import { AutoChatEndpoint, resolveAutoChatEndpoint } from '../../../platform/end
 import { ChatEndpoint } from '../../../platform/endpoint/node/chatEndpoint';
 import { EmbeddingEndpoint } from '../../../platform/endpoint/node/embeddingsEndpoint';
 import { IModelMetadataFetcher, ModelMetadataFetcher } from '../../../platform/endpoint/node/modelMetadataFetcher';
-import { ProxyExperimentEndpoint } from '../../../platform/endpoint/node/proxyExperimentEndpoint';
+import { applyExperimentModifications, getCustomDefaultModelExperimentConfig, ProxyExperimentEndpoint } from '../../../platform/endpoint/node/proxyExperimentEndpoint';
 import { ExtensionContributedChatEndpoint } from '../../../platform/endpoint/vscode-node/extChatEndpoint';
 import { IEnvService } from '../../../platform/env/common/envService';
 import { ILogService } from '../../../platform/log/common/logService';
@@ -24,6 +24,7 @@ import { ITelemetryService } from '../../../platform/telemetry/common/telemetry'
 import { TokenizerType } from '../../../util/common/tokenizer';
 import { IInstantiationService, ServicesAccessor } from '../../../util/vs/platform/instantiation/common/instantiation';
 
+
 export class ProductionEndpointProvider implements IEndpointProvider {
 
 	declare readonly _serviceBrand: undefined;
@@ -90,7 +91,7 @@ export class ProductionEndpointProvider implements IEndpointProvider {
 	private getOrCreateProxyExperimentEndpointInstance(name: string, id: string, endpoint: IChatEndpoint): IChatEndpoint {
 		let chatEndpoint = this._chatEndpoints.get(id);
 		if (!chatEndpoint) {
-			chatEndpoint = new ProxyExperimentEndpoint(name, id, endpoint);
+			chatEndpoint = new ProxyExperimentEndpoint(name, id, endpoint, /* isDefault: */ true);
 			this._chatEndpoints.set(id, chatEndpoint);
 		}
 		return chatEndpoint;
@@ -104,22 +105,11 @@ export class ProductionEndpointProvider implements IEndpointProvider {
 			this._embeddingEndpoints.set(modelId, embeddingEndpoint);
 		}
 		return embeddingEndpoint;
-
-	}
-
-	private getExperimentData(): { selected: string; name: string; id: string } | undefined {
-		const selected = this._expService.getTreatmentVariable<string>('vscode', 'custommodel1');
-		const id = this._expService.getTreatmentVariable<string>('vscode', 'custommodel1.id');
-		const name = this._expService.getTreatmentVariable<string>('vscode', 'custommodel1.name');
-		if (selected && id && name) {
-			return { selected, id, name };
-		}
-		return undefined;
 	}
 
 	async getChatEndpoint(requestOrFamilyOrModel: LanguageModelChat | ChatRequest | ChatEndpointFamily): Promise<IChatEndpoint> {
 		this._logService.logger.trace(`Resolving chat model`);
-		const experimentModelConfig = this.getExperimentData();
+		const experimentModelConfig = getCustomDefaultModelExperimentConfig(this._expService);
 
 		if (this._overridenChatModel) {
 			// Override, only allowed by internal users. Sets model based on setting
@@ -142,16 +132,20 @@ export class ProductionEndpointProvider implements IEndpointProvider {
 		let endpoint: IChatEndpoint;
 		if (typeof requestOrFamilyOrModel === 'string') {
 			// The family case, resolve the chat model for the passed in family
-			const modelMetadata = await this._modelFetcher.getChatModelFromFamily(requestOrFamilyOrModel);
-			endpoint = this.getOrCreateChatEndpointInstance(modelMetadata);
+			let modelMetadata = await this._modelFetcher.getChatModelFromFamily(requestOrFamilyOrModel);
+			modelMetadata = applyExperimentModifications(modelMetadata, experimentModelConfig);
+			endpoint = this.getOrCreateChatEndpointInstance(modelMetadata!);
 		} else {
 			const model = 'model' in requestOrFamilyOrModel ? requestOrFamilyOrModel.model : requestOrFamilyOrModel;
 			if (experimentModelConfig && model && model.id === experimentModelConfig.id) {
 				endpoint = (await this.getAllChatEndpoints()).find(e => e.model === experimentModelConfig.selected) || await this.getChatEndpoint('gpt-4.1');
 			} else if (model && model.vendor === 'copilot' && model.id === AutoChatEndpoint.id) {
 				return resolveAutoChatEndpoint(this, this._expService, (requestOrFamilyOrModel as ChatRequest)?.prompt);
 			} else if (model && model.vendor === 'copilot') {
-				const modelMetadata = await this._modelFetcher.getChatModelFromApiModel(model);
+				let modelMetadata = await this._modelFetcher.getChatModelFromApiModel(model);
+				if (modelMetadata) {
+					modelMetadata = applyExperimentModifications(modelMetadata, experimentModelConfig);
+				}
 				// If we fail to resolve a model since this is panel we give GPT-4.1. This really should never happen as the picker is powered by the same service.
 				endpoint = modelMetadata ? this.getOrCreateChatEndpointInstance(modelMetadata) : await this.getChatEndpoint('gpt-4.1');
 			} else if (model) {
@@ -194,9 +188,10 @@ export class ProductionEndpointProvider implements IEndpointProvider {
 		const models: IChatModelInformation[] = await this._modelFetcher.getAllChatModels();
 		const chatEndpoints = [];
 
-		const experimentModelConfig = this.getExperimentData();
+		const experimentModelConfig = getCustomDefaultModelExperimentConfig(this._expService);
 
-		for (const model of models) {
+		for (let model of models) {
+			model = applyExperimentModifications(model, experimentModelConfig) ?? model;
 			const chatEndpoint = this.getOrCreateChatEndpointInstance(model);
 			chatEndpoints.push(chatEndpoint);
 			if (experimentModelConfig && chatEndpoint.model === experimentModelConfig.selected) {

src/platform/endpoint/node/proxyExperimentEndpoint.ts

@@ -15,14 +15,21 @@ import { FinishedCallback, OptionalChatRequestParams } from '../../networking/co
 import { Response } from '../../networking/common/fetcherService';
 import { IChatEndpoint, IEndpointBody } from '../../networking/common/networking';
 import { ChatCompletion } from '../../networking/common/openai';
+import { IExperimentationService } from '../../telemetry/common/nullExperimentationService';
 import { ITelemetryService, TelemetryProperties } from '../../telemetry/common/telemetry';
 import { TelemetryData } from '../../telemetry/common/telemetryData';
+import { IChatModelInformation } from '../common/endpointProvider';
 
 export class ProxyExperimentEndpoint implements IChatEndpoint {
 	public readonly showInModelPicker: boolean;
 	public readonly family: string;
 
-	constructor(public readonly name: string, public readonly model: string, public readonly selectedEndpoint: IChatEndpoint) {
+	constructor(
+		public readonly name: string,
+		public readonly model: string,
+		public readonly selectedEndpoint: IChatEndpoint,
+		private readonly _isDefault: boolean
+	) {
 		// This is a proxy endpoint that wraps another endpoint, typically used for experiments.
 		// This should be used to show the endpoint in the model picker, when in experiment.
 		this.showInModelPicker = true;
@@ -69,6 +76,9 @@ export class ProxyExperimentEndpoint implements IChatEndpoint {
 	}
 
 	get isDefault(): boolean {
+		if (this._isDefault !== undefined) {
+			return this._isDefault;
+		}
 		return this.selectedEndpoint.isDefault;
 	}
 
@@ -115,5 +125,32 @@ export class ProxyExperimentEndpoint implements IChatEndpoint {
 	acquireTokenizer(): ITokenizer {
 		return this.selectedEndpoint.acquireTokenizer();
 	}
+}
+
+
+interface ExperimentConfig {
+	selected: string;
+	name: string;
+	id: string;
+}
+
+export function getCustomDefaultModelExperimentConfig(expService: IExperimentationService): ExperimentConfig | undefined {
+	const selected = expService.getTreatmentVariable<string>('vscode', 'custommodel1');
+	const id = expService.getTreatmentVariable<string>('vscode', 'custommodel1.id');
+	const name = expService.getTreatmentVariable<string>('vscode', 'custommodel1.name');
+	if (selected && id && name) {
+		return { selected, id, name };
+	}
+	return undefined;
+}
 
+export function applyExperimentModifications(
+	modelMetadata: IChatModelInformation,
+	experimentConfig: ExperimentConfig | undefined
+): IChatModelInformation {
+	const knownDefaults = ['gpt-4.1'];
+	if (modelMetadata && experimentConfig && modelMetadata.is_chat_default && knownDefaults.includes(modelMetadata.id)) {
+		return { ...modelMetadata, is_chat_default: false };
+	}
+	return modelMetadata;
 }