From fffcc804779a84293ae96f2fec8011dff2589ff4 Mon Sep 17 00:00:00 2001 From: monotykamary Date: Wed, 23 Apr 2025 06:50:59 +0700 Subject: [PATCH] feat: update gemini sdk and add thinking budget support (#2964) * feat: update gemini sdk and add thinking budget support * chore: remove redundant comments * refactor(thinking-budget): abstract calculation for budget slider * chore: remove some more redundant comments --- package-lock.json | 23 +++-- package.json | 2 +- src/api/providers/gemini.ts | 92 ++++++++++++++----- src/api/transform/gemini-format.ts | 19 ++-- src/shared/api.ts | 29 ++++++ src/utils/cost.ts | 30 ++++-- .../src/components/settings/ApiOptions.tsx | 91 ++++++++++++------ .../settings/ThinkingBudgetSlider.tsx | 32 +++++-- 8 files changed, 226 insertions(+), 92 deletions(-) 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