Updates AI provider/model selection flow

Author: axosoft-ramint

src/plus/ai/aiProviderService.ts

@@ -1,5 +1,5 @@
 import type { CancellationToken, Disposable, Event, MessageItem, ProgressOptions } from 'vscode';
-import { env, EventEmitter, window } from 'vscode';
+import { env, EventEmitter, ThemeIcon, window } from 'vscode';
 import type { AIPrimaryProviders, AIProviderAndModel, AIProviders, SupportedAIModels } from '../../constants.ai';
 import { primaryAIProviders } from '../../constants.ai';
 import type { AIGenerateDraftEventData, Source, TelemetryEvents } from '../../constants.telemetry';
@@ -11,7 +11,7 @@ import type { GitRevisionReference } from '../../git/models/reference';
 import type { Repository } from '../../git/models/repository';
 import { uncommitted, uncommittedStaged } from '../../git/models/revision';
 import { assertsCommitHasFullDetails } from '../../git/utils/commit.utils';
-import { showAIModelPicker } from '../../quickpicks/aiModelPicker';
+import { showAIModelPicker, showAIProviderPicker } from '../../quickpicks/aiPicker';
 import { configuration } from '../../system/-webview/configuration';
 import type { Storage } from '../../system/-webview/storage';
 import { supportedInVSCodeVersion } from '../../system/-webview/vscode';
@@ -22,6 +22,7 @@ import { lazy } from '../../system/lazy';
 import type { Deferred } from '../../system/promise';
 import { getSettledValue } from '../../system/promise';
 import type { ServerConnection } from '../gk/serverConnection';
+import { ensureAccountQuickPick } from '../gk/utils/-webview/acount.utils';
 import type { AIActionType, AIModel, AIModelDescriptor } from './models/model';
 import type { PromptTemplateContext } from './models/promptTemplates';
 import type { AIProvider, AIRequestResult } from './models/provider';
