Skip to content

Commit 2c2547a

Browse files
shariqriazzshariqriazz
committed
loadbalance
1 parent a2d441c commit 2c2547a

File tree

11 files changed

+471
-48
lines changed

11 files changed

+471
-48
lines changed

‎src/api/index.ts‎

Lines changed: 32 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { VertexHandler } from "./providers/vertex"
1111
import { OpenAiHandler } from "./providers/openai"
1212
import { OllamaHandler } from "./providers/ollama"
1313
import { LmStudioHandler } from "./providers/lmstudio"
14-
import { GeminiHandler } from "./providers/gemini"
14+
import { GeminiHandler, ApiKeyRotationCallback, RequestCountUpdateCallback } from "./providers/gemini"
1515
import { OpenAiNativeHandler } from "./providers/openai-native"
1616
import { DeepSeekHandler } from "./providers/deepseek"
1717
import { MistralHandler } from "./providers/mistral"
@@ -29,41 +29,54 @@ export interface ApiHandler {
2929
getModel(): { id: string; info: ModelInfo }
3030
}
3131

32-
export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
33-
const { apiProvider, ...options } = configuration
32+
/**
33+
* Callbacks that can be passed to API handlers
34+
*/
35+
export interface ApiHandlerCallbacks {
36+
onGeminiApiKeyRotation?: ApiKeyRotationCallback
37+
onGeminiRequestCountUpdate?: RequestCountUpdateCallback
38+
geminiInitialRequestCount?: number
39+
}
40+
41+
export function buildApiHandler(configuration: ApiConfiguration, callbacks?: ApiHandlerCallbacks): ApiHandler {
42+
const { apiProvider, ...handlerOptions } = configuration
3443
switch (apiProvider) {
3544
case "anthropic":
36-
return new AnthropicHandler(options)
45+
return new AnthropicHandler(handlerOptions)
3746
case "glama":
38-
return new GlamaHandler(options)
47+
return new GlamaHandler(handlerOptions)
3948
case "openrouter":
40-
return new OpenRouterHandler(options)
49+
return new OpenRouterHandler(handlerOptions)
4150
case "bedrock":
42-
return new AwsBedrockHandler(options)
51+
return new AwsBedrockHandler(handlerOptions)
4352
case "vertex":
44-
return new VertexHandler(options)
53+
return new VertexHandler(handlerOptions)
4554
case "openai":
46-
return new OpenAiHandler(options)
55+
return new OpenAiHandler(handlerOptions)
4756
case "ollama":
48-
return new OllamaHandler(options)
57+
return new OllamaHandler(handlerOptions)
4958
case "lmstudio":
50-
return new LmStudioHandler(options)
59+
return new LmStudioHandler(handlerOptions)
5160
case "gemini":
52-
return new GeminiHandler(options)
61+
return new GeminiHandler(handlerOptions, {
62+
onApiKeyRotation: callbacks?.onGeminiApiKeyRotation,
63+
onRequestCountUpdate: callbacks?.onGeminiRequestCountUpdate,
64+
initialRequestCount: callbacks?.geminiInitialRequestCount,
65+
})
5366
case "openai-native":
54-
return new OpenAiNativeHandler(options)
67+
return new OpenAiNativeHandler(handlerOptions)
5568
case "deepseek":
56-
return new DeepSeekHandler(options)
69+
return new DeepSeekHandler(handlerOptions)
5770
case "vscode-lm":
58-
return new VsCodeLmHandler(options)
71+
return new VsCodeLmHandler(handlerOptions)
5972
case "mistral":
60-
return new MistralHandler(options)
73+
return new MistralHandler(handlerOptions)
6174
case "unbound":
62-
return new UnboundHandler(options)
75+
return new UnboundHandler(handlerOptions)
6376
case "requesty":
64-
return new RequestyHandler(options)
77+
return new RequestyHandler(handlerOptions)
6578
default:
66-
return new AnthropicHandler(options)
79+
return new AnthropicHandler(handlerOptions)
6780
}
6881
}
6982

‎src/api/providers/gemini.ts‎

Lines changed: 144 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,156 @@ import { convertAnthropicMessageToGemini } from "../transform/gemini-format"
66
import { ApiStream } from "../transform/stream"
77

88
const GEMINI_DEFAULT_TEMPERATURE = 0
9+
const DEFAULT_REQUEST_COUNT = 10 // Default number of requests before switching API keys
10+
11+
// Define a callback type for API key rotation
12+
export type ApiKeyRotationCallback = (newIndex: number, totalKeys: number, apiKey: string) => void
13+
export type RequestCountUpdateCallback = (newCount: number) => void
914

