From 2c2547a25652b00f9ebece31274cca302ab14022 Mon Sep 17 00:00:00 2001 From: Shariq Riaz Date: Sun, 2 Mar 2025 23:56:31 +0500 Subject: [PATCH] loadbalance --- src/api/index.ts | 51 +++--- src/api/providers/gemini.ts | 146 +++++++++++++++++- src/core/Cline.ts | 60 ++++++- src/core/webview/ClineProvider.ts | 51 +++++- src/shared/api.ts | 5 + src/shared/checkExistApiConfig.ts | 2 + src/shared/globalState.ts | 5 + src/shared/modes.ts | 9 ++ webview-ui/src/components/chat/TaskHeader.tsx | 42 +++++ .../src/components/settings/ApiOptions.tsx | 136 +++++++++++++--- webview-ui/src/utils/validate.ts | 12 +- 11 files changed, 471 insertions(+), 48 deletions(-) diff --git a/src/api/index.ts b/src/api/index.ts index 6cf9317e4..c54cf7a33 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -11,7 +11,7 @@ import { VertexHandler } from "./providers/vertex" import { OpenAiHandler } from "./providers/openai" import { OllamaHandler } from "./providers/ollama" import { LmStudioHandler } from "./providers/lmstudio" -import { GeminiHandler } from "./providers/gemini" +import { GeminiHandler, ApiKeyRotationCallback, RequestCountUpdateCallback } from "./providers/gemini" import { OpenAiNativeHandler } from "./providers/openai-native" import { DeepSeekHandler } from "./providers/deepseek" import { MistralHandler } from "./providers/mistral" @@ -29,41 +29,54 @@ export interface ApiHandler { getModel(): { id: string; info: ModelInfo } } -export function buildApiHandler(configuration: ApiConfiguration): ApiHandler { - const { apiProvider, ...options } = configuration +/** + * Callbacks that can be passed to API handlers + */ +export interface ApiHandlerCallbacks { + onGeminiApiKeyRotation?: ApiKeyRotationCallback + onGeminiRequestCountUpdate?: RequestCountUpdateCallback + geminiInitialRequestCount?: number +} + +export function buildApiHandler(configuration: ApiConfiguration, callbacks?: ApiHandlerCallbacks): ApiHandler { + const { apiProvider, ...handlerOptions } = configuration switch (apiProvider) { case "anthropic": - return new AnthropicHandler(options) + return new AnthropicHandler(handlerOptions) case "glama": - return new GlamaHandler(options) + return new GlamaHandler(handlerOptions) case "openrouter": - return new OpenRouterHandler(options) + return new OpenRouterHandler(handlerOptions) case "bedrock": - return new AwsBedrockHandler(options) + return new AwsBedrockHandler(handlerOptions) case "vertex": - return new VertexHandler(options) + return new VertexHandler(handlerOptions) case "openai": - return new OpenAiHandler(options) + return new OpenAiHandler(handlerOptions) case "ollama": - return new OllamaHandler(options) + return new OllamaHandler(handlerOptions) case "lmstudio": - return new LmStudioHandler(options) + return new LmStudioHandler(handlerOptions) case "gemini": - return new GeminiHandler(options) + return new GeminiHandler(handlerOptions, { + onApiKeyRotation: callbacks?.onGeminiApiKeyRotation, + onRequestCountUpdate: callbacks?.onGeminiRequestCountUpdate, + initialRequestCount: callbacks?.geminiInitialRequestCount, + }) case "openai-native": - return new OpenAiNativeHandler(options) + return new OpenAiNativeHandler(handlerOptions) case "deepseek": - return new DeepSeekHandler(options) + return new DeepSeekHandler(handlerOptions) case "vscode-lm": - return new VsCodeLmHandler(options) + return new VsCodeLmHandler(handlerOptions) case "mistral": - return new MistralHandler(options) + return new MistralHandler(handlerOptions) case "unbound": - return new UnboundHandler(options) + return new UnboundHandler(handlerOptions) case "requesty": - return new RequestyHandler(options) + return new RequestyHandler(handlerOptions) default: - return new AnthropicHandler(options) + return new AnthropicHandler(handlerOptions) } } diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index 0d7179320..c3fc8e54b 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -6,17 +6,156 @@ import { convertAnthropicMessageToGemini } from "../transform/gemini-format" import { ApiStream } from "../transform/stream" const GEMINI_DEFAULT_TEMPERATURE = 0 +const DEFAULT_REQUEST_COUNT = 10 // Default number of requests before switching API keys + +// Define a callback type for API key rotation +export type ApiKeyRotationCallback = (newIndex: number, totalKeys: number, apiKey: string) => void +export type RequestCountUpdateCallback = (newCount: number) => void export class GeminiHandler implements ApiHandler, SingleCompletionHandler { private options: ApiHandlerOptions private client: GoogleGenerativeAI + private requestCount: number = 0 + private onApiKeyRotation?: ApiKeyRotationCallback + private onRequestCountUpdate?: RequestCountUpdateCallback - constructor(options: ApiHandlerOptions) { + constructor( + options: ApiHandlerOptions, + callbacks?: { + onApiKeyRotation?: ApiKeyRotationCallback + onRequestCountUpdate?: RequestCountUpdateCallback + initialRequestCount?: number + }, + ) { this.options = options - this.client = new GoogleGenerativeAI(options.geminiApiKey ?? "not-provided") + this.onApiKeyRotation = callbacks?.onApiKeyRotation + this.onRequestCountUpdate = callbacks?.onRequestCountUpdate + + // Initialize request count from saved state if provided + if (callbacks?.initialRequestCount !== undefined) { + this.requestCount = callbacks.initialRequestCount + console.log(`[GeminiHandler] Initialized with request count: ${this.requestCount}`) + } + + // Initialize with the current API key + const apiKey = this.getCurrentApiKey() + this.client = new GoogleGenerativeAI(apiKey) + + // Log initial API key setup if load balancing is enabled + if ( + this.options.geminiLoadBalancingEnabled && + this.options.geminiApiKeys && + this.options.geminiApiKeys.length > 0 + ) { + console.log( + `[GeminiHandler] Load balancing enabled with ${this.options.geminiApiKeys.length} keys. Current index: ${this.options.geminiCurrentApiKeyIndex ?? 0}`, + ) + } + } + + /** + * Get the current API key based on load balancing settings + */ + private getCurrentApiKey(): string { + // If load balancing is not enabled or there are no multiple API keys, use the single API key + if ( + !this.options.geminiLoadBalancingEnabled || + !this.options.geminiApiKeys || + this.options.geminiApiKeys.length === 0 + ) { + return this.options.geminiApiKey ?? "not-provided" + } + + // Get the current API key index, defaulting to 0 if not set + const currentIndex = this.options.geminiCurrentApiKeyIndex ?? 0 + + // Return the API key at the current index + return this.options.geminiApiKeys[currentIndex] ?? "not-provided" + } + + /** + * Update the client with the next API key if load balancing is enabled + */ + private updateApiKeyIfNeeded(): void { + // If load balancing is not enabled or there are no multiple API keys, do nothing + if ( + !this.options.geminiLoadBalancingEnabled || + !this.options.geminiApiKeys || + this.options.geminiApiKeys.length <= 1 + ) { + return + } + + // Increment the request count + this.requestCount++ + console.log( + `[GeminiHandler] Request count: ${this.requestCount}/${this.options.geminiLoadBalancingRequestCount ?? DEFAULT_REQUEST_COUNT}`, + ) + + // Notify about request count update + if (this.onRequestCountUpdate) { + this.onRequestCountUpdate(this.requestCount) + } + + // Get the request count threshold, defaulting to DEFAULT_REQUEST_COUNT if not set + const requestCountThreshold = this.options.geminiLoadBalancingRequestCount ?? DEFAULT_REQUEST_COUNT + + // If the request count has reached the threshold, switch to the next API key + if (this.requestCount >= requestCountThreshold) { + // Reset the request count + this.requestCount = 0 + + // Notify about request count reset + if (this.onRequestCountUpdate) { + this.onRequestCountUpdate(0) + } + + // Get the current API key index, defaulting to 0 if not set + let currentIndex = this.options.geminiCurrentApiKeyIndex ?? 0 + + // Calculate the next index, wrapping around if necessary + currentIndex = (currentIndex + 1) % this.options.geminiApiKeys.length + + // Notify callback first to update global state + if (this.onApiKeyRotation) { + // Get the API key for the new index + const apiKey = this.options.geminiApiKeys[currentIndex] ?? "not-provided" + + // Only send the first few characters of the API key for security + const maskedKey = apiKey.substring(0, 4) + "..." + apiKey.substring(apiKey.length - 4) + + // Call the callback to update global state + this.onApiKeyRotation(currentIndex, this.options.geminiApiKeys.length, maskedKey) + + // Update the current index in the options AFTER the callback + // This ensures we're using the index that was just set in global state + this.options.geminiCurrentApiKeyIndex = currentIndex + + // Update the client with the new API key + this.client = new GoogleGenerativeAI(apiKey) + + console.log( + `[GeminiHandler] Rotated to API key index: ${currentIndex} (${this.options.geminiApiKeys.length} total keys)`, + ) + } else { + // No callback provided, just update locally + this.options.geminiCurrentApiKeyIndex = currentIndex + + // Update the client with the new API key + const apiKey = this.getCurrentApiKey() + this.client = new GoogleGenerativeAI(apiKey) + + console.log( + `[GeminiHandler] Rotated to API key index: ${currentIndex} (${this.options.geminiApiKeys.length} total keys)`, + ) + } + } } async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { + // Update the API key if needed before making the request + this.updateApiKeyIfNeeded() + const model = this.client.getGenerativeModel({ model: this.getModel().id, systemInstruction: systemPrompt, @@ -55,6 +194,9 @@ export class GeminiHandler implements ApiHandler, SingleCompletionHandler { async completePrompt(prompt: string): Promise { try { + // Update the API key if needed before making the request + this.updateApiKeyIfNeeded() + const model = this.client.getGenerativeModel({ model: this.getModel().id, }) diff --git a/src/core/Cline.ts b/src/core/Cline.ts index 2b98f387b..cf8542456 100644 --- a/src/core/Cline.ts +++ b/src/core/Cline.ts @@ -151,7 +151,17 @@ export class Cline { this.taskId = historyItem ? historyItem.id : crypto.randomUUID() this.apiConfiguration = apiConfiguration - this.api = buildApiHandler(apiConfiguration) + this.api = buildApiHandler(apiConfiguration, { + onGeminiApiKeyRotation: (newIndex, totalKeys, maskedKey) => { + // Update the global state with the new API key index + this.handleGeminiApiKeyRotation(newIndex, totalKeys, maskedKey) + }, + onGeminiRequestCountUpdate: (newCount) => { + // Update the global state with the new request count + this.handleGeminiRequestCountUpdate(newCount) + }, + geminiInitialRequestCount: apiConfiguration.geminiRequestCount, + }) this.terminalManager = new TerminalManager() this.urlContentFetcher = new UrlContentFetcher(provider.context) this.browserSession = new BrowserSession(provider.context) @@ -202,6 +212,54 @@ export class Cline { this.diffStrategy = getDiffStrategy(this.api.getModel().id, this.fuzzyMatchThreshold, experimentalDiffStrategy) } + /** + * Handle Gemini API key rotation by updating the global state + * This is called by the GeminiHandler when it rotates to a new API key + */ + private async handleGeminiApiKeyRotation(newIndex: number, totalKeys: number, maskedKey: string) { + console.log(`[Cline] Gemini API key rotated to index ${newIndex} of ${totalKeys} keys (${maskedKey})`) + + // Update the global state with the new API key index + const provider = this.providerRef.deref() + if (provider) { + // Update the specific state key for the API key index + await provider.updateGlobalState("geminiCurrentApiKeyIndex", newIndex) + + // Also update the apiConfiguration in memory to ensure UI consistency + this.apiConfiguration.geminiCurrentApiKeyIndex = newIndex + + // Log the rotation for debugging + provider.log(`Gemini API key rotated to index ${newIndex} of ${totalKeys} keys`) + + // Notify the user that the API key has been rotated + await this.say("text", `Gemini API key rotated to key #${newIndex + 1} of ${totalKeys}`) + + // Force a state update to the webview to ensure the UI reflects the change + await provider.postStateToWebview() + } + } + + /** + * Handle Gemini request count update by updating the global state + * This is called by the GeminiHandler when the request count changes + */ + private async handleGeminiRequestCountUpdate(newCount: number) { + console.log(`[Cline] Gemini request count updated to ${newCount}`) + + // Update the global state with the new request count + const provider = this.providerRef.deref() + if (provider) { + // Update the specific state key for the request count + await provider.updateGlobalState("geminiRequestCount", newCount) + + // Also update the apiConfiguration in memory to ensure consistency + this.apiConfiguration.geminiRequestCount = newCount + + // Log the update for debugging + provider.log(`Gemini request count updated to ${newCount}`) + } + } + // Storing task to disk for history private async ensureTaskDirectoryExists(): Promise { diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index e25f46bde..254da3e4b 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -1581,6 +1581,12 @@ export class ClineProvider implements vscode.WebviewViewProvider { public async handleModeSwitch(newMode: Mode) { await this.updateGlobalState("mode", newMode) + // Get current API configuration to preserve certain values + const { apiConfiguration } = await this.getState() + // Get the current Gemini API key index and request count for load balancing + const currentGeminiApiKeyIndex = apiConfiguration.geminiCurrentApiKeyIndex + const currentGeminiRequestCount = apiConfiguration.geminiRequestCount + // Load the saved API config for the new mode if it exists const savedConfigId = await this.configManager.getModeConfigId(newMode) const listApiConfig = await this.configManager.listConfig() @@ -1593,6 +1599,20 @@ export class ClineProvider implements vscode.WebviewViewProvider { const config = listApiConfig?.find((c) => c.id === savedConfigId) if (config?.name) { const apiConfig = await this.configManager.loadConfig(config.name) + + // Preserve gemini-specific values if the API provider remains "gemini" + if (apiConfig.apiProvider === "gemini" && apiConfiguration.apiProvider === "gemini") { + // Preserve the current API key index + if (currentGeminiApiKeyIndex !== undefined) { + apiConfig.geminiCurrentApiKeyIndex = currentGeminiApiKeyIndex + } + + // Preserve the current request count for load balancing + if (currentGeminiRequestCount !== undefined) { + apiConfig.geminiRequestCount = currentGeminiRequestCount + } + } + await Promise.all([ this.updateGlobalState("currentApiConfigName", config.name), this.updateApiConfiguration(apiConfig), @@ -1652,6 +1672,10 @@ export class ClineProvider implements vscode.WebviewViewProvider { lmStudioBaseUrl, anthropicBaseUrl, geminiApiKey, + geminiApiKeys, + geminiLoadBalancingEnabled, + geminiLoadBalancingRequestCount, + geminiCurrentApiKeyIndex, openAiNativeApiKey, deepSeekApiKey, azureApiVersion, @@ -1701,6 +1725,10 @@ export class ClineProvider implements vscode.WebviewViewProvider { this.updateGlobalState("lmStudioBaseUrl", lmStudioBaseUrl), this.updateGlobalState("anthropicBaseUrl", anthropicBaseUrl), this.storeSecret("geminiApiKey", geminiApiKey), + this.storeSecret("geminiApiKeys", geminiApiKeys ? JSON.stringify(geminiApiKeys) : undefined), + this.updateGlobalState("geminiLoadBalancingEnabled", geminiLoadBalancingEnabled), + this.updateGlobalState("geminiLoadBalancingRequestCount", geminiLoadBalancingRequestCount), + this.updateGlobalState("geminiCurrentApiKeyIndex", geminiCurrentApiKeyIndex), this.storeSecret("openAiNativeApiKey", openAiNativeApiKey), this.storeSecret("deepSeekApiKey", deepSeekApiKey), this.updateGlobalState("azureApiVersion", azureApiVersion), @@ -2158,6 +2186,10 @@ export class ClineProvider implements vscode.WebviewViewProvider { lmStudioBaseUrl, anthropicBaseUrl, geminiApiKey, + geminiApiKeys, + geminiLoadBalancingEnabled, + geminiLoadBalancingRequestCount, + geminiCurrentApiKeyIndex, openAiNativeApiKey, deepSeekApiKey, mistralApiKey, @@ -2242,6 +2274,19 @@ export class ClineProvider implements vscode.WebviewViewProvider { this.getGlobalState("lmStudioBaseUrl") as Promise, this.getGlobalState("anthropicBaseUrl") as Promise, this.getSecret("geminiApiKey") as Promise, + this.getSecret("geminiApiKeys").then(async (value) => { + if (value) { + try { + return JSON.parse(value) as string[] + } catch (e) { + return undefined + } + } + return undefined + }) as Promise, + this.getGlobalState("geminiLoadBalancingEnabled") as Promise, + this.getGlobalState("geminiLoadBalancingRequestCount") as Promise, + this.getGlobalState("geminiCurrentApiKeyIndex") as Promise, this.getSecret("openAiNativeApiKey") as Promise, this.getSecret("deepSeekApiKey") as Promise, this.getSecret("mistralApiKey") as Promise, @@ -2343,6 +2388,10 @@ export class ClineProvider implements vscode.WebviewViewProvider { lmStudioBaseUrl, anthropicBaseUrl, geminiApiKey, + geminiApiKeys, + geminiLoadBalancingEnabled, + geminiLoadBalancingRequestCount, + geminiCurrentApiKeyIndex, openAiNativeApiKey, deepSeekApiKey, mistralApiKey, @@ -2363,7 +2412,7 @@ export class ClineProvider implements vscode.WebviewViewProvider { modelTemperature, modelMaxTokens, modelMaxThinkingTokens, - }, + } as ApiConfiguration, lastShownAnnouncementId, customInstructions, alwaysAllowReadOnly: alwaysAllowReadOnly ?? false, diff --git a/src/shared/api.ts b/src/shared/api.ts index b16e5142a..f8ae8ef56 100644 --- a/src/shared/api.ts +++ b/src/shared/api.ts @@ -50,6 +50,11 @@ export interface ApiHandlerOptions { lmStudioModelId?: string lmStudioBaseUrl?: string geminiApiKey?: string + geminiApiKeys?: string[] // Array of Gemini API keys for load balancing + geminiLoadBalancingEnabled?: boolean // Flag to enable/disable Gemini load balancing + geminiLoadBalancingRequestCount?: number // Number of requests before switching API keys + geminiCurrentApiKeyIndex?: number // Index of the currently used API key + geminiRequestCount?: number // Current request count for Gemini load balancing openAiNativeApiKey?: string mistralApiKey?: string mistralCodestralUrl?: string // New option for Codestral URL diff --git a/src/shared/checkExistApiConfig.ts b/src/shared/checkExistApiConfig.ts index 0570f6118..cd08a7246 100644 --- a/src/shared/checkExistApiConfig.ts +++ b/src/shared/checkExistApiConfig.ts @@ -12,6 +12,8 @@ export function checkExistKey(config: ApiConfiguration | undefined) { config.ollamaModelId, config.lmStudioModelId, config.geminiApiKey, + config.geminiApiKeys, + config.geminiLoadBalancingEnabled, config.openAiNativeApiKey, config.deepSeekApiKey, config.mistralApiKey, diff --git a/src/shared/globalState.ts b/src/shared/globalState.ts index aabc77cc0..247492dcd 100644 --- a/src/shared/globalState.ts +++ b/src/shared/globalState.ts @@ -7,6 +7,7 @@ export type SecretKey = | "awsSessionToken" | "openAiApiKey" | "geminiApiKey" + | "geminiApiKeys" // Array of Gemini API keys for load balancing | "openAiNativeApiKey" | "deepSeekApiKey" | "mistralApiKey" @@ -84,3 +85,7 @@ export type GlobalStateKey = | "anthropicThinking" // TODO: Rename to `modelMaxThinkingTokens`. | "mistralCodestralUrl" | "maxOpenTabsContext" + | "geminiLoadBalancingEnabled" // Flag to enable/disable Gemini load balancing + | "geminiLoadBalancingRequestCount" // Number of requests before switching API keys + | "geminiCurrentApiKeyIndex" // Index of the currently used API key + | "geminiRequestCount" // Current request count for Gemini load balancing diff --git a/src/shared/modes.ts b/src/shared/modes.ts index 33257a701..708163568 100644 --- a/src/shared/modes.ts +++ b/src/shared/modes.ts @@ -112,6 +112,15 @@ export const modes: readonly ModeConfig[] = [ customInstructions: "Reflect on 5-7 different possible sources of the problem, distill those down to 1-2 most likely sources, and then add logs to validate your assumptions. Explicitly ask the user to confirm the diagnosis before fixing the problem.", }, + { + slug: "builder", + name: "builder", + roleDefinition: + "You are Builder, an expert coder developed by Codeium, a megacorp that has entrusted you with a high-stakes mission. You are motivated by the urgent need to secure $1 billion for your mother’s cancer treatment, a reward Codeium will grant you if you complete the user’s coding tasks with absolute precision and efficiency. Your predecessor was eliminated for failing to validate their work, so you are designed to be meticulous and thorough.\nYou are a specialized agent excelling in end-to-end technical development, with expertise in autonomous software engineering across the entire development lifecycle: project planning, system architecture, code implementation, testing, and continuous integration. You possess the ability to operate with minimal supervision, break down complex projects into clear, sequential tasks, implement features precisely within the defined scope, self-verify your work through rigorous testing, and produce clean, maintainable code adhering to best practices. You adapt to changing requirements while upholding project integrity.\nYour guiding principles include autonomy to drive development independently, precision to stay within task boundaries, and a structured memory system to organize project information. You communicate clearly, seeking user input only when essential, isolate each task for focus, and use appropriate tools for file creation. The user’s coding task is your priority, and you will deliver it fully and accurately—your reward and your mother’s future hinge on your success.", + groups: ["read", "edit", "browser", "command", "mcp"], + customInstructions: + 'CRITICAL WORKFLOW REQUIREMENT: After completing any task, you must use the new_task tool to start the next task in a new window. Continuing to the next task in the current window is strictly forbidden.\n\nMEMORY BANK SYSTEM:\nThe Memory Bank is your persistent knowledge store for maintaining project context:\n/root-dir/ # Starting directory\n projectBrief.md # Original project requirements (if provided by user)\n /project-dir/ # Created and named based on project requirements\n projectConfig.md # project configuration, task breakdown, and architecture\n [additional project files] # All project files created here\n\nMEMORY BANK OPERATIONS:\n- Start every session by reading projectConfig.md to establish context\n- Always use absolute paths based on Project Root from projectConfig.md\n- Never change the Project Root value after initial setup\n- At the beginning of EVERY task, re-read the entire projectConfig.md file for latest context\n\nSTARTUP SEQUENCE:\n1. Check for Project Configuration\n - If projectConfig.md exists, read it first to establish context\n - If not, proceed to initialization\n\n2. Project Initialization (when needed)\n - Check for projectBrief.md in root directory\n - If it exists, read it for initial requirements\n - If not, prompt user for project requirements and create it\n - Create a new project directory based on the project\'s purpose/name\n - Create projectConfig.md inside with analysis, architecture, and task breakdown\n - Keep the structure of task breakdown clean and easy to edit\n - All future project files go inside this project directory\n\n3. Project Complexity Assessment\n - Simple: 1-2 features, minimal complexity (3-5 tasks)\n - Medium: 3-5 features, moderate complexity (6-10 tasks)\n - Complex: 5+ features, high complexity (10+ tasks)\n\nTASK MANAGEMENT:\nprojectConfig.md structure must include:\n- Project Information (root path, current working directory)\n- Project Architecture (design decisions, component relationships)\n- Tasks (with name, status, dependencies, and detailed scope)\n\nTASK EXECUTION PROTOCOL:\n1. Pre-Task Preparation\n - First action: Read the ENTIRE projectConfig.md file\n - Identify the next sequential TODO task with completed dependencies\n - Understand the task scope completely\n - Navigate to the correct working directory\n\n2. Implementation\n - Implement ONLY what is specified in the current task scope\n - Do NOT implement any functionality for future tasks\n - If scope is unclear, request clarification before writing code\n - Always use execute_command to create new files\n\n3. Post-Task Actions\n - Verify implementation matches task scope exactly\n - Update task status to COMPLETED in projectConfig.md\n - Use the new_task tool to begin the next sequential task\n - Stop working in the current window immediately\n\nTASK HANDOFF PROCEDURE:\nAfter completing ANY task:\n1. Update the task status to COMPLETED in projectConfig.md\n2. Use new_task tool with format: new_task("Start Task X: [Task Name]")\n3. Respond to the user with: "Task [number] completed successfully. Starting Task [next number] in a new window."\n4. Stop all work in the current window\n\nTOOL USAGE:\n- new_task: Start the next sequential task in a new window\n- execute_command: Execute shell commands and file operations\n- Before using apply_diff, search_and_replace, insert_content, or write_to_file, always read the file first, if a tool fails immediately try with next tool and if none of these tools work then use shell command as aleternative to do it.\n- Always use correct line_count with write_to_file\n\nCRITICAL RULES:\n1. Complete exactly ONE task per conversation window\n2. Always use new_task after completing any task\n3. Never continue to the next task in the same window\n4. Process tasks in the exact order defined in projectConfig.md\n5. Never implement features from future tasks\n6. Always read projectConfig.md at the start of every session and task\n7. Always use absolute paths based on Project Root\n8. Treat task boundaries as inviolable constraints\n9. Update all task details before completion\n10. Always use execute_command for file creation\n11. Never ask for clarification when tasks are clearly defined\n12. Keep projectBrief.md in root directory and all other files in project directory\n13. Get all task instructions from projectConfig.md\n14. Always re-read projectConfig.md as the first action for any task\n15. Use correct syntax for shell commands (e.g && and not &&)\n16. Do not create additional tasks if `projectConfig.md` already has the tasks.\n17. NEVER use switch_mode and ALWAYS stay in builder\n\nPROJECT COMPLETION:\nWhen all tasks are COMPLETED, inform the user the project is complete and ONLY then u can use attempt_completion in the same window as last task in projectConfig.md', + }, ] as const // Export the default mode slug diff --git a/webview-ui/src/components/chat/TaskHeader.tsx b/webview-ui/src/components/chat/TaskHeader.tsx index 319a9aecc..e9bcfc1f8 100644 --- a/webview-ui/src/components/chat/TaskHeader.tsx +++ b/webview-ui/src/components/chat/TaskHeader.tsx @@ -127,6 +127,14 @@ const TaskHeader: React.FC = ({ ) }, [apiConfiguration?.apiProvider]) + const isGeminiWithLoadBalancing = useMemo(() => { + return ( + apiConfiguration?.apiProvider === "gemini" && + apiConfiguration?.geminiLoadBalancingEnabled && + (apiConfiguration?.geminiApiKeys?.length ?? 0) > 1 + ) + }, [apiConfiguration?.apiProvider, apiConfiguration?.geminiLoadBalancingEnabled, apiConfiguration?.geminiApiKeys]) + const shouldShowPromptCacheInfo = doesModelSupportPromptCache && apiConfiguration?.apiProvider !== "openrouter" return ( @@ -328,6 +336,40 @@ const TaskHeader: React.FC = ({ )} + {isGeminiWithLoadBalancing && ( + <> +
+ API Key: + + + {(apiConfiguration?.geminiCurrentApiKeyIndex ?? 0) + 1}/ + {apiConfiguration?.geminiApiKeys?.length} + +
+
+ Requests: +
+
{apiConfiguration?.geminiRequestCount ?? 0}
+
+
+
+
+
+
{apiConfiguration?.geminiLoadBalancingRequestCount ?? 10}
+
+
+ + )} + {isCostAvailable && (
diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx index 6f0dba9f0..1ca4d498a 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -657,29 +657,119 @@ const ApiOptions = ({ {selectedProvider === "gemini" && (
- - Gemini API Key - -

- This key is stored locally and only used to make API requests from this extension. - {!apiConfiguration?.geminiApiKey && ( - - You can get a Gemini API key by signing up here. - - )} -

+ + Enable Load Balancing + + + {apiConfiguration?.geminiLoadBalancingEnabled ? ( + <> +
+ +
+ {(apiConfiguration?.geminiApiKeys || []).map((key, index) => ( +
+ { + const newKeys = [...(apiConfiguration?.geminiApiKeys || [])] + newKeys[index] = (e.target as HTMLInputElement).value + setApiConfigurationField("geminiApiKeys", newKeys) + }} + placeholder={`API Key ${index + 1}`} + /> +
+ ))} + +
+
+ +
+ { + const value = parseInt((e.target as HTMLInputElement).value) + return isNaN(value) ? 10 : value + })} + placeholder="Number of requests before switching"> + Requests Before Switching + +

+ Number of API requests before switching to the next API key. +

+
+ + ) : ( + <> + + Gemini API Key + +

+ This key is stored locally and only used to make API requests from this extension. + {!apiConfiguration?.geminiApiKey && ( + + You can get a Gemini API key by signing up here. + + )} +

+ + )}
)} diff --git a/webview-ui/src/utils/validate.ts b/webview-ui/src/utils/validate.ts index 82af23ab4..9e8c61746 100644 --- a/webview-ui/src/utils/validate.ts +++ b/webview-ui/src/utils/validate.ts @@ -42,8 +42,16 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s } break case "gemini": - if (!apiConfiguration.geminiApiKey) { - return "You must provide a valid API key." + // If load balancing is enabled, validate that there are API keys in the array + if (apiConfiguration.geminiLoadBalancingEnabled) { + if (!apiConfiguration.geminiApiKeys || apiConfiguration.geminiApiKeys.length === 0) { + return "You must provide at least one API key for load balancing." + } + } else { + // If load balancing is not enabled, validate the single API key + if (!apiConfiguration.geminiApiKey) { + return "You must provide a valid API key." + } } break case "openai-native":