From d67a72b85b61695a489567bfdd71c047d53c6f3a Mon Sep 17 00:00:00 2001
From: Eric Amodio <eamodio@gmail.com>
Date: Tue, 6 Jun 2023 14:58:09 -0400
Subject: [PATCH] Adds switch AI model command

---
 package.json                    |  9 +++++++
 src/commands.ts                 |  1 +
 src/commands/switchAIModel.ts   | 21 ++++++++++++++++
 src/constants.ts                |  1 +
 src/quickpicks/aiModelPicker.ts | 55 +++++++++++++++++++++++++++++++++++++++++
 5 files changed, 87 insertions(+)
 create mode 100644 src/commands/switchAIModel.ts
 create mode 100644 src/quickpicks/aiModelPicker.ts

diff --git a/package.json b/package.json
index e6e03f7..0e1aedb 100644
--- a/package.json
+++ b/package.json
@@ -5059,6 +5059,11 @@
 				"category": "GitLens+"
 			},
 			{
+				"command": "gitlens.switchAIModel",
+				"title": "Switch AI Model",
+				"category": "GitLens"
+			},
+			{
 				"command": "gitlens.switchMode",
 				"title": "Switch Mode",
 				"category": "GitLens"
@@ -8124,6 +8129,10 @@
 					"when": "!gitlens:disabled && !gitlens:readonly && !gitlens:untrusted && !gitlens:hasVirtualFolders"
 				},
 				{
+					"command": "gitlens.switchAIModel",
+					"when": "gitlens:enabled"
+				},
+				{
 					"command": "gitlens.switchMode",
 					"when": "gitlens:enabled"
 				},
diff --git a/src/commands.ts b/src/commands.ts
index c383f45..f3c1a18 100644
--- a/src/commands.ts
+++ b/src/commands.ts
@@ -58,6 +58,7 @@ export * from './commands/showQuickStashList';
 export * from './commands/showView';
 export * from './commands/stashApply';
 export * from './commands/stashSave';
+export * from './commands/switchAIModel';
 export * from './commands/switchMode';
 export * from './commands/toggleCodeLens';
 export * from './commands/toggleFileAnnotations';
diff --git a/src/commands/switchAIModel.ts b/src/commands/switchAIModel.ts
new file mode 100644
index 0000000..ac60005
--- /dev/null
+++ b/src/commands/switchAIModel.ts
@@ -0,0 +1,21 @@
+import { Commands } from '../constants';
+import type { Container } from '../container';
+import { showAIModelPicker } from '../quickpicks/aiModelPicker';
+import { command } from '../system/command';
+import { configuration } from '../system/configuration';
+import { Command } from './base';
+
+@command()
+export class SwitchAIModelCommand extends Command {
+	constructor(private readonly container: Container) {
+		super(Commands.SwitchAIModel);
+	}
+
+	async execute() {
+		const pick = await showAIModelPicker();
+		if (pick == null) return;
+
+		await configuration.updateEffective('ai.experimental.provider', pick.provider);
+		await configuration.updateEffective(`ai.experimental.${pick.provider}.model`, pick.model);
+	}
+}
diff --git a/src/constants.ts b/src/constants.ts
index cff0462..ae61ed3 100644
--- a/src/constants.ts
+++ b/src/constants.ts
@@ -284,6 +284,7 @@ export const enum Commands {
 	StashApply = 'gitlens.stashApply',
 	StashSave = 'gitlens.stashSave',
 	StashSaveFiles = 'gitlens.stashSaveFiles',
+	SwitchAIModel = 'gitlens.switchAIModel',
 	SwitchMode = 'gitlens.switchMode',
 	ToggleCodeLens = 'gitlens.toggleCodeLens',
 	ToggleFileBlame = 'gitlens.toggleFileBlame',
diff --git a/src/quickpicks/aiModelPicker.ts b/src/quickpicks/aiModelPicker.ts
new file mode 100644
index 0000000..cedc4c7
--- /dev/null
+++ b/src/quickpicks/aiModelPicker.ts
@@ -0,0 +1,55 @@
+import type { QuickPickItem } from 'vscode';
+import { QuickPickItemKind, window } from 'vscode';
+import type { AnthropicModels } from '../ai/anthropicProvider';
+import type { OpenAIModels } from '../ai/openaiProvider';
+import type { AIProviders } from '../constants';
+import { configuration } from '../system/configuration';
+
+export interface ModelQuickPickItem extends QuickPickItem {
+	provider: AIProviders;
+	model: OpenAIModels | AnthropicModels;
+}
+
+export async function showAIModelPicker(): Promise<ModelQuickPickItem | undefined> {
+	const provider = configuration.get('ai.experimental.provider') ?? 'openai';
+	let model = configuration.get(`ai.experimental.${provider}.model`);
+	if (model == null) {
+		model = provider === 'anthropic' ? 'claude-v1' : 'gpt-3.5-turbo';
+	}
+
+	type QuickPickSeparator = { label: string; kind: QuickPickItemKind.Separator };
+
+	const items: (ModelQuickPickItem | QuickPickSeparator)[] = [
+		{ label: 'OpenAI', kind: QuickPickItemKind.Separator },
+		{ label: 'OpenAI', description: 'GPT 3.5 Turbo', provider: 'openai', model: 'gpt-3.5-turbo' },
+		{ label: 'OpenAI', description: 'GPT 4', provider: 'openai', model: 'gpt-4' },
+		{ label: 'OpenAI', description: 'GPT 4 33k', provider: 'openai', model: 'gpt-4-32k' },
+		{ label: 'Anthropic', kind: QuickPickItemKind.Separator },
+		{ label: 'Anthropic', description: 'Claude v1', provider: 'anthropic', model: 'claude-v1' },
+		{ label: 'Anthropic', description: 'Claude v1 100k', provider: 'anthropic', model: 'claude-v1-100k' },
+		{ label: 'Anthropic', description: 'Claude Instant v1', provider: 'anthropic', model: 'claude-instant-v1' },
+		{
+			label: 'Anthropic',
+			description: 'Claude Instant v1 100k',
+			provider: 'anthropic',
+			model: 'claude-instant-v1-100k',
+		},
+	];
+
+	for (const item of items) {
+		if (item.kind === QuickPickItemKind.Separator) continue;
+
+		if (item.model === model) {
+			item.picked = true;
+			break;
+		}
+	}
+
+	const pick = (await window.showQuickPick(items, {
+		title: 'Switch AI Model',
+		placeHolder: 'select an AI model to use for experimental AI features',
+		matchOnDescription: true,
+	})) as ModelQuickPickItem | undefined;
+
+	return pick;
+}