diff --git a/package-lock.json b/package-lock.json
index 6e24a1cf85..88ce2d38a8 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -15,7 +15,7 @@
"@aws-sdk/client-bedrock-runtime": "^3.758.0",
"@bufbuild/protobuf": "^2.2.5",
"@google-cloud/vertexai": "^1.9.3",
- "@google/generative-ai": "^0.18.0",
+ "@google/genai": "^0.9.0",
"@grpc/grpc-js": "^1.9.15",
"@mistralai/mistralai": "^1.5.0",
"@modelcontextprotocol/sdk": "^1.7.0",
@@ -5567,11 +5567,17 @@
"node": ">=18.0.0"
}
},
- "node_modules/@google/generative-ai": {
- "version": "0.18.0",
- "resolved": "https://registry.npmjs.org/@google/generative-ai/-/generative-ai-0.18.0.tgz",
- "integrity": "sha512-AhaIWSpk2tuhYHrBhUqC0xrWWznmYEja1/TRDIb+5kruBU5kUzMlFsXCQNO9PzyTZ4clUJ3CX/Rvy+Xm9x+w3g==",
+ "node_modules/@google/genai": {
+ "version": "0.9.0",
+ "resolved": "https://registry.npmjs.org/@google/genai/-/genai-0.9.0.tgz",
+ "integrity": "sha512-FD2RizYGInsvfjeaN6O+wQGpRnGVglS1XWrGQr8K7D04AfMmvPodDSw94U9KyFtsVLzWH9kmlPyFM+G4jbmkqg==",
"license": "Apache-2.0",
+ "dependencies": {
+ "google-auth-library": "^9.14.2",
+ "ws": "^8.18.0",
+ "zod": "^3.22.4",
+ "zod-to-json-schema": "^3.22.4"
+ },
"engines": {
"node": ">=18.0.0"
}
@@ -12498,9 +12504,10 @@
}
},
"node_modules/google-auth-library": {
- "version": "9.14.0",
- "resolved": "https://registry.npmjs.org/google-auth-library/-/google-auth-library-9.14.0.tgz",
- "integrity": "sha512-Y/eq+RWVs55Io/anIsm24sDS8X79Tq948zVLGaa7+KlJYYqaGwp1YI37w48nzrNi12RgnzMrQD4NzdmCowT90g==",
+ "version": "9.15.1",
+ "resolved": "https://registry.npmjs.org/google-auth-library/-/google-auth-library-9.15.1.tgz",
+ "integrity": "sha512-Jb6Z0+nvECVz+2lzSMt9u98UsoakXxA2HGHMCxh+so3n90XgYWkq5dur19JAJV7ONiJY22yBTyJB1TSkvPq9Ng==",
+ "license": "Apache-2.0",
"dependencies": {
"base64-js": "^1.3.0",
"ecdsa-sig-formatter": "^1.0.11",
diff --git a/package.json b/package.json
index cd7f34e765..fd41c460be 100644
--- a/package.json
+++ b/package.json
@@ -355,7 +355,7 @@
"@aws-sdk/client-bedrock-runtime": "^3.758.0",
"@bufbuild/protobuf": "^2.2.5",
"@google-cloud/vertexai": "^1.9.3",
- "@google/generative-ai": "^0.18.0",
+ "@google/genai": "^0.9.0",
"@grpc/grpc-js": "^1.9.15",
"@mistralai/mistralai": "^1.5.0",
"@modelcontextprotocol/sdk": "^1.7.0",
diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts
index 21dad02a0b..ddb9e74d47 100644
--- a/src/api/providers/gemini.ts
+++ b/src/api/providers/gemini.ts
@@ -1,5 +1,12 @@
-import { Anthropic } from "@anthropic-ai/sdk"
-import { GoogleGenerativeAI } from "@google/generative-ai"
+import type { Anthropic } from "@anthropic-ai/sdk"
+// Restore GenerateContentConfig import and add GenerateContentResponseUsageMetadata
+import {
+ GoogleGenAI,
+ type GenerationConfig,
+ type Content,
+ type GenerateContentConfig,
+ type GenerateContentResponseUsageMetadata,
+} from "@google/genai"
import { withRetry } from "../retry"
import { ApiHandler } from "../"
import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "@shared/api"
@@ -8,45 +15,80 @@ import { ApiStream } from "../transform/stream"
export class GeminiHandler implements ApiHandler {
private options: ApiHandlerOptions
- private client: GoogleGenerativeAI
+ private client: GoogleGenAI // Updated client type
constructor(options: ApiHandlerOptions) {
if (!options.geminiApiKey) {
throw new Error("API key is required for Google Gemini")
}
this.options = options
- this.client = new GoogleGenerativeAI(options.geminiApiKey)
+ // Updated client initialization
+ this.client = new GoogleGenAI({ apiKey: options.geminiApiKey })
}
@withRetry()
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
- const modelOptions = {
- model: this.getModel().id,
- systemInstruction: systemPrompt,
+ const { id: modelId, info: modelInfo } = this.getModel()
+
+ // Re-implement thinking budget logic based on new SDK structure
+ const thinkingBudget = this.options.thinkingBudgetTokens ?? 0
+ const maxBudget = modelInfo.thinkingConfig?.maxBudget ?? 0
+
+ // port add baseUrl configuration for gemini api requests (#2843)
+ const httpOptions = this.options.geminiBaseUrl ? { baseUrl: this.options.geminiBaseUrl } : undefined
+
+ // Base generation config - Restore type and systemInstruction
+ const generationConfig: GenerateContentConfig = {
+ httpOptions,
+ temperature: 0, // Default temperature
+ systemInstruction: systemPrompt, // System prompt belongs here
}
- const clientOptions = this.options.geminiBaseUrl ? { baseUrl: this.options.geminiBaseUrl } : undefined
- const model = this.client.getGenerativeModel(modelOptions, clientOptions)
- const result = await model.generateContentStream({
- contents: messages.map(convertAnthropicMessageToGemini),
- generationConfig: {
- // maxOutputTokens: this.getModel().info.maxTokens,
- temperature: 0,
- },
- })
+ // Convert messages to the format expected by @google/genai
+ // Note: convertAnthropicMessageToGemini might need adjustments
+ const contents: Content[] = messages.map(convertAnthropicMessageToGemini)
- for await (const chunk of result.stream) {
- yield {
- type: "text",
- text: chunk.text(),
+ // Construct the main request config - Type as GenerateContentConfig
+ const requestConfig: GenerateContentConfig = {
+ ...generationConfig,
+ }
+
+ // Add thinking config if the model supports it
+ if (modelInfo.thinkingConfig?.outputPrice !== undefined && maxBudget > 0) {
+ requestConfig.thinkingConfig = {
+ thinkingBudget: thinkingBudget,
}
}
- const response = await result.response
- yield {
- type: "usage",
- inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
- outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0,
+ // Generate content using the new SDK structure via client.models
+ const result = await this.client.models.generateContentStream({
+ model: modelId, // Pass model ID directly
+ contents,
+ config: requestConfig, // Pass the combined config (which includes systemInstruction)
+ })
+
+ // Declare variable to hold the last usage metadata found
+ let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
+
+ // Iterate directly over the stream
+ for await (const chunk of result) {
+ if (chunk.text) {
+ yield {
+ type: "text",
+ text: chunk.text,
+ }
+ }
+ if (chunk.usageMetadata) {
+ lastUsageMetadata = chunk.usageMetadata
+ }
+ }
+
+ if (lastUsageMetadata) {
+ yield {
+ type: "usage",
+ inputTokens: lastUsageMetadata.promptTokenCount ?? 0,
+ outputTokens: lastUsageMetadata.candidatesTokenCount ?? 0,
+ }
}
}
diff --git a/src/api/transform/gemini-format.ts b/src/api/transform/gemini-format.ts
index bb2c207e9a..f69935c6a8 100644
--- a/src/api/transform/gemini-format.ts
+++ b/src/api/transform/gemini-format.ts
@@ -1,14 +1,14 @@
import { Anthropic } from "@anthropic-ai/sdk"
-import { Content, EnhancedGenerateContentResponse, InlineDataPart, Part, TextPart } from "@google/generative-ai"
+import { Content, GenerateContentResponse, Part } from "@google/genai"
export function convertAnthropicContentToGemini(content: string | Anthropic.ContentBlockParam[]): Part[] {
if (typeof content === "string") {
- return [{ text: content } as TextPart]
+ return [{ text: content }]
}
- return content.flatMap((block) => {
+ return content.flatMap((block): Part => {
switch (block.type) {
case "text":
- return { text: block.text } as TextPart
+ return { text: block.text }
case "image":
if (block.source.type !== "base64") {
throw new Error("Unsupported image source type")
@@ -18,7 +18,7 @@ export function convertAnthropicContentToGemini(content: string | Anthropic.Cont
data: block.source.data,
mimeType: block.source.media_type,
},
- } as InlineDataPart
+ }
default:
throw new Error(`Unsupported content block type: ${block.type}`)
}
@@ -39,16 +39,14 @@ export function unescapeGeminiContent(content: string) {
return content.replace(/\\n/g, "\n").replace(/\\'/g, "'").replace(/\\"/g, '"').replace(/\\r/g, "\r").replace(/\\t/g, "\t")
}
-export function convertGeminiResponseToAnthropic(response: EnhancedGenerateContentResponse): Anthropic.Messages.Message {
+export function convertGeminiResponseToAnthropic(response: GenerateContentResponse): Anthropic.Messages.Message {
const content: Anthropic.Messages.ContentBlock[] = []
- // Add the main text response
- const text = response.text()
+ const text = response.text
if (text) {
content.push({ type: "text", text, citations: null })
}
- // Determine stop reason
let stop_reason: Anthropic.Messages.Message["stop_reason"] = null
const finishReason = response.candidates?.[0]?.finishReason
if (finishReason) {
@@ -64,12 +62,11 @@ export function convertGeminiResponseToAnthropic(response: EnhancedGenerateConte
case "OTHER":
stop_reason = "stop_sequence"
break
- // Add more cases if needed
}
}
return {
- id: `msg_${Date.now()}`, // Generate a unique ID
+ id: `msg_${Date.now()}`,
type: "message",
role: "assistant",
content,
diff --git a/src/shared/api.ts b/src/shared/api.ts
index 6f445ba6fc..cd5480325c 100644
--- a/src/shared/api.ts
+++ b/src/shared/api.ts
@@ -104,6 +104,11 @@ export interface ModelInfo {
inputPriceTiers?: PriceTier[] // Add for tiered input pricing
outputPrice?: number // Keep for non-tiered output models
outputPriceTiers?: PriceTier[] // Add for tiered output pricing
+ thinkingConfig?: {
+ maxBudget?: number // Max allowed thinking budget tokens
+ outputPrice?: number // Output price per million tokens when budget > 0
+ outputPriceTiers?: PriceTier[] // Optional: Tiered output price when budget > 0
+ }
cacheWritesPrice?: number
cacheReadsPrice?: number
description?: string
@@ -409,6 +414,18 @@ export const vertexModels = {
{ tokenLimit: Infinity, price: 15.0 }, // Output price for > 200k input tokens
],
},
+ "gemini-2.5-flash-preview-04-17": {
+ maxTokens: 65536,
+ contextWindow: 1_048_576,
+ supportsImages: true,
+ supportsPromptCache: false,
+ inputPrice: 0.15,
+ outputPrice: 0.6,
+ thinkingConfig: {
+ maxBudget: 24576,
+ outputPrice: 3.5,
+ },
+ },
"gemini-2.0-flash-thinking-exp-01-21": {
maxTokens: 65_536,
contextWindow: 1_048_576,
@@ -505,6 +522,18 @@ export const geminiModels = {
{ tokenLimit: Infinity, price: 15.0 }, // Output price for > 200k input tokens
],
},
+ "gemini-2.5-flash-preview-04-17": {
+ maxTokens: 65536,
+ contextWindow: 1_048_576,
+ supportsImages: true,
+ supportsPromptCache: false,
+ inputPrice: 0.15,
+ outputPrice: 0.6,
+ thinkingConfig: {
+ maxBudget: 24576,
+ outputPrice: 3.5,
+ },
+ },
"gemini-2.0-flash-001": {
maxTokens: 8192,
contextWindow: 1_048_576,
diff --git a/src/utils/cost.ts b/src/utils/cost.ts
index d383575620..facf788ec7 100644
--- a/src/utils/cost.ts
+++ b/src/utils/cost.ts
@@ -7,7 +7,10 @@ function calculateApiCostInternal(
cacheCreationInputTokens: number,
cacheReadInputTokens: number,
totalInputTokensForPricing?: number, // The *total* input tokens, used for tiered pricing lookup
+ thinkingBudgetTokens?: number, // Add thinking budget info
): number {
+ const usedThinkingBudget = thinkingBudgetTokens && thinkingBudgetTokens > 0
+
// Determine effective input price
let effectiveInputPrice = modelInfo.inputPrice || 0
if (modelInfo.inputPriceTiers && modelInfo.inputPriceTiers.length > 0 && totalInputTokensForPricing !== undefined) {
@@ -23,10 +26,15 @@ function calculateApiCostInternal(
}
}
- // Determine effective output price (based on total *input* tokens for pricing)
- let effectiveOutputPrice = modelInfo.outputPrice || 0
- if (modelInfo.outputPriceTiers && modelInfo.outputPriceTiers.length > 0 && totalInputTokensForPricing !== undefined) {
- // Ensure tiers are sorted by tokenLimit ascending before finding
+ // Determine effective output price
+ let effectiveOutputPrice = 0
+ // Check if thinking budget was used and has a specific price
+ if (usedThinkingBudget && modelInfo.thinkingConfig?.outputPrice !== undefined) {
+ effectiveOutputPrice = modelInfo.thinkingConfig.outputPrice
+ // TODO: Add support for tiered thinking budget output pricing if needed in the future
+ // } else if (usedThinkingBudget && modelInfo.thinkingConfig?.outputPriceTiers) { ... }
+ } else if (modelInfo.outputPriceTiers && modelInfo.outputPriceTiers.length > 0 && totalInputTokensForPricing !== undefined) {
+ // Use standard tiered output pricing (based on total *input* tokens for pricing)
const sortedOutputTiers = [...modelInfo.outputPriceTiers].sort((a, b) => a.tokenLimit - b.tokenLimit)
const tier = sortedOutputTiers.find((t) => totalInputTokensForPricing! <= t.tokenLimit)
if (tier) {
@@ -55,9 +63,12 @@ export function calculateApiCostAnthropic(
outputTokens: number,
cacheCreationInputTokens?: number,
cacheReadInputTokens?: number,
+ thinkingBudgetTokens?: number,
): number {
const cacheCreationInputTokensNum = cacheCreationInputTokens || 0
const cacheReadInputTokensNum = cacheReadInputTokens || 0
+ // Anthropic style: inputTokens already represents the total, so pass it directly for tiered pricing lookup if needed
+ // (though Anthropic models currently don't use tiered pricing based on input size)
// Anthropic style doesn't need totalInputTokensForPricing as its inputTokens already represents the total
return calculateApiCostInternal(
modelInfo,
@@ -65,17 +76,19 @@ export function calculateApiCostAnthropic(
outputTokens,
cacheCreationInputTokensNum,
cacheReadInputTokensNum,
- undefined, // Pass undefined for totalInputTokensForPricing
+ inputTokens,
+ thinkingBudgetTokens,
)
}
// For OpenAI compliant usage, the input tokens count INCLUDES the cached tokens
export function calculateApiCostOpenAI(
modelInfo: ModelInfo,
- inputTokens: number,
+ inputTokens: number, // For OpenAI-style, this includes cached tokens
outputTokens: number,
cacheCreationInputTokens?: number,
cacheReadInputTokens?: number,
+ thinkingBudgetTokens?: number, // Pass thinking budget info
): number {
const cacheCreationInputTokensNum = cacheCreationInputTokens || 0
const cacheReadInputTokensNum = cacheReadInputTokens || 0
@@ -84,10 +97,11 @@ export function calculateApiCostOpenAI(
// Pass the original 'inputTokens' as 'totalInputTokensForPricing' for tier lookup
return calculateApiCostInternal(
modelInfo,
- nonCachedInputTokens, // Pass the adjusted token count here
+ nonCachedInputTokens,
outputTokens,
cacheCreationInputTokensNum,
cacheReadInputTokensNum,
- inputTokens, // Pass the original total input tokens for pricing tier lookup
+ inputTokens,
+ thinkingBudgetTokens,
)
}
diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx
index 771bc651a9..8894f28c02 100644
--- a/webview-ui/src/components/settings/ApiOptions.tsx
+++ b/webview-ui/src/components/settings/ApiOptions.tsx
@@ -834,6 +834,15 @@ const ApiOptions = ({
)}
+
+ {/* Add Thinking Budget Slider specifically for gemini-2.5-flash-preview-04-17 */}
+ {selectedProvider === "gemini" && selectedModelId === "gemini-2.5-flash-preview-04-17" && (
+
+ )}
)}
@@ -1748,17 +1757,30 @@ export const formatPrice = (price: number) => {
}
// Returns an array of formatted tier strings
-const formatTiers = (tiers: ModelInfo["inputPriceTiers"]): string[] => {
+const formatTiers = (tiers: ModelInfo["inputPriceTiers"]): JSX.Element[] => {
if (!tiers || tiers.length === 0) {
return []
}
return tiers.map((tier, index, arr) => {
const prevLimit = index > 0 ? arr[index - 1].tokenLimit : 0
- const limitText =
- tier.tokenLimit === Infinity
- ? `> ${prevLimit.toLocaleString()}` // Assumes sorted and Infinity is last
- : `<= ${tier.tokenLimit.toLocaleString()}`
- return `${formatPrice(tier.price)}/million tokens (${limitText} tokens)`
+ return (
+
+ {formatPrice(tier.price)}/million tokens (
+ {tier.tokenLimit === Number.POSITIVE_INFINITY ? (
+
+ {"> "}
+ {prevLimit.toLocaleString()}
+
+ ) : (
+
+ {"<= "}
+ {tier.tokenLimit.toLocaleString()}
+
+ )}
+ {" tokens)"}
+ {index < arr.length - 1 &&
}
+
+ )
})
}
@@ -1776,18 +1798,14 @@ export const ModelInfoView = ({
isPopup?: boolean
}) => {
const isGemini = Object.keys(geminiModels).includes(selectedModelId)
+ const hasThinkingConfig = !!modelInfo.thinkingConfig
// Create elements for tiered pricing separately
const inputPriceElement = modelInfo.inputPriceTiers ? (
Input price:
- {formatTiers(modelInfo.inputPriceTiers).map((tierString, i, arr) => (
-
- {tierString}
- {i < arr.length - 1 &&
}
-
- ))}
+ {formatTiers(modelInfo.inputPriceTiers)}
) : modelInfo.inputPrice !== undefined && modelInfo.inputPrice > 0 ? (
@@ -1795,23 +1813,38 @@ export const ModelInfoView = ({
) : null
- const outputPriceElement = modelInfo.outputPriceTiers ? (
-
- Output price:
- (based on input tokens)
-
- {formatTiers(modelInfo.outputPriceTiers).map((tierString, i, arr) => (
-
- {tierString}
- {i < arr.length - 1 &&
}
-
- ))}
-
- ) : modelInfo.outputPrice !== undefined && modelInfo.outputPrice > 0 ? (
-
- Output price: {formatPrice(modelInfo.outputPrice)}/million tokens
-
- ) : null
+ // --- Output Price Logic ---
+ let outputPriceElement = null
+ if (hasThinkingConfig && modelInfo.outputPrice !== undefined && modelInfo.thinkingConfig?.outputPrice !== undefined) {
+ // Display both standard and thinking budget prices
+ outputPriceElement = (
+
+ Output price (Standard): {formatPrice(modelInfo.outputPrice)}/million
+ tokens
+
+ Output price (Thinking Budget > 0):{" "}
+ {formatPrice(modelInfo.thinkingConfig.outputPrice)}/million tokens
+
+ )
+ } else if (modelInfo.outputPriceTiers) {
+ // Display tiered output pricing
+ outputPriceElement = (
+
+ Output price:
+ (based on input tokens)
+
+ {formatTiers(modelInfo.outputPriceTiers)}
+
+ )
+ } else if (modelInfo.outputPrice !== undefined && modelInfo.outputPrice > 0) {
+ // Display single standard output price
+ outputPriceElement = (
+
+ Output price: {formatPrice(modelInfo.outputPrice)}/million tokens
+
+ )
+ }
+ // --- End Output Price Logic ---
const infoItems = [
modelInfo.description && (
diff --git a/webview-ui/src/components/settings/ThinkingBudgetSlider.tsx b/webview-ui/src/components/settings/ThinkingBudgetSlider.tsx
index 0486f4e186..fc21a85845 100644
--- a/webview-ui/src/components/settings/ThinkingBudgetSlider.tsx
+++ b/webview-ui/src/components/settings/ThinkingBudgetSlider.tsx
@@ -1,10 +1,10 @@
import { memo, useCallback, useState } from "react"
-import { anthropicModels, ApiConfiguration } from "@shared/api"
+import { anthropicModels, ApiConfiguration, geminiDefaultModelId, geminiModels, ModelInfo } from "@shared/api"
import { VSCodeCheckbox } from "@vscode/webview-ui-toolkit/react"
import styled from "styled-components"
// Constants
-const MIN_VALID_TOKENS = 1024
+const DEFAULT_MIN_VALID_TOKENS = 1024
const MAX_PERCENTAGE = 0.8
const THUMB_SIZE = 16
@@ -82,11 +82,23 @@ const RangeInput = styled.input<{ $value: number; $min: number; $max: number }>`
interface ThinkingBudgetSliderProps {
apiConfiguration: ApiConfiguration | undefined
setApiConfiguration: (apiConfiguration: ApiConfiguration) => void
+ maxBudget?: number
}
-const ThinkingBudgetSlider = ({ apiConfiguration, setApiConfiguration }: ThinkingBudgetSliderProps) => {
- const maxTokens = anthropicModels["claude-3-7-sonnet-20250219"].maxTokens
- const maxSliderValue = Math.floor(maxTokens * MAX_PERCENTAGE)
+const ThinkingBudgetSlider = ({ apiConfiguration, setApiConfiguration, maxBudget }: ThinkingBudgetSliderProps) => {
+ const maxTokens =
+ apiConfiguration?.apiProvider === "gemini"
+ ? geminiModels[geminiDefaultModelId].maxTokens
+ : anthropicModels["claude-3-7-sonnet-20250219"].maxTokens
+
+ // use maxBudget prop if provided, otherwise apply the percentage cap to maxTokens
+ const maxSliderValue = (() => {
+ if (maxBudget !== undefined) {
+ return maxBudget
+ }
+ return Math.floor(maxTokens * MAX_PERCENTAGE)
+ })()
+
const isEnabled = (apiConfiguration?.thinkingBudgetTokens || 0) > 0
// Add local state for the slider value
@@ -106,7 +118,7 @@ const ThinkingBudgetSlider = ({ apiConfiguration, setApiConfiguration }: Thinkin
const handleToggleChange = (event: any) => {
const isChecked = (event.target as HTMLInputElement).checked
- const newValue = isChecked ? MIN_VALID_TOKENS : 0
+ const newValue = isChecked ? DEFAULT_MIN_VALID_TOKENS : 0
setLocalValue(newValue)
setApiConfiguration({
...apiConfiguration,
@@ -130,18 +142,18 @@ const ThinkingBudgetSlider = ({ apiConfiguration, setApiConfiguration }: Thinkin