Add option to choose different models

src/ClaudeDev.ts

@@ -414,12 +414,9 @@ export class ClaudeDev {
 		}
 	}
 
-	// Calculates cost of a Claude 3.5 Sonnet API request
 	calculateApiCost(inputTokens: number, outputTokens: number): number {
-		const INPUT_COST_PER_MILLION = 3.0 // $3 per million input tokens
-		const OUTPUT_COST_PER_MILLION = 15.0 // $15 per million output tokens
-		const inputCost = (inputTokens / 1_000_000) * INPUT_COST_PER_MILLION
-		const outputCost = (outputTokens / 1_000_000) * OUTPUT_COST_PER_MILLION
+		const inputCost = (this.api.getModel().info.inputPrice / 1_000_000) * inputTokens
+		const outputCost = (this.api.getModel().info.outputPrice / 1_000_000) * outputTokens
 		const totalCost = inputCost + outputCost
 		return totalCost
 	}

src/api/anthropic.ts

@@ -1,6 +1,6 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import { ApiHandler, withoutImageData } from "."
-import { ApiHandlerOptions } from "../shared/api"
+import { anthropicDefaultModelId, AnthropicModelId, anthropicModels, ApiHandlerOptions, ModelInfo } from "../shared/api"
 
 export class AnthropicHandler implements ApiHandler {
 	private options: ApiHandlerOptions
@@ -18,17 +18,20 @@ export class AnthropicHandler implements ApiHandler {
 	): Promise<Anthropic.Messages.Message> {
 		return await this.client.messages.create(
 			{
-				model: "claude-3-5-sonnet-20240620", // https://docs.anthropic.com/en/docs/about-claude/models
-				max_tokens: 8192, // beta max tokens
+				model: this.getModel().id,
+				max_tokens: this.getModel().info.maxTokens,
 				system: systemPrompt,
 				messages,
 				tools,
 				tool_choice: { type: "auto" },
 			},
-			{
-				// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
-				headers: { "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15" },
-			}
+			// https://x.com/alexalbert__/status/1812921642143900036
+			// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
+			this.getModel().id === "claude-3-5-sonnet-20240620"
+				? {
+						headers: { "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15" },
+				  }
+				: undefined
 		)
 	}
 
@@ -41,12 +44,21 @@ export class AnthropicHandler implements ApiHandler {
 		>
 	): any {
 		return {
-			model: "claude-3-5-sonnet-20240620",
-			max_tokens: 8192,
+			model: this.getModel().id,
+			max_tokens: this.getModel().info.maxTokens,
 			system: "(see SYSTEM_PROMPT in src/ClaudeDev.ts)",
 			messages: [{ conversation_history: "..." }, { role: "user", content: withoutImageData(userContent) }],
 			tools: "(see tools in src/ClaudeDev.ts)",
 			tool_choice: { type: "auto" },
 		}
 	}
+
+	getModel(): { id: AnthropicModelId; info: ModelInfo } {
+		const modelId = this.options.apiModelId
+		if (modelId && modelId in anthropicModels) {
+			const id = modelId as AnthropicModelId
+			return { id, info: anthropicModels[id] }
+		}
+		return { id: anthropicDefaultModelId, info: anthropicModels[anthropicDefaultModelId] }
+	}
 }

src/api/bedrock.ts

@@ -1,7 +1,7 @@
 import AnthropicBedrock from "@anthropic-ai/bedrock-sdk"
 import { Anthropic } from "@anthropic-ai/sdk"
-import { ApiHandlerOptions } from "../shared/api"
 import { ApiHandler, withoutImageData } from "."
