Skip to content

Commit 8d952cb

Browse files
jddunnjddunn
committed
feat: add key pool rotation to all remaining providers (20 files)
Deepgram (streaming + batch STT), AssemblyAI STT, Gemini LLM, Stability/Fal/Replicate/Flux image, Fal/Replicate/Runway video, Suno/Udio/Fal/StableAudio/Replicate audio, Serper/Tavily/Brave/Firecrawl search. Groq/Mistral/Together/XAI delegate to OpenAIProvider (already pooled).
1 parent 2b4304b commit 8d952cb

20 files changed

Lines changed: 116 additions & 49 deletions

‎src/core/llm/providers/implementations/GeminiProvider.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import {
3434
ProviderEmbeddingResponse,
3535
} from '../IProvider';
3636
import { GeminiProviderError } from '../errors/GeminiProviderError';
37+
import { ApiKeyPool } from '../../../providers/ApiKeyPool.js';
3738

3839
// ---------------------------------------------------------------------------
3940
// Configuration
@@ -259,6 +260,7 @@ export class GeminiProvider implements IProvider {
259260
public defaultModelId?: string;
260261

261262
private config!: GeminiProviderConfig;
263+
private keyPool: ApiKeyPool | null = null;
262264

263265
constructor() {}
264266

@@ -292,6 +294,7 @@ export class GeminiProvider implements IProvider {
292294
defaultModelId: 'gemini-2.5-flash',
293295
...config,
294296
};
297+
this.keyPool = new ApiKeyPool(config.apiKey);
295298
this.defaultModelId = this.config.defaultModelId;
296299
this.isInitialized = true;
297300

@@ -1082,7 +1085,7 @@ export class GeminiProvider implements IProvider {
10821085
body: Record<string, unknown>,
10831086
): Promise<T> {
10841087
// API key is passed as query parameter — Gemini's auth convention
1085-
const url = `${this.config.baseURL}${endpoint}?key=${this.config.apiKey}`;
1088+
const url = `${this.config.baseURL}${endpoint}?key=${this.keyPool?.hasKeys ? this.keyPool.next() : this.config.apiKey}`;
10861089
const headers: Record<string, string> = {
10871090
'Content-Type': 'application/json',
10881091
'User-Agent': 'AgentOS/1.0 (GeminiProvider)',
@@ -1203,7 +1206,7 @@ export class GeminiProvider implements IProvider {
12031206
body: Record<string, unknown>,
12041207
): Promise<ReadableStream<Uint8Array>> {
12051208
// Both alt=sse and key= are query params
1206-
const url = `${this.config.baseURL}${endpoint}?alt=sse&key=${this.config.apiKey}`;
1209+
const url = `${this.config.baseURL}${endpoint}?alt=sse&key=${this.keyPool?.hasKeys ? this.keyPool.next() : this.config.apiKey}`;
12071210
const headers: Record<string, string> = {
12081211
'Content-Type': 'application/json',
12091212
'User-Agent': 'AgentOS/1.0 (GeminiProvider)',

‎src/hearing/providers/AssemblyAISTTProvider.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import type {
55
SpeechTranscriptionResult,
66
SpeechTranscriptionSegment,
77
} from '../../speech/types.js';
8+
import { ApiKeyPool } from '../../core/providers/ApiKeyPool.js';
89

910
/**
1011
* Configuration for the {@link AssemblyAISTTProvider}.
@@ -228,8 +229,11 @@ export class AssemblyAISTTProvider implements SpeechToTextProvider {
228229
* });
229230
* ```
230231
*/
232+
private readonly keyPool: ApiKeyPool;
233+
231234
constructor(private readonly config: AssemblyAISTTProviderConfig) {
232235
this.fetchImpl = config.fetchImpl ?? fetch;
236+
this.keyPool = new ApiKeyPool(config.apiKey);
233237
}
234238

235239
/**
@@ -290,7 +294,7 @@ export class AssemblyAISTTProvider implements SpeechToTextProvider {
290294
const uploadResponse = await this.fetchImpl(`${ASSEMBLYAI_BASE}/upload`, {
291295
method: 'POST',
292296
headers: {
293-
Authorization: this.config.apiKey,
297+
Authorization: this.keyPool.next(),
294298
'Content-Type': audio.mimeType ?? 'audio/wav',
295299
},
296300
body: audio.data as any,
@@ -316,7 +320,7 @@ export class AssemblyAISTTProvider implements SpeechToTextProvider {
316320
const submitResponse = await this.fetchImpl(`${ASSEMBLYAI_BASE}/transcript`, {
317321
method: 'POST',
318322
headers: {
319-
Authorization: this.config.apiKey,
323+
Authorization: this.keyPool.next(),
320324
'Content-Type': 'application/json',
321325
},
322326
body: JSON.stringify(submitPayload),
@@ -349,7 +353,7 @@ export class AssemblyAISTTProvider implements SpeechToTextProvider {
349353
}
350354

351355
const pollResponse = await this.fetchImpl(`${ASSEMBLYAI_BASE}/transcript/${id}`, {
352-
headers: { Authorization: this.config.apiKey },
356+
headers: { Authorization: this.keyPool.next() },
353357
signal,
354358
});
355359

‎src/hearing/providers/DeepgramBatchSTTProvider.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import type {
55
SpeechTranscriptionResult,
66
SpeechTranscriptionSegment,
77
} from '../../speech/types.js';
8+
import { ApiKeyPool } from '../../core/providers/ApiKeyPool.js';
89

910
/**
1011
* Configuration for the {@link DeepgramBatchSTTProvider}.
@@ -204,8 +205,11 @@ export class DeepgramBatchSTTProvider implements SpeechToTextProvider {
204205
* });
205206
* ```
206207
*/
208+
private readonly keyPool: ApiKeyPool;
209+
207210
constructor(private readonly config: DeepgramBatchSTTProviderConfig) {
208211
this.fetchImpl = config.fetchImpl ?? fetch;
212+
this.keyPool = new ApiKeyPool(config.apiKey);
209213
}
210214

211215
/**
@@ -271,7 +275,7 @@ export class DeepgramBatchSTTProvider implements SpeechToTextProvider {
271275
const response = await this.fetchImpl(url, {
272276
method: 'POST',
273277
headers: {
274-
Authorization: `Token ${this.config.apiKey}`,
278+
Authorization: `Token ${this.keyPool.next()}`,
275279
'Content-Type': contentType,
276280
},
277281
// Cast needed because SpeechAudioInput.data is typed as Buffer but

‎src/media/audio/providers/FalAudioProvider.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
*/
3434

3535
import type { IAudioGenerator } from '../IAudioGenerator.js';
36+
import { ApiKeyPool } from '../../../core/providers/ApiKeyPool.js';
3637
import type { MusicGenerateRequest, SFXGenerateRequest, AudioResult } from '../types.js';
3738

3839
// ---------------------------------------------------------------------------
@@ -176,6 +177,7 @@ export class FalAudioProvider implements IAudioGenerator {
176177

177178
/** Internal resolved configuration. */
178179
private _config!: Required<Pick<FalAudioProviderConfig, 'apiKey' | 'baseURL' | 'pollIntervalMs' | 'timeoutMs'>> & FalAudioProviderConfig;
180+
private keyPool!: ApiKeyPool;
179181

180182
// -------------------------------------------------------------------------
181183
// Lifecycle
@@ -214,6 +216,7 @@ export class FalAudioProvider implements IAudioGenerator {
214216
};
215217

216218
this.defaultModelId = this._config.defaultModelId;
219+
this.keyPool = new ApiKeyPool(apiKey);
217220
this.isInitialized = true;
218221
}
219222

@@ -332,7 +335,7 @@ export class FalAudioProvider implements IAudioGenerator {
332335
const response = await fetch(url, {
333336
method: 'POST',
334337
headers: {
335-
Authorization: `Key ${this._config.apiKey}`,
338+
Authorization: `Key ${this.keyPool.next()}`,
336339
'Content-Type': 'application/json',
337340
},
338341
body: JSON.stringify(body),
@@ -368,7 +371,7 @@ export class FalAudioProvider implements IAudioGenerator {
368371

369372
const response = await fetch(url, {
370373
headers: {
371-
Authorization: `Key ${this._config.apiKey}`,
374+
Authorization: `Key ${this.keyPool.next()}`,
372375
},
373376
});
374377

@@ -413,7 +416,7 @@ export class FalAudioProvider implements IAudioGenerator {
413416

414417
const response = await fetch(url, {
415418
headers: {
416-
Authorization: `Key ${this._config.apiKey}`,
419+
Authorization: `Key ${this.keyPool.next()}`,
417420
},
418421
});
419422

‎src/media/audio/providers/ReplicateAudioProvider.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
*/
3535

3636
import type { IAudioGenerator } from '../IAudioGenerator.js';
37+
import { ApiKeyPool } from '../../../core/providers/ApiKeyPool.js';
3738
import type { MusicGenerateRequest, SFXGenerateRequest, AudioResult } from '../types.js';
3839

3940
// ---------------------------------------------------------------------------
@@ -205,6 +206,7 @@ export class ReplicateAudioProvider implements IAudioGenerator {
205206

206207
/** Internal resolved configuration. */
207208
private _config!: Required<Pick<ReplicateAudioProviderConfig, 'apiKey' | 'baseURL' | 'defaultMusicModel' | 'defaultSfxModel' | 'pollIntervalMs' | 'timeoutMs'>>;
209+
private keyPool!: ApiKeyPool;
208210

209211
// -------------------------------------------------------------------------
210212
// Lifecycle
@@ -247,6 +249,7 @@ export class ReplicateAudioProvider implements IAudioGenerator {
247249
};
248250

249251
this.defaultModelId = this._config.defaultMusicModel;
252+
this.keyPool = new ApiKeyPool(apiKey);
250253
this.isInitialized = true;
251254
}
252255

@@ -394,7 +397,7 @@ export class ReplicateAudioProvider implements IAudioGenerator {
394397
const response = await fetch(`${this._config.baseURL}/predictions`, {
395398
method: 'POST',
396399
headers: {
397-
Authorization: `Token ${this._config.apiKey}`,
400+
Authorization: `Token ${this.keyPool.next()}`,
398401
'Content-Type': 'application/json',
399402
Prefer: 'wait=60',
400403
},
@@ -424,7 +427,7 @@ export class ReplicateAudioProvider implements IAudioGenerator {
424427
while (Date.now() - startedAt < this._config.timeoutMs) {
425428
const response = await fetch(url, {
426429
headers: {
427-
Authorization: `Token ${this._config.apiKey}`,
430+
Authorization: `Token ${this.keyPool.next()}`,
428431
},
429432
});
430433

‎src/media/audio/providers/StableAudioProvider.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
*/
2727

2828
import type { IAudioGenerator } from '../IAudioGenerator.js';
29+
import { ApiKeyPool } from '../../../core/providers/ApiKeyPool.js';
2930
import type { MusicGenerateRequest, SFXGenerateRequest, AudioResult } from '../types.js';
3031

3132
// ---------------------------------------------------------------------------
@@ -96,6 +97,7 @@ export class StableAudioProvider implements IAudioGenerator {
9697

9798
/** Internal resolved configuration. */
9899
private _config!: Required<Pick<StableAudioProviderConfig, 'apiKey' | 'baseURL'>> & StableAudioProviderConfig;
100+
private keyPool!: ApiKeyPool;
99101

100102
// -------------------------------------------------------------------------
101103
// Lifecycle
@@ -126,6 +128,7 @@ export class StableAudioProvider implements IAudioGenerator {
126128
};
127129

128130
this.defaultModelId = this._config.defaultModelId;
131+
this.keyPool = new ApiKeyPool(apiKey);
129132
this.isInitialized = true;
130133
}
131134

@@ -225,7 +228,7 @@ export class StableAudioProvider implements IAudioGenerator {
225228
const response = await fetch(url, {
226229
method: 'POST',
227230
headers: {
228-
Authorization: `Bearer ${this._config.apiKey}`,
231+
Authorization: `Bearer ${this.keyPool.next()}`,
229232
'Content-Type': 'application/json',
230233
},
231234
body: JSON.stringify(body),

‎src/media/audio/providers/SunoProvider.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
*/
2828

2929
import type { IAudioGenerator } from '../IAudioGenerator.js';
30+
import { ApiKeyPool } from '../../../core/providers/ApiKeyPool.js';
3031
import type { MusicGenerateRequest, SFXGenerateRequest, AudioResult } from '../types.js';
3132

3233
// ---------------------------------------------------------------------------
@@ -182,6 +183,7 @@ export class SunoProvider implements IAudioGenerator {
182183

183184
/** Internal resolved configuration. */
184185
private _config!: Required<Pick<SunoProviderConfig, 'apiKey' | 'baseURL' | 'pollIntervalMs' | 'timeoutMs'>> & SunoProviderConfig;
186+
private keyPool!: ApiKeyPool;
185187

186188
// -------------------------------------------------------------------------
187189
// Lifecycle
@@ -220,6 +222,7 @@ export class SunoProvider implements IAudioGenerator {
220222
};
221223

222224
this.defaultModelId = this._config.defaultModelId;
225+
this.keyPool = new ApiKeyPool(apiKey);
223226
this.isInitialized = true;
224227
}
225228

@@ -355,7 +358,7 @@ export class SunoProvider implements IAudioGenerator {
355358
const response = await fetch(`${this._config.baseURL}/predictions`, {
356359
method: 'POST',
357360
headers: {
358-
Authorization: `Token ${this._config.apiKey}`,
361+
Authorization: `Token ${this.keyPool.next()}`,
359362
'Content-Type': 'application/json',
360363
Prefer: 'wait=60',
361364
},
@@ -385,7 +388,7 @@ export class SunoProvider implements IAudioGenerator {
385388
while (Date.now() - startedAt < this._config.timeoutMs) {
386389
const response = await fetch(url, {
387390
headers: {
388-
Authorization: `Token ${this._config.apiKey}`,
391+
Authorization: `Token ${this.keyPool.next()}`,
389392
},
390393
});
391394

‎src/media/audio/providers/UdioProvider.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
*/
2626

2727
import type { IAudioGenerator } from '../IAudioGenerator.js';
28+
import { ApiKeyPool } from '../../../core/providers/ApiKeyPool.js';
2829
import type { MusicGenerateRequest, SFXGenerateRequest, AudioResult } from '../types.js';
2930

3031
// ---------------------------------------------------------------------------
@@ -173,6 +174,7 @@ export class UdioProvider implements IAudioGenerator {
173174

174175
/** Internal resolved configuration. */
175176
private _config!: Required<Pick<UdioProviderConfig, 'apiKey' | 'baseURL' | 'pollIntervalMs' | 'timeoutMs'>> & UdioProviderConfig;
177+
private keyPool!: ApiKeyPool;
176178

177179
// -------------------------------------------------------------------------
178180
// Lifecycle
@@ -211,6 +213,7 @@ export class UdioProvider implements IAudioGenerator {
211213
};
212214

213215
this.defaultModelId = this._config.defaultModelId;
216+
this.keyPool = new ApiKeyPool(apiKey);
214217
this.isInitialized = true;
215218
}
216219

@@ -345,7 +348,7 @@ export class UdioProvider implements IAudioGenerator {
345348
const response = await fetch(`${this._config.baseURL}/predictions`, {
346349
method: 'POST',
347350
headers: {
348-
Authorization: `Token ${this._config.apiKey}`,
351+
Authorization: `Token ${this.keyPool.next()}`,
349352
'Content-Type': 'application/json',
350353
Prefer: 'wait=60',
351354
},
@@ -375,7 +378,7 @@ export class UdioProvider implements IAudioGenerator {
375378
while (Date.now() - startedAt < this._config.timeoutMs) {
376379
const response = await fetch(url, {
377380
headers: {
378-
Authorization: `Token ${this._config.apiKey}`,
381+
Authorization: `Token ${this.keyPool.next()}`,
379382
},
380383
});
381384

‎src/media/images/providers/FalImageProvider.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,7 @@ export class FalImageProvider implements IImageProvider {
220220

221221
/** Internal resolved configuration. */
222222
private _config!: Required<Pick<FalImageProviderConfig, 'apiKey' | 'baseURL' | 'pollIntervalMs' | 'timeoutMs'>> & FalImageProviderConfig;
223+
private keyPool!: ApiKeyPool;
223224

224225
/**
225226
* Initialize the provider with API credentials and optional configuration.
@@ -259,6 +260,7 @@ export class FalImageProvider implements IImageProvider {
259260
};
260261

261262
this.defaultModelId = this._config.defaultModelId;
263+
this.keyPool = new ApiKeyPool(apiKey);
262264
this.isInitialized = true;
263265
}
264266

@@ -467,7 +469,7 @@ export class FalImageProvider implements IImageProvider {
467469
const response = await fetch(url, {
468470
method: 'POST',
469471
headers: {
470-
Authorization: `Key ${this._config.apiKey}`,
472+
Authorization: `Key ${this.keyPool.next()}`,
471473
'Content-Type': 'application/json',
472474
},
473475
body: JSON.stringify(body),
@@ -503,7 +505,7 @@ export class FalImageProvider implements IImageProvider {
503505

504506
const response = await fetch(url, {
505507
headers: {
506-
Authorization: `Key ${this._config.apiKey}`,
508+
Authorization: `Key ${this.keyPool.next()}`,
507509
},
508510
});
509511

@@ -548,7 +550,7 @@ export class FalImageProvider implements IImageProvider {
548550

549551
const response = await fetch(url, {
550552
headers: {
551-
Authorization: `Key ${this._config.apiKey}`,
553+
Authorization: `Key ${this.keyPool.next()}`,
552554
},
553555
});
554556

0 commit comments

Comments
 (0)