token pools

src/antrhopic.ts

@@ -1,16 +1,27 @@
 import {Anthropic as anthropic} from '@anthropic-ai/sdk';
-import {findByProp, objectMap, JSONSanitize, JSONAttemptParse} from '@ztimson/utils';
+import {findByProp, objectMap, JSONSanitize, JSONAttemptParse, makeArray} from '@ztimson/utils';
 import {AbortablePromise, Ai} from './ai.ts';
 import {LLMMessage, LLMRequest} from './llm.ts';
 import {LLMProvider} from './provider.ts';
+import {TokenPool} from './token-pool.ts';
 import {convertSchema} from './tools.ts';
 
 export class Anthropic extends LLMProvider {
-	client!: anthropic;
+	private clients = new Map<string, anthropic>();
+	tokenPool!: TokenPool;
 
-	constructor(public readonly ai: Ai, public readonly apiToken: string, public model: string) {
+	constructor(public readonly ai: Ai, public readonly apiToken: string | string[], public model: string) {
 		super();
-		this.client = new anthropic({apiKey: apiToken});
+		this.tokenPool = new TokenPool(...makeArray(apiToken).filter(Boolean));
+	}
+
+	private getClient(token: string): anthropic {
+		let client = this.clients.get(token);
+		if(!client) {
+			client = new anthropic({apiKey: token});
+			this.clients.set(token, client);
+		}
+		return client;
 	}
 
 	private toStandard(history: any[]): LLMMessage[] {
@@ -90,7 +101,7 @@ export class Anthropic extends LLMProvider {
 			do {
 				requestParams.messages = history.map(({timestamp, ...m}) => m);
 				const callStart = Date.now();
-				resp = await this.client.messages.create(requestParams).catch(err => {
+				resp = await this.tokenPool.run(token => this.getClient(token).messages.create(requestParams)).catch(err => {
 					err.message += `\n\nMessages:\n${JSON.stringify(history, null, 2)}`;
 					throw err;
 				});

src/index.ts

@@ -5,5 +5,6 @@ export * from './llm';
 export * from './memory';
 export * from './open-ai';
 export * from './provider';
+export * from './token-pool'
 export * from './tools';
 export * from './vision';

src/llm.ts

@@ -11,8 +11,8 @@ import {Memory, MemoryCache, MemoryManager, MemoryOptions} from './memory.ts';
 
 const MAX_AGENT_DEPTH = 5;
 
-export type AnthropicConfig = {proto: 'anthropic', token: string};
+export type AnthropicConfig = {proto: 'anthropic', token: string | string[]};
-export type OpenAiConfig = {proto: 'openai', host?: string, token: string};
+export type OpenAiConfig = {proto: 'openai', host?: string, token: string | string[]};
 
 export type Agent = {
 	name: string;

src/open-ai.ts

@@ -1,19 +1,28 @@
 import {OpenAI as openAI} from 'openai';
-import {findByProp, objectMap, JSONSanitize, JSONAttemptParse, clean} from '@ztimson/utils';
+import {findByProp, objectMap, JSONSanitize, JSONAttemptParse, clean, makeArray} from '@ztimson/utils';
 import {AbortablePromise, Ai} from './ai.ts';
 import {LLMMessage, LLMRequest} from './llm.ts';
 import {LLMProvider} from './provider.ts';
+import {TokenPool} from './token-pool.ts';
 import {convertSchema} from './tools.ts';
 
 export class OpenAi extends LLMProvider {
-	client!: openAI;
+	tokenPool!: TokenPool;
+	private clients = new Map<string, openAI>();
 
-	constructor(public readonly ai: Ai, public readonly host: string | null, public readonly token: string, public model: string) {
+	constructor(public readonly ai: Ai, public readonly host: string | null, public readonly token: string | string[], public model: string) {
 		super();
-		this.client = new openAI(clean({
+		const tokens = makeArray(token).filter(Boolean);
-			baseURL: host,
+		this.tokenPool = new TokenPool(...(tokens.length ? tokens : [host ? 'ignored' : '']));
-			apiKey: token || (host ? 'ignored' : undefined)
+	}
-		}));
+
+	private getClient(token: string): openAI {
+		let client = this.clients.get(token);
+		if(!client) {
+			client = new openAI(clean({baseURL: this.host, apiKey: token || undefined}));
+			this.clients.set(token, client);
+		}
+		return client;
 	}
 
 	private toStandard(history: any[]): LLMMessage[] {
@@ -117,7 +126,7 @@ export class OpenAi extends LLMProvider {
 			do {
 				requestParams.messages = history.map(({timestamp, ...m}) => m);
 				const callStart = Date.now();
-				resp = await this.client.chat.completions.create(requestParams).catch(err => {
+				resp = await this.tokenPool.run(token => this.getClient(token).chat.completions.create(requestParams)).catch(err => {
 					err.message += `\n\nMessages:\n${JSON.stringify(history, null, 2)}`;
 					throw err;
 				});

src/token-pool.ts

@@ -0,0 +1,65 @@
+const DEFAULT_COOLDOWN = 15 * 60 * 1000;
+
+type TokenState = {
+	token: string;
+	cooldownUntil: number; // 0 = available now
+	lastError?: {code: number, message: string};
+};
+
+export class TokenPoolExhaustedError extends Error {
+	constructor(public tokens: Record<string, {code: number, message: string}>) {
+		super(`All tokens exhausted:\n${Object.entries(tokens).map(([t, e]) => `${t}: [${e.code}] ${e.message}`).join('\n')}`);
+		this.name = 'TokenPoolExhaustedError';
+	}
+}
+
+export class TokenPool {
+	private states: TokenState[];
+
+	constructor(...tokens: string[]) {
+		this.states = tokens.map(token => ({token, cooldownUntil: 0}));
+	}
+
+	private preview(token: string): string {
+		return token.length <= 8 ? '****' : `${token.slice(0, 4)}...${token.slice(-4)}`;
+	}
+
+	/** Anthropic & OpenAI SDKs both attach `status` to thrown errors */
+	private statusCode(err: any): number {
+		return err?.status ?? err?.response?.status ?? err?.statusCode;
+	}
+
+	private retryAfter(err: any): number {
+		const headers = err?.headers || err?.response?.headers;
+		const raw = headers?.get?.('retry-after') ?? headers?.['retry-after'];
+		if(raw) {
+			const seconds = Number(raw);
+			if(!isNaN(seconds)) return Date.now() + seconds * 1000;
+			const date = new Date(raw).getTime();
+			if(!isNaN(date)) return date;
+		}
+		return Date.now() + DEFAULT_COOLDOWN;
+	}
+
+	async run<T>(fn: (token: string) => Promise<T>): Promise<T> {
+		const now = Date.now();
+		for(const state of this.states) {
+			if(state.cooldownUntil > now) continue;
+			try {
+				const result = await fn(state.token);
+				state.cooldownUntil = 0;
+				state.lastError = undefined;
+				return result;
+			} catch(err: any) {
+				const code = this.statusCode(err);
+				if(![401, 403, 429].includes(code)) throw err;
+				state.cooldownUntil = code === 429 ? this.retryAfter(err) : Date.now() + DEFAULT_COOLDOWN;
+				state.lastError = {code, message: err?.message || 'Unknown error'};
+			}
+		}
+
+		const failures: Record<string, {code: number, message: string}> = {};
+		this.states.forEach(s => { if(s.lastError) failures[this.preview(s.token)] = s.lastError; });
+		throw new TokenPoolExhaustedError(failures);
+	}
+}