1015
export class GeminiHandler implements ApiHandler, SingleCompletionHandler {
1116
private options: ApiHandlerOptions
1217
private client: GoogleGenerativeAI
18+
private requestCount: number = 0
19+
private onApiKeyRotation?: ApiKeyRotationCallback
20+
private onRequestCountUpdate?: RequestCountUpdateCallback
1321

14-
constructor(options: ApiHandlerOptions) {
22+
constructor(
23+
options: ApiHandlerOptions,
24+
callbacks?: {
25+
onApiKeyRotation?: ApiKeyRotationCallback
26+
onRequestCountUpdate?: RequestCountUpdateCallback
27+
initialRequestCount?: number
28+
},
29+
) {
1530
this.options = options
16-
this.client = new GoogleGenerativeAI(options.geminiApiKey ?? "not-provided")
31+
this.onApiKeyRotation = callbacks?.onApiKeyRotation
32+
this.onRequestCountUpdate = callbacks?.onRequestCountUpdate
33+
34+
// Initialize request count from saved state if provided
35+
if (callbacks?.initialRequestCount !== undefined) {
36+
this.requestCount = callbacks.initialRequestCount
37+
console.log(`[GeminiHandler] Initialized with request count: ${this.requestCount}`)
38+
}
39+
40+
// Initialize with the current API key
41+
const apiKey = this.getCurrentApiKey()
42+
this.client = new GoogleGenerativeAI(apiKey)
43+
44+
// Log initial API key setup if load balancing is enabled
45+
if (
46+
this.options.geminiLoadBalancingEnabled &&
47+
this.options.geminiApiKeys &&
48+
this.options.geminiApiKeys.length > 0
49+
) {
50+
console.log(
51+
`[GeminiHandler] Load balancing enabled with ${this.options.geminiApiKeys.length} keys. Current index: ${this.options.geminiCurrentApiKeyIndex ?? 0}`,
52+
)
53+
}
54+
}
55+
56+
/**
57+
* Get the current API key based on load balancing settings
58+
*/
59+
private getCurrentApiKey(): string {
60+
// If load balancing is not enabled or there are no multiple API keys, use the single API key
61+
if (
62+
!this.options.geminiLoadBalancingEnabled ||
63+
!this.options.geminiApiKeys ||
64+
this.options.geminiApiKeys.length === 0
65+
) {
66+
return this.options.geminiApiKey ?? "not-provided"
67+
}
68+
69+
// Get the current API key index, defaulting to 0 if not set
70+
const currentIndex = this.options.geminiCurrentApiKeyIndex ?? 0
71+
72+
// Return the API key at the current index
73+
return this.options.geminiApiKeys[currentIndex] ?? "not-provided"
74+
}
75+
76+
/**
77+
* Update the client with the next API key if load balancing is enabled
78+
*/
79+
private updateApiKeyIfNeeded(): void {
80+
// If load balancing is not enabled or there are no multiple API keys, do nothing
81+
if (
82+
!this.options.geminiLoadBalancingEnabled ||
83+
!this.options.geminiApiKeys ||
84+
this.options.geminiApiKeys.length <= 1
85+
) {
86+
return
87+
}
88+
89+
// Increment the request count
90+
this.requestCount++
91+
console.log(
92+
`[GeminiHandler] Request count: ${this.requestCount}/${this.options.geminiLoadBalancingRequestCount ?? DEFAULT_REQUEST_COUNT}`,
93+
)
94+
95+
// Notify about request count update
96+
if (this.onRequestCountUpdate) {
97+
this.onRequestCountUpdate(this.requestCount)
98+
}
99+
100+
// Get the request count threshold, defaulting to DEFAULT_REQUEST_COUNT if not set
101+
const requestCountThreshold = this.options.geminiLoadBalancingRequestCount ?? DEFAULT_REQUEST_COUNT
102+
103+
// If the request count has reached the threshold, switch to the next API key
104+
if (this.requestCount >= requestCountThreshold) {
105+
// Reset the request count
106+
this.requestCount = 0
107+
108+
// Notify about request count reset
109+
if (this.onRequestCountUpdate) {
110+
this.onRequestCountUpdate(0)
111+
}
112+
113+
// Get the current API key index, defaulting to 0 if not set
114+
let currentIndex = this.options.geminiCurrentApiKeyIndex ?? 0
115+
116+
// Calculate the next index, wrapping around if necessary
117+
currentIndex = (currentIndex + 1) % this.options.geminiApiKeys.length
118+
119+
// Notify callback first to update global state
120+
if (this.onApiKeyRotation) {
121+
// Get the API key for the new index
122+
const apiKey = this.options.geminiApiKeys[currentIndex] ?? "not-provided"
123+
124+
// Only send the first few characters of the API key for security
125+
const maskedKey = apiKey.substring(0, 4) + "..." + apiKey.substring(apiKey.length - 4)
126+
127+
// Call the callback to update global state
128+
this.onApiKeyRotation(currentIndex, this.options.geminiApiKeys.length, maskedKey)
129+
130+
// Update the current index in the options AFTER the callback
131+
// This ensures we're using the index that was just set in global state
132+
this.options.geminiCurrentApiKeyIndex = currentIndex
133+
134+
// Update the client with the new API key
135+
this.client = new GoogleGenerativeAI(apiKey)
136+
137+
console.log(
138+
`[GeminiHandler] Rotated to API key index: ${currentIndex} (${this.options.geminiApiKeys.length} total keys)`,
139+
)
140+
} else {
141+
// No callback provided, just update locally
142+
this.options.geminiCurrentApiKeyIndex = currentIndex
143+
144+
// Update the client with the new API key
145+
const apiKey = this.getCurrentApiKey()
146+
this.client = new GoogleGenerativeAI(apiKey)
147+
148+
console.log(
149+
`[GeminiHandler] Rotated to API key index: ${currentIndex} (${this.options.geminiApiKeys.length} total keys)`,
150+
)
151+
}
152+
}
17153
}
18154

19155
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
156+
// Update the API key if needed before making the request
157+
this.updateApiKeyIfNeeded()
158+
20159
const model = this.client.getGenerativeModel({
21160
model: this.getModel().id,
22161
systemInstruction: systemPrompt,
@@ -55,6 +194,9 @@ export class GeminiHandler implements ApiHandler, SingleCompletionHandler {
55194

56195
async completePrompt(prompt: string): Promise<string> {
57196
try {
197+
// Update the API key if needed before making the request
198+
this.updateApiKeyIfNeeded()
199+
58200
const model = this.client.getGenerativeModel({
59201
model: this.getModel().id,
60202
})

‎src/core/Cline.ts‎

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,17 @@ export class Cline {
151151
this.taskId = historyItem ? historyItem.id : crypto.randomUUID()
152152

153153
this.apiConfiguration = apiConfiguration
154-
this.api = buildApiHandler(apiConfiguration)
154+
this.api = buildApiHandler(apiConfiguration, {
155+
onGeminiApiKeyRotation: (newIndex, totalKeys, maskedKey) => {
156+
// Update the global state with the new API key index
157+
this.handleGeminiApiKeyRotation(newIndex, totalKeys, maskedKey)
158+
},
159+
onGeminiRequestCountUpdate: (newCount) => {
160+
// Update the global state with the new request count
161+
this.handleGeminiRequestCountUpdate(newCount)
162+
},
163+
geminiInitialRequestCount: apiConfiguration.geminiRequestCount,
164+
})
155165
this.terminalManager = new TerminalManager()
156166
this.urlContentFetcher = new UrlContentFetcher(provider.context)
157167
this.browserSession = new BrowserSession(provider.context)
@@ -202,6 +212,54 @@ export class Cline {
202212
this.diffStrategy = getDiffStrategy(this.api.getModel().id, this.fuzzyMatchThreshold, experimentalDiffStrategy)
203213
}
204214

215+
/**
216+
* Handle Gemini API key rotation by updating the global state
217+
* This is called by the GeminiHandler when it rotates to a new API key
218+
*/
219+
private async handleGeminiApiKeyRotation(newIndex: number, totalKeys: number, maskedKey: string) {
220+
console.log(`[Cline] Gemini API key rotated to index ${newIndex} of ${totalKeys} keys (${maskedKey})`)
221+
222+
// Update the global state with the new API key index
223+
const provider = this.providerRef.deref()
224+
if (provider) {
225+
// Update the specific state key for the API key index
226+
await provider.updateGlobalState("geminiCurrentApiKeyIndex", newIndex)
227+
228+
// Also update the apiConfiguration in memory to ensure UI consistency
229+
this.apiConfiguration.geminiCurrentApiKeyIndex = newIndex
230+
231+
// Log the rotation for debugging
232+
provider.log(`Gemini API key rotated to index ${newIndex} of ${totalKeys} keys`)
233+
234+
// Notify the user that the API key has been rotated
235+
await this.say("text", `Gemini API key rotated to key #${newIndex + 1} of ${totalKeys}`)
236+
237+
// Force a state update to the webview to ensure the UI reflects the change
238+
await provider.postStateToWebview()
239+
}
240+
}
241+
242+
/**
243+
* Handle Gemini request count update by updating the global state
244+
* This is called by the GeminiHandler when the request count changes
245+
*/
246+
private async handleGeminiRequestCountUpdate(newCount: number) {
247+
console.log(`[Cline] Gemini request count updated to ${newCount}`)
248+
249+
// Update the global state with the new request count
250+
const provider = this.providerRef.deref()
251+
if (provider) {
252+
// Update the specific state key for the request count
253+
await provider.updateGlobalState("geminiRequestCount", newCount)
254+
255+
// Also update the apiConfiguration in memory to ensure consistency
256+
this.apiConfiguration.geminiRequestCount = newCount
257+
258+
// Log the update for debugging
259+
provider.log(`Gemini request count updated to ${newCount}`)
260+
}
261+
}
262+
205263
// Storing task to disk for history
206264

207265
private async ensureTaskDirectoryExists(): Promise<string> {

0 commit comments

Comments
 (0)