token pools

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({
-			baseURL: host,
-			apiKey: token || (host ? 'ignored' : undefined)
-		}));
+		const tokens = makeArray(token).filter(Boolean);
+		this.tokenPool = new TokenPool(...(tokens.length ? tokens : [host ? 'ignored' : '']));
+	}
+
+	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;
 				});