+import { ApiHandlerOptions, bedrockDefaultModelId, BedrockModelId, bedrockModels, ModelInfo } from "../shared/api"
 
 // https://docs.anthropic.com/en/api/claude-on-amazon-bedrock
 export class AwsBedrockHandler implements ApiHandler {
@@ -28,8 +28,8 @@ export class AwsBedrockHandler implements ApiHandler {
 		tools: Anthropic.Messages.Tool[]
 	): Promise<Anthropic.Messages.Message> {
 		return await this.client.messages.create({
-			model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
-			max_tokens: 4096,
+			model: this.getModel().id,
+			max_tokens: this.getModel().info.maxTokens,
 			system: systemPrompt,
 			messages,
 			tools,
@@ -46,12 +46,21 @@ export class AwsBedrockHandler implements ApiHandler {
 		>
 	): any {
 		return {
-			model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
-			max_tokens: 4096,
+			model: this.getModel().id,
+			max_tokens: this.getModel().info.maxTokens,
 			system: "(see SYSTEM_PROMPT in src/ClaudeDev.ts)",
 			messages: [{ conversation_history: "..." }, { role: "user", content: withoutImageData(userContent) }],
 			tools: "(see tools in src/ClaudeDev.ts)",
 			tool_choice: { type: "auto" },
 		}
 	}
+
+	getModel(): { id: BedrockModelId; info: ModelInfo } {
+		const modelId = this.options.apiModelId
+		if (modelId && modelId in bedrockModels) {
+			const id = modelId as BedrockModelId
+			return { id, info: bedrockModels[id] }
+		}
+		return { id: bedrockDefaultModelId, info: bedrockModels[bedrockDefaultModelId] }
+	}
 }

src/api/index.ts

@@ -1,5 +1,5 @@
 import { Anthropic } from "@anthropic-ai/sdk"
-import { ApiConfiguration } from "../shared/api"
+import { ApiConfiguration, ApiModelId, ModelInfo } from "../shared/api"
 import { AnthropicHandler } from "./anthropic"
 import { AwsBedrockHandler } from "./bedrock"
 import { OpenRouterHandler } from "./openrouter"
@@ -19,6 +19,8 @@ export interface ApiHandler {
 			| Anthropic.ToolResultBlockParam
 		>
 	): any
+
+	getModel(): { id: ApiModelId; info: ModelInfo }
 }
 
 export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {

src/api/openrouter.ts

@@ -1,7 +1,13 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import OpenAI from "openai"
 import { ApiHandler, withoutImageData } from "."
-import { ApiHandlerOptions } from "../shared/api"
+import {
+	ApiHandlerOptions,
+	ModelInfo,
+	openRouterDefaultModelId,
+	OpenRouterModelId,
+	openRouterModels,
+} from "../shared/api"
 
 export class OpenRouterHandler implements ApiHandler {
 	private options: ApiHandlerOptions
@@ -41,8 +47,8 @@ export class OpenRouterHandler implements ApiHandler {
 		}))
 
 		const completion = await this.client.chat.completions.create({
-			model: "anthropic/claude-3.5-sonnet:beta",
-			max_tokens: 4096,
+			model: this.getModel().id,
+			max_tokens: this.getModel().info.maxTokens,
 			messages: openAiMessages,
 			tools: openAiTools,
 			tool_choice: "auto",
@@ -258,11 +264,20 @@ export class OpenRouterHandler implements ApiHandler {
 		>
 	): any {
 		return {
-			model: "anthropic/claude-3.5-sonnet:beta",
-			max_tokens: 4096,
+			model: this.getModel().id,
+			max_tokens: this.getModel().info.maxTokens,
 			messages: [{ conversation_history: "..." }, { role: "user", content: withoutImageData(userContent) }],
 			tools: "(see tools in src/ClaudeDev.ts)",
 			tool_choice: "auto",
 		}
 	}
+
+	getModel(): { id: OpenRouterModelId; info: ModelInfo } {
+		const modelId = this.options.apiModelId
+		if (modelId && modelId in openRouterModels) {
+			const id = modelId as OpenRouterModelId
+			return { id, info: openRouterModels[id] }
+		}
+		return { id: openRouterDefaultModelId, info: openRouterModels[openRouterDefaultModelId] }
+	}
 }

src/providers/ClaudeDevProvider.ts

@@ -1,6 +1,6 @@
 import * as vscode from "vscode"
 import { ClaudeDev } from "../ClaudeDev"
-import { ApiProvider } from "../shared/api"
+import { ApiModelId, ApiProvider } from "../shared/api"
 import { ExtensionMessage } from "../shared/ExtensionMessage"
 import { WebviewMessage } from "../shared/WebviewMessage"
 import { downloadTask, getNonce, getUri, selectImages } from "../utils"
@@ -11,7 +11,7 @@ https://github.com/KumarVariable/vscode-extension-sidebar-html/blob/master/src/c
 */
 
 type SecretKey = "apiKey" | "openRouterApiKey" | "awsAccessKey" | "awsSecretKey"
-type GlobalStateKey = "apiProvider" | "awsRegion" | "maxRequestsPerTask" | "lastShownAnnouncementId"
+type GlobalStateKey = "apiProvider" | "apiModelId" | "awsRegion" | "maxRequestsPerTask" | "lastShownAnnouncementId"
 
 export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 	public static readonly sideBarId = "claude-dev.SidebarProvider" // used in package.json as the view's id. This value cannot be changed due to how vscode caches views based on their id, and updating the id would break existing instances of the extension.
@@ -132,15 +132,8 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 
 	async initClaudeDevWithTask(task?: string, images?: string[]) {
 		await this.clearTask() // ensures that an exising task doesn't exist before starting a new one, although this shouldn't be possible since user must clear task before starting a new one
-		const { apiProvider, apiKey, openRouterApiKey, awsAccessKey, awsSecretKey, awsRegion, maxRequestsPerTask } =
-			await this.getState()
-		this.claudeDev = new ClaudeDev(
-			this,
-			{ apiProvider, apiKey, openRouterApiKey, awsAccessKey, awsSecretKey, awsRegion },
-			maxRequestsPerTask,
-			task,
-			images
-		)
+		const { maxRequestsPerTask, apiConfiguration } = await this.getState()
+		this.claudeDev = new ClaudeDev(this, apiConfiguration, maxRequestsPerTask, task, images)
 	}
 
 	// Send any JSON serializable data to the react app
@@ -255,9 +248,17 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 						break
 					case "apiConfiguration":
 						if (message.apiConfiguration) {
-							const { apiProvider, apiKey, openRouterApiKey, awsAccessKey, awsSecretKey, awsRegion } =
-								message.apiConfiguration
+							const {
+								apiProvider,
+								apiModelId,
+								apiKey,
+								openRouterApiKey,
+								awsAccessKey,
+								awsSecretKey,
+								awsRegion,
+							} = message.apiConfiguration
 							await this.updateGlobalState("apiProvider", apiProvider)
+							await this.updateGlobalState("apiModelId", apiModelId)
 							await this.storeSecret("apiKey", apiKey)
 							await this.storeSecret("openRouterApiKey", openRouterApiKey)
 							await this.storeSecret("awsAccessKey", awsAccessKey)
@@ -308,21 +309,12 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 	}
 
 	async postStateToWebview() {
-		const {
-			apiProvider,
-			apiKey,
-			openRouterApiKey,
-			awsAccessKey,
-			awsSecretKey,
-			awsRegion,
-			maxRequestsPerTask,
-			lastShownAnnouncementId,
-		} = await this.getState()
+		const { apiConfiguration, maxRequestsPerTask, lastShownAnnouncementId } = await this.getState()
 		this.postMessageToWebview({
 			type: "state",
 			state: {
 				version: this.context.extension?.packageJSON?.version ?? "",
-				apiConfiguration: { apiProvider, apiKey, openRouterApiKey, awsAccessKey, awsSecretKey, awsRegion },
+				apiConfiguration,
 				maxRequestsPerTask,
 				themeName: vscode.workspace.getConfiguration("workbench").get<string>("colorTheme"),
 				claudeMessages: this.claudeDev?.claudeMessages || [],
@@ -420,6 +412,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 	async getState() {
 		const [
 			apiProvider,
+			apiModelId,
 			apiKey,
 			openRouterApiKey,
 			awsAccessKey,
@@ -429,6 +422,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 			lastShownAnnouncementId,
 		] = await Promise.all([
 			this.getGlobalState("apiProvider") as Promise<ApiProvider | undefined>,
+			this.getGlobalState("apiModelId") as Promise<ApiModelId | undefined>,
 			this.getSecret("apiKey") as Promise<string | undefined>,
 			this.getSecret("openRouterApiKey") as Promise<string | undefined>,
 			this.getSecret("awsAccessKey") as Promise<string | undefined>,
@@ -438,12 +432,15 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 			this.getGlobalState("lastShownAnnouncementId") as Promise<string | undefined>,
 		])
 		return {
-			apiProvider: apiProvider || "anthropic", // for legacy users that were using Anthropic by default
-			apiKey,
-			openRouterApiKey,
-			awsAccessKey,
-			awsSecretKey,
-			awsRegion,
+			apiConfiguration: {
+				apiProvider: apiProvider || "anthropic", // for legacy users that were using Anthropic by default
+				apiModelId,
+				apiKey,
+				openRouterApiKey,
+				awsAccessKey,
+				awsSecretKey,
+				awsRegion,
+			},
 			maxRequestsPerTask,
 			lastShownAnnouncementId,
 		}

src/shared/api.ts

@@ -1,6 +1,7 @@
 export type ApiProvider = "anthropic" | "openrouter" | "bedrock"
 
 export interface ApiHandlerOptions {
+	apiModelId?: ApiModelId
 	apiKey?: string // anthropic
 	openRouterApiKey?: string
 	awsAccessKey?: string
@@ -11,3 +12,69 @@ export interface ApiHandlerOptions {
 export type ApiConfiguration = ApiHandlerOptions & {
 	apiProvider?: ApiProvider
 }
+
+// Models
+
+export interface ModelInfo {
+	maxTokens: number
+	supportsImages: boolean
+	inputPrice: number
+	outputPrice: number
+}
+
+export type ApiModelId = AnthropicModelId | OpenRouterModelId | BedrockModelId
+
+// Anthropic
+export type AnthropicModelId = keyof typeof anthropicModels
+export const anthropicDefaultModelId: AnthropicModelId = "claude-3-5-sonnet-20240620"
+// https://docs.anthropic.com/en/docs/about-claude/models
+export const anthropicModels = {
+	"claude-3-5-sonnet-20240620": {
+		maxTokens: 8192,
+		supportsImages: true,
+		inputPrice: 3.0, // $3 per million input tokens
+		outputPrice: 15.0, // $15 per million output tokens
+	},
+	"claude-3-opus-20240229": {
+		maxTokens: 4096,
+		supportsImages: true,
+		inputPrice: 15.0,
+		outputPrice: 75.0,
+	},
+	"claude-3-sonnet-20240229": {
+		maxTokens: 4096,
+		supportsImages: true,
+		inputPrice: 2.5,
+		outputPrice: 12.5,
+	},
+	"claude-3-haiku-20240307": {
+		maxTokens: 4096,
+		supportsImages: true,
+		inputPrice: 2.5,
+		outputPrice: 12.5,
+	},
+} as const satisfies Record<string, ModelInfo>
+
+// OpenRouter
+export type OpenRouterModelId = keyof typeof openRouterModels
+export const openRouterDefaultModelId: OpenRouterModelId = "anthropic/claude-3.5-sonnet:beta"
+export const openRouterModels = {
+	"anthropic/claude-3.5-sonnet:beta": {
+		maxTokens: 4096,
+		supportsImages: true,
+		inputPrice: 3.0,
+		outputPrice: 15.0,
+	},
+} as const satisfies Record<string, ModelInfo>
+
+// AWS Bedrock
+export type BedrockModelId = keyof typeof bedrockModels
+export const bedrockDefaultModelId: BedrockModelId = "anthropic.claude-3-5-sonnet-20240620-v1:0"
+export const bedrockModels = {
+	"anthropic.claude-3-5-sonnet-20240620-v1:0": {
+		maxTokens: 4096,
+		supportsImages: true,
+		inputPrice: 3.0,
+		outputPrice: 15.0,
+	},
+} as const satisfies Record<string, ModelInfo> // as const assertion makes the object deeply readonly (just declaring it as const makes it mutable)