feat: add multi-key auth profile rotation for model providers
This commit is contained in:
@@ -550,6 +550,20 @@ describe('configSchema — models auth_mode', () => {
|
||||
});
|
||||
expect(moonshot.models.default.provider).toBe('moonshot');
|
||||
});
|
||||
|
||||
it('accepts multiple api_keys per model tier', () => {
|
||||
const result = configSchema.parse({
|
||||
...minimalConfig,
|
||||
models: {
|
||||
default: {
|
||||
provider: 'openai',
|
||||
model: 'gpt-4o',
|
||||
api_keys: ['sk-1', 'sk-2'],
|
||||
},
|
||||
},
|
||||
});
|
||||
expect(result.models.default.api_keys).toEqual(['sk-1', 'sk-2']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('configSchema — matrix', () => {
|
||||
|
||||
@@ -132,6 +132,7 @@ const modelConfigBaseSchema = z.object({
|
||||
model: z.string(),
|
||||
endpoint: z.string().optional(),
|
||||
api_key: z.string().optional(),
|
||||
api_keys: z.array(z.string().min(1)).optional(),
|
||||
auth_token: z.string().optional(),
|
||||
/** Credential selection strategy for this tier (provider-specific). */
|
||||
auth_mode: z.enum(['auto', 'api_key', 'oauth']).optional(),
|
||||
|
||||
@@ -28,6 +28,16 @@ describe('createClientFromConfig', () => {
|
||||
expect(client.constructor.name).toBe('AnthropicClient');
|
||||
});
|
||||
|
||||
it('creates RotatingModelClient for anthropic when api_keys are provided', async () => {
|
||||
const { createClientFromConfig } = await loadFactory();
|
||||
const client = createClientFromConfig({
|
||||
provider: 'anthropic',
|
||||
model: 'claude-sonnet-4-5-20250514',
|
||||
api_keys: ['sk-ant-1', 'sk-ant-2'],
|
||||
});
|
||||
expect(client.constructor.name).toBe('RotatingModelClient');
|
||||
});
|
||||
|
||||
it('creates OpenAIClient for openai provider', async () => {
|
||||
const { createClientFromConfig } = await loadFactory();
|
||||
const client = createClientFromConfig({
|
||||
@@ -38,6 +48,16 @@ describe('createClientFromConfig', () => {
|
||||
expect(client.constructor.name).toBe('OpenAIClient');
|
||||
});
|
||||
|
||||
it('creates RotatingModelClient for openai when api_keys are provided', async () => {
|
||||
const { createClientFromConfig } = await loadFactory();
|
||||
const client = createClientFromConfig({
|
||||
provider: 'openai',
|
||||
model: 'gpt-4o',
|
||||
api_keys: ['sk-1', 'sk-2'],
|
||||
});
|
||||
expect(client.constructor.name).toBe('RotatingModelClient');
|
||||
});
|
||||
|
||||
it('creates OllamaClient for ollama provider', async () => {
|
||||
const { createClientFromConfig } = await loadFactory();
|
||||
const client = createClientFromConfig({
|
||||
|
||||
+105
-33
@@ -1,5 +1,5 @@
|
||||
import type { Config, ModelConfig } from '../config/index.js';
|
||||
import { AnthropicClient, OpenAIClient, OllamaClient, LlamaCppClient, GeminiClient, BedrockClient, GitHubModelsClient, SyntheticClient, ModelRouter, DEFAULT_RETRY_CONFIG } from '../models/index.js';
|
||||
import { AnthropicClient, OpenAIClient, OllamaClient, LlamaCppClient, GeminiClient, BedrockClient, GitHubModelsClient, SyntheticClient, RotatingModelClient, ModelRouter, DEFAULT_RETRY_CONFIG } from '../models/index.js';
|
||||
import type { ModelClient, RetryConfig, ModelTier } from '../models/index.js';
|
||||
import { logger } from '../logger.js';
|
||||
import { getZaiApiKey } from '../auth/zai.js';
|
||||
@@ -33,6 +33,30 @@ function requireApiKey(cfg: ModelConfig, envVar: string): string {
|
||||
return key;
|
||||
}
|
||||
|
||||
function resolveApiKeyPool(cfg: ModelConfig, envVar?: string): string[] {
|
||||
const configured = (cfg.api_keys ?? []).map((key) => key.trim()).filter(Boolean);
|
||||
if (configured.length > 0) {
|
||||
return configured;
|
||||
}
|
||||
if (cfg.api_key?.trim()) {
|
||||
return [cfg.api_key.trim()];
|
||||
}
|
||||
if (envVar && process.env[envVar]?.trim()) {
|
||||
return [process.env[envVar]!.trim()];
|
||||
}
|
||||
return [];
|
||||
}
|
||||
|
||||
function createApiKeyClient(
|
||||
keys: string[],
|
||||
build: (apiKey: string) => ModelClient,
|
||||
): ModelClient {
|
||||
if (keys.length === 1) {
|
||||
return build(keys[0]);
|
||||
}
|
||||
return new RotatingModelClient(keys.map((key) => build(key)));
|
||||
}
|
||||
|
||||
function resolveZaiCredential(cfg: ModelConfig): string {
|
||||
const raw = cfg.api_key
|
||||
?? cfg.auth_token
|
||||
@@ -75,26 +99,34 @@ export function createClientFromConfig(cfg: ModelConfig): ModelClient {
|
||||
}
|
||||
|
||||
if (authMode === 'api_key') {
|
||||
const apiKey = cfg.api_key ?? getAnthropicApiKey();
|
||||
if (!apiKey) {
|
||||
const keys = resolveApiKeyPool(cfg);
|
||||
const envKey = getAnthropicApiKey();
|
||||
const allKeys = keys.length > 0
|
||||
? keys
|
||||
: (envKey ? [envKey] : []);
|
||||
if (allKeys.length === 0) {
|
||||
throw new Error(
|
||||
'Anthropic API key not configured (auth_mode: api_key). ' +
|
||||
'Set ANTHROPIC_API_KEY, run `flynn anthropic-auth`, or provide api_key in config.',
|
||||
'Set ANTHROPIC_API_KEY, run `flynn anthropic-auth`, or provide api_key/api_keys in config.',
|
||||
);
|
||||
}
|
||||
return new AnthropicClient({
|
||||
return createApiKeyClient(allKeys, (apiKey) => new AnthropicClient({
|
||||
model: cfg.model,
|
||||
apiKey,
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
// auto: prefer API key, then token
|
||||
const apiKey = cfg.api_key ?? getAnthropicApiKey();
|
||||
if (apiKey) {
|
||||
return new AnthropicClient({
|
||||
// auto: prefer API keys, then token
|
||||
const configuredKeys = resolveApiKeyPool(cfg);
|
||||
const envKey = getAnthropicApiKey();
|
||||
const allKeys = configuredKeys.length > 0
|
||||
? configuredKeys
|
||||
: (envKey ? [envKey] : []);
|
||||
if (allKeys.length > 0) {
|
||||
return createApiKeyClient(allKeys, (apiKey) => new AnthropicClient({
|
||||
model: cfg.model,
|
||||
apiKey,
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
const token = cfg.auth_token ?? getAnthropicAuthToken();
|
||||
@@ -130,26 +162,34 @@ export function createClientFromConfig(cfg: ModelConfig): ModelClient {
|
||||
}
|
||||
|
||||
if (authMode === 'api_key') {
|
||||
const apiKey = cfg.api_key ?? getOpenAIApiKey();
|
||||
if (!apiKey) {
|
||||
const keys = resolveApiKeyPool(cfg);
|
||||
const envKey = getOpenAIApiKey();
|
||||
const allKeys = keys.length > 0
|
||||
? keys
|
||||
: (envKey ? [envKey] : []);
|
||||
if (allKeys.length === 0) {
|
||||
throw new Error(
|
||||
'OpenAI API key not configured (auth_mode: api_key). ' +
|
||||
'Set OPENAI_API_KEY, run `flynn openai-key`, or provide api_key in config.',
|
||||
'Set OPENAI_API_KEY, run `flynn openai-key`, or provide api_key/api_keys in config.',
|
||||
);
|
||||
}
|
||||
return new OpenAIClient({
|
||||
return createApiKeyClient(allKeys, (apiKey) => new OpenAIClient({
|
||||
model: cfg.model,
|
||||
apiKey,
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
// auto: prefer API key, then OAuth
|
||||
const apiKey = cfg.api_key ?? getOpenAIApiKey();
|
||||
if (apiKey) {
|
||||
return new OpenAIClient({
|
||||
// auto: prefer API keys, then OAuth
|
||||
const configuredKeys = resolveApiKeyPool(cfg);
|
||||
const envKey = getOpenAIApiKey();
|
||||
const allKeys = configuredKeys.length > 0
|
||||
? configuredKeys
|
||||
: (envKey ? [envKey] : []);
|
||||
if (allKeys.length > 0) {
|
||||
return createApiKeyClient(allKeys, (apiKey) => new OpenAIClient({
|
||||
model: cfg.model,
|
||||
apiKey,
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
const existing = loadStoredOpenAIAuth();
|
||||
@@ -184,11 +224,19 @@ export function createClientFromConfig(cfg: ModelConfig): ModelClient {
|
||||
apiKey: cfg.api_key,
|
||||
});
|
||||
case 'openrouter':
|
||||
return new OpenAIClient({
|
||||
{
|
||||
const keys = resolveApiKeyPool(cfg, 'OPENROUTER_API_KEY');
|
||||
if (keys.length === 0) {
|
||||
throw new Error(
|
||||
'API key required for openrouter. Set OPENROUTER_API_KEY or provide api_key/api_keys in config.',
|
||||
);
|
||||
}
|
||||
return createApiKeyClient(keys, (apiKey) => new OpenAIClient({
|
||||
model: cfg.model,
|
||||
apiKey: requireApiKey(cfg, 'OPENROUTER_API_KEY'),
|
||||
apiKey,
|
||||
baseURL: cfg.endpoint ?? 'https://openrouter.ai/api/v1',
|
||||
});
|
||||
}));
|
||||
}
|
||||
case 'vercel':
|
||||
return new OpenAIClient({
|
||||
model: cfg.model,
|
||||
@@ -202,23 +250,47 @@ export function createClientFromConfig(cfg: ModelConfig): ModelClient {
|
||||
baseURL: cfg.endpoint ?? 'https://api.z.ai/api/paas/v4',
|
||||
});
|
||||
case 'xai':
|
||||
return new OpenAIClient({
|
||||
{
|
||||
const keys = resolveApiKeyPool(cfg, 'XAI_API_KEY');
|
||||
if (keys.length === 0) {
|
||||
throw new Error(
|
||||
'API key required for xai. Set XAI_API_KEY or provide api_key/api_keys in config.',
|
||||
);
|
||||
}
|
||||
return createApiKeyClient(keys, (apiKey) => new OpenAIClient({
|
||||
model: cfg.model,
|
||||
apiKey: requireApiKey(cfg, 'XAI_API_KEY'),
|
||||
apiKey,
|
||||
baseURL: cfg.endpoint ?? 'https://api.x.ai/v1',
|
||||
});
|
||||
}));
|
||||
}
|
||||
case 'minimax':
|
||||
return new OpenAIClient({
|
||||
{
|
||||
const keys = resolveApiKeyPool(cfg, 'MINIMAX_API_KEY');
|
||||
if (keys.length === 0) {
|
||||
throw new Error(
|
||||
'API key required for minimax. Set MINIMAX_API_KEY or provide api_key/api_keys in config.',
|
||||
);
|
||||
}
|
||||
return createApiKeyClient(keys, (apiKey) => new OpenAIClient({
|
||||
model: cfg.model,
|
||||
apiKey: requireApiKey(cfg, 'MINIMAX_API_KEY'),
|
||||
apiKey,
|
||||
baseURL: cfg.endpoint ?? 'https://api.minimax.io/v1',
|
||||
});
|
||||
}));
|
||||
}
|
||||
case 'moonshot':
|
||||
return new OpenAIClient({
|
||||
{
|
||||
const keys = resolveApiKeyPool(cfg, 'MOONSHOT_API_KEY');
|
||||
if (keys.length === 0) {
|
||||
throw new Error(
|
||||
'API key required for moonshot. Set MOONSHOT_API_KEY or provide api_key/api_keys in config.',
|
||||
);
|
||||
}
|
||||
return createApiKeyClient(keys, (apiKey) => new OpenAIClient({
|
||||
model: cfg.model,
|
||||
apiKey: requireApiKey(cfg, 'MOONSHOT_API_KEY'),
|
||||
apiKey,
|
||||
baseURL: cfg.endpoint ?? 'https://api.moonshot.cn/v1',
|
||||
});
|
||||
}));
|
||||
}
|
||||
case 'bedrock':
|
||||
return new BedrockClient({
|
||||
model: cfg.model,
|
||||
|
||||
@@ -7,6 +7,7 @@ export { SyntheticClient, type SyntheticClientConfig } from './synthetic.js';
|
||||
export { OllamaClient, type OllamaClientConfig } from './local/index.js';
|
||||
export { LlamaCppClient, type LlamaCppClientConfig } from './local/index.js';
|
||||
export { ModelRouter, type ModelRouterConfig, type ModelTier } from './router.js';
|
||||
export { RotatingModelClient } from './rotating.js';
|
||||
export { withRetry, isRetryable, DEFAULT_RETRY_CONFIG, type RetryConfig } from './retry.js';
|
||||
export { estimateCost, MODEL_COSTS_PER_MILLION } from './costs.js';
|
||||
export { supportsAudioInput } from './capabilities.js';
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import type { ModelClient } from './types.js';
|
||||
import { RotatingModelClient } from './rotating.js';
|
||||
|
||||
function makeClient(chatImpl: ModelClient['chat']): ModelClient {
|
||||
return { chat: chatImpl };
|
||||
}
|
||||
|
||||
describe('RotatingModelClient', () => {
|
||||
it('throws when created with no clients', () => {
|
||||
expect(() => new RotatingModelClient([])).toThrow(/at least one client/i);
|
||||
});
|
||||
|
||||
it('falls through to the next profile when the first fails', async () => {
|
||||
const first = makeClient(vi.fn().mockRejectedValue(new Error('rate limited')));
|
||||
const second = makeClient(vi.fn().mockResolvedValue({ content: 'ok' }));
|
||||
const rotating = new RotatingModelClient([first, second]);
|
||||
|
||||
const response = await rotating.chat({ messages: [{ role: 'user', content: 'hello' }] });
|
||||
expect(response.content).toBe('ok');
|
||||
expect(first.chat).toHaveBeenCalledTimes(1);
|
||||
expect(second.chat).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('sticks to the last successful profile until it fails', async () => {
|
||||
const first = makeClient(vi.fn().mockRejectedValue(new Error('429')));
|
||||
const second = makeClient(vi.fn().mockResolvedValue({ content: 'ok' }));
|
||||
const rotating = new RotatingModelClient([first, second]);
|
||||
|
||||
await rotating.chat({ messages: [{ role: 'user', content: 'a' }] });
|
||||
await rotating.chat({ messages: [{ role: 'user', content: 'b' }] });
|
||||
|
||||
expect(first.chat).toHaveBeenCalledTimes(1);
|
||||
expect(second.chat).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,64 @@
|
||||
import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient } from './types.js';
|
||||
|
||||
/**
|
||||
* Model client wrapper that rotates across equivalent auth profiles (e.g. API keys).
|
||||
* Sticky-by-success behavior: keep using the last successful profile until it fails.
|
||||
*/
|
||||
export class RotatingModelClient implements ModelClient {
|
||||
private readonly clients: ModelClient[];
|
||||
private currentIndex = 0;
|
||||
|
||||
constructor(clients: ModelClient[]) {
|
||||
if (clients.length === 0) {
|
||||
throw new Error('RotatingModelClient requires at least one client');
|
||||
}
|
||||
this.clients = clients;
|
||||
}
|
||||
|
||||
async chat(request: ChatRequest): Promise<ChatResponse> {
|
||||
const start = this.currentIndex;
|
||||
const errors: Error[] = [];
|
||||
|
||||
for (let offset = 0; offset < this.clients.length; offset += 1) {
|
||||
const index = (start + offset) % this.clients.length;
|
||||
const client = this.clients[index];
|
||||
try {
|
||||
const response = await client.chat(request);
|
||||
this.currentIndex = index;
|
||||
return response;
|
||||
} catch (error) {
|
||||
errors.push(error instanceof Error ? error : new Error(String(error)));
|
||||
}
|
||||
}
|
||||
|
||||
throw new Error(`All auth profiles failed: ${errors.map((e) => e.message).join(', ')}`);
|
||||
}
|
||||
|
||||
async *chatStream(request: ChatRequest): AsyncIterable<ChatStreamEvent> {
|
||||
const start = this.currentIndex;
|
||||
|
||||
for (let offset = 0; offset < this.clients.length; offset += 1) {
|
||||
const index = (start + offset) % this.clients.length;
|
||||
const client = this.clients[index];
|
||||
if (!client.chatStream) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let failed = false;
|
||||
for await (const event of client.chatStream(request)) {
|
||||
if (event.type === 'error') {
|
||||
failed = true;
|
||||
break;
|
||||
}
|
||||
yield event;
|
||||
}
|
||||
|
||||
if (!failed) {
|
||||
this.currentIndex = index;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
yield { type: 'error', error: new Error('All auth profiles failed for streaming') };
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user