@@ -184,10 +185,48 @@ export class AIProviderService implements Disposable {
 
 		if (options?.silent) return undefined;
 
-		const pick = await showAIModelPicker(this.container, cfg);
-		if (pick == null) return undefined;
+		let chosenProvider: AIProviders | undefined = undefined;
+		let chosenModel: AIModel | undefined = undefined;
+
+		if (!options?.force) {
+			const vsCodeModels = await this.getModels('vscode');
+			if (vsCodeModels.length !== 0) {
+				chosenProvider = 'vscode';
+			} else if ((await this.container.subscription.getSubscription()).account?.verified) {
+				chosenProvider = 'gitkraken';
+				const gitkrakenModels = await this.getModels('gitkraken');
+				chosenModel = gitkrakenModels.find(m => m.default);
+			}
+		}
+
+		if (chosenProvider == null) {
+			chosenProvider = (await showAIProviderPicker(this.container, cfg))?.provider;
+			if (chosenProvider == null) return;
+			if (
+				(chosenProvider === 'gitkraken' ||
+					(chosenProvider !== 'vscode' &&
+						(await this.container.storage.getSecret(`gitlens.${chosenProvider}.key`)) == null)) &&
+				!(await ensureAccountQuickPick(
+					this.container,
+					{
+						label: 'Use AI-powered GitLens features like Generate Commit Message, Explain Commit, and more.',
+						iconPath: new ThemeIcon('sparkle'),
+					},
+					source,
+				))
+			) {
+				return;
+			}
+		}
 
-		const model = await this.getOrUpdateModel(pick.model);
+		if (!(await this.ensureProviderConfigured(chosenProvider))) return;
+
+		if (chosenModel == null) {
+			chosenModel = (await showAIModelPicker(this.container, chosenProvider, cfg))?.model;
+			if (chosenModel == null) return;
+		}
+
+		const model = await this.getOrUpdateModel(chosenModel);
 
 		this.container.telemetry.sendEvent(
 			'ai/switchModel',
@@ -204,6 +243,24 @@ export class AIProviderService implements Disposable {
 		return model;
 	}
 
+	private async ensureProviderConfigured(providerId: AIProviders): Promise<boolean> {
+		const key = await this.container.storage.getSecret(`gitlens.${providerId}.key`);
+		if (key != null) return true;
+
+		if (this._provider != null && providerId === this._provider.id) return this._provider.ensureConfigured();
+		const type = await _supportedProviderTypes.get(providerId)?.value;
+		if (type == null) {
+			return false;
+		}
+
+		const p = new type(this.container, this.connection);
+		try {
+			return await p.ensureConfigured();
+		} finally {
+			p.dispose();
+		}
+	}
+
 	private getOrUpdateModel(model: AIModel): Promise<AIModel | undefined>;
 	private getOrUpdateModel<T extends AIProviders>(providerId: T, modelId: string): Promise<AIModel | undefined>;
 	private async getOrUpdateModel(
@@ -592,6 +649,14 @@ export class AIProviderService implements Disposable {
 		return changes;
 	}
 
+	async resetProvider(provider: AIProviders): Promise<void> {
+		void env.clipboard.writeText((await this.container.storage.getSecret(`gitlens.${provider}.key`)) ?? '');
+		void this.container.storage.deleteSecret(`gitlens.${provider}.key`);
+
+		void this.container.storage.delete(`confirm:ai:tos:${provider}`);
+		void this.container.storage.deleteWorkspace(`confirm:ai:tos:${provider}`);
+	}
+
 	async reset(all?: boolean): Promise<void> {
 		let { _provider: provider } = this;
 		if (provider == null) {
@@ -625,11 +690,7 @@ export class AIProviderService implements Disposable {
 		}
 
 		if (provider != null && result === resetCurrent) {
-			void env.clipboard.writeText((await this.container.storage.getSecret(`gitlens.${provider.id}.key`)) ?? '');
-			void this.container.storage.deleteSecret(`gitlens.${provider.id}.key`);
-
-			void this.container.storage.delete(`confirm:ai:tos:${provider.id}`);
-			void this.container.storage.deleteWorkspace(`confirm:ai:tos:${provider.id}`);
+			void this.resetProvider(provider.id);
 		} else if (result === resetAll) {
 			const keys = [];
 			for (const [providerId] of _supportedProviderTypes) {

src/plus/ai/models/provider.ts

@@ -27,6 +27,7 @@ export interface AIProvider<Provider extends AIProviders = AIProviders> extends
 
 	onDidChange?: Event<void>;
 
+	ensureConfigured(): Promise<boolean>;
 	getModels(): Promise<readonly AIModel<Provider>[]>;
 	getPromptTemplate(action: AIActionType, model: AIModel<Provider>): Promise<PromptTemplate | undefined>;
 

src/plus/ai/openAICompatibleProvider.ts

@@ -59,6 +59,10 @@ export abstract class OpenAICompatibleProvider<T extends AIProviders> implements
 		});
 	}
 
+	async ensureConfigured(): Promise<boolean> {
+		return (await this.getApiKey()) != null;
+	}
+
 	protected getHeaders<TAction extends AIActionType>(
 		_action: TAction,
 		_model: AIModel<T>,

src/plus/ai/vscodeProvider.ts

@@ -54,6 +54,10 @@ export class VSCodeAIProvider implements AIProvider<typeof provider.id> {
 		return models.map(getModelFromChatModel);
 	}
 
+	async ensureConfigured(): Promise<boolean> {
+		return (await this.getModels()).length !== 0;
+	}
+
 	async getPromptTemplate(action: AIActionType, model: VSCodeAIModel): Promise<PromptTemplate | undefined> {
 		return Promise.resolve(getLocalPromptTemplate(action, model));
 	}

src/plus/gk/utils/-webview/acount.utils.ts

@@ -1,6 +1,10 @@
+import type { QuickPickItem } from 'vscode';
 import { window } from 'vscode';
 import type { Source } from '../../../../constants.telemetry';
 import type { Container } from '../../../../container';
+import { createQuickPickSeparator } from '../../../../quickpicks/items/common';
+import type { DirectiveQuickPickItem } from '../../../../quickpicks/items/directive';
+import { createDirectiveQuickPickItem, Directive } from '../../../../quickpicks/items/directive';
 
 export async function ensureAccount(container: Container, title: string, source: Source): Promise<boolean> {
 	while (true) {
@@ -52,3 +56,54 @@ export async function ensureAccount(container: Container, title: string, source:
 
 	return true;
 }
+
+export async function ensureAccountQuickPick(
+	container: Container,
+	descriptionItem: QuickPickItem,
+	source?: Source,
+): Promise<boolean> {
+	while (true) {
+		const account = (await container.subscription.getSubscription()).account;
+		if (account?.verified === true) break;
+		const directives: DirectiveQuickPickItem[] = [
+			createDirectiveQuickPickItem(Directive.Noop, undefined, descriptionItem),
+		];
+		let placeholder = 'Requires an account to continue';
+		if (account?.verified === false) {
+			directives.push(
+				createDirectiveQuickPickItem(Directive.RequiresVerification, true),
+				createQuickPickSeparator(),
+				createDirectiveQuickPickItem(Directive.Cancel),
+			);
+			placeholder = 'You must verify your email before you can continue';
+		} else {
+			directives.push(
+				createDirectiveQuickPickItem(Directive.SignIn, true),
+				createQuickPickSeparator(),
+				createDirectiveQuickPickItem(Directive.Cancel),
+			);
+		}
+
+		const result = await window.showQuickPick(directives, {
+			placeHolder: placeholder,
+			ignoreFocusOut: true,
+		});
+
+		if (result == null) return false;
+		if (result.directive === Directive.Noop) continue;
+		if (result.directive === Directive.SignIn) {
+			if (await container.subscription.loginOrSignUp(false, source)) {
+				continue;
+			}
+		}
+		if (result.directive === Directive.RequiresVerification) {
+			if (await container.subscription.resendVerification(source)) {
+				continue;
+			}
+		}
+
+		return false;
+	}
+
+	return true;
+}