UNPKG

@just-every/ensemble

Version:

LLM provider abstraction layer with unified streaming interface

307 lines 13.4 kB
export class QuotaTracker { quotas = {}; updateCallback; openAIFreeQuota = { gpt4Family: { limit: 1000000, used: 0, }, gptMiniFamily: { limit: 10000000, used: 0, }, gpt4Models: ['gpt-4.5-preview', 'gpt-4.1', 'gpt-4o', 'o1', 'o3'], gptMiniModels: ['gpt-4.1-mini', 'gpt-4.1-nano', 'gpt-4o-mini', 'o1-mini', 'o4-mini'], }; constructor(updateCallback) { this.updateCallback = updateCallback; this.initializeProviderQuotas(); } setUpdateCallback(callback) { this.updateCallback = callback; } initializeProviderQuotas() { this.quotas['google'] = { provider: 'google', creditBalance: 0, creditLimit: 0, lastResetDate: new Date(), models: { 'gemini-2.5-pro-exp-03-25': { model: 'gemini-2.5-pro-exp-03-25', dailyTokenLimit: 1000000, dailyTokensUsed: 0, dailyRequestLimit: 25, dailyRequestsUsed: 0, rateLimit: { requestsPerMinute: 5, tokensPerMinute: 250000, }, lastResetDate: new Date(), }, 'gemini-2.5-flash-preview-04-17': { model: 'gemini-2.5-flash-preview-04-17', dailyTokenLimit: 0, dailyTokensUsed: 0, dailyRequestLimit: 500, dailyRequestsUsed: 0, rateLimit: { requestsPerMinute: 10, tokensPerMinute: 250000, }, lastResetDate: new Date(), }, 'gemini-2.0-flash': { model: 'gemini-2.0-flash', dailyTokenLimit: 0, dailyTokensUsed: 0, dailyRequestLimit: 1500, dailyRequestsUsed: 0, rateLimit: { requestsPerMinute: 15, tokensPerMinute: 1000000, }, lastResetDate: new Date(), }, 'gemini-2.0-flash-lite': { model: 'gemini-2.0-flash-lite', dailyTokenLimit: 0, dailyTokensUsed: 0, dailyRequestLimit: 1500, dailyRequestsUsed: 0, rateLimit: { requestsPerMinute: 30, tokensPerMinute: 1000000, }, lastResetDate: new Date(), }, }, }; this.quotas['openai'] = { provider: 'openai', creditBalance: 0, creditLimit: 0, lastResetDate: new Date(), info: { freeQuota: this.openAIFreeQuota, }, models: {}, }; this.quotas['xai'] = { provider: 'xai', creditBalance: 150, creditLimit: 150, lastResetDate: new Date(), models: {}, }; } getProviderQuota(provider) { return (this.quotas[provider] || { provider, models: {}, }); } getModelQuota(provider, model) { const providerQuota = this.getProviderQuota(provider); if (providerQuota.models[model]) { return providerQuota.models[model]; } return null; } trackUsage(provider, model, inputTokens, outputTokens) { if (!this.quotas[provider]) { return true; } const providerQuota = this.quotas[provider]; let significantChange = false; const totalTokens = inputTokens + outputTokens; const today = new Date(); if (providerQuota.lastResetDate && (today.getUTCDate() !== providerQuota.lastResetDate.getUTCDate() || today.getUTCMonth() !== providerQuota.lastResetDate.getUTCMonth() || today.getUTCFullYear() !== providerQuota.lastResetDate.getUTCFullYear())) { for (const modelKey in providerQuota.models) { const modelQuota = providerQuota.models[modelKey]; modelQuota.dailyTokensUsed = 0; modelQuota.dailyRequestsUsed = 0; modelQuota.lastResetDate = today; } if (provider === 'openai' && providerQuota.info?.freeQuota) { const freeQuota = providerQuota.info.freeQuota; freeQuota.gpt4Family.used = 0; freeQuota.gptMiniFamily.used = 0; } providerQuota.lastResetDate = today; significantChange = true; } if (provider === 'openai' && providerQuota.info?.freeQuota) { const freeQuota = providerQuota.info.freeQuota; if (freeQuota.gpt4Models.includes(model)) { const prevUsed = freeQuota.gpt4Family.used; freeQuota.gpt4Family.used += totalTokens; console.log(`[QuotaTracker] OpenAI GPT-4 family usage: ${freeQuota.gpt4Family.used}/${freeQuota.gpt4Family.limit}`); const prevPercent = Math.floor((prevUsed / freeQuota.gpt4Family.limit) * 10); const currPercent = Math.floor((freeQuota.gpt4Family.used / freeQuota.gpt4Family.limit) * 10); if (prevPercent !== currPercent) { significantChange = true; } if (freeQuota.gpt4Family.used >= freeQuota.gpt4Family.limit) { console.log(`[QuotaTracker] OpenAI GPT-4 family daily limit reached: ${freeQuota.gpt4Family.used} > ${freeQuota.gpt4Family.limit}`); significantChange = true; this.sendQuotaUpdate(); return false; } } if (freeQuota.gptMiniModels.includes(model)) { const prevUsed = freeQuota.gptMiniFamily.used; freeQuota.gptMiniFamily.used += totalTokens; console.log(`[QuotaTracker] OpenAI GPT-Mini family usage: ${freeQuota.gptMiniFamily.used}/${freeQuota.gptMiniFamily.limit}`); const prevPercent = Math.floor((prevUsed / freeQuota.gptMiniFamily.limit) * 10); const currPercent = Math.floor((freeQuota.gptMiniFamily.used / freeQuota.gptMiniFamily.limit) * 10); if (prevPercent !== currPercent) { significantChange = true; } if (freeQuota.gptMiniFamily.used >= freeQuota.gptMiniFamily.limit) { console.log(`[QuotaTracker] OpenAI GPT-Mini family daily limit reached: ${freeQuota.gptMiniFamily.used} > ${freeQuota.gptMiniFamily.limit}`); significantChange = true; this.sendQuotaUpdate(); return false; } } } const modelQuota = this.getModelQuota(provider, model); if (!modelQuota && provider !== 'openai') { return true; } if (modelQuota) { const previousDailyTokensUsed = modelQuota.dailyTokensUsed; modelQuota.dailyTokensUsed += totalTokens; const previousDailyRequestsUsed = modelQuota.dailyRequestsUsed; modelQuota.dailyRequestsUsed += 1; if (modelQuota.dailyTokenLimit > 0) { const previousPercent = Math.floor((previousDailyTokensUsed / modelQuota.dailyTokenLimit) * 10); const currentPercent = Math.floor((modelQuota.dailyTokensUsed / modelQuota.dailyTokenLimit) * 10); if (previousPercent !== currentPercent) { significantChange = true; } } if (modelQuota.dailyRequestLimit > 0) { const previousPercent = Math.floor((previousDailyRequestsUsed / modelQuota.dailyRequestLimit) * 10); const currentPercent = Math.floor((modelQuota.dailyRequestsUsed / modelQuota.dailyRequestLimit) * 10); if (previousPercent !== currentPercent) { significantChange = true; } } if (modelQuota.dailyTokenLimit > 0 && modelQuota.dailyTokensUsed >= modelQuota.dailyTokenLimit) { console.log(`[QuotaTracker] ${provider} model ${model} daily token limit reached: ${modelQuota.dailyTokensUsed} > ${modelQuota.dailyTokenLimit}`); significantChange = true; this.sendQuotaUpdate(); return false; } if (modelQuota.dailyRequestLimit > 0 && modelQuota.dailyRequestsUsed >= modelQuota.dailyRequestLimit) { console.log(`[QuotaTracker] ${provider} model ${model} daily request limit reached: ${modelQuota.dailyRequestsUsed} > ${modelQuota.dailyRequestLimit}`); significantChange = true; this.sendQuotaUpdate(); return false; } } if (significantChange) { this.sendQuotaUpdate(); } return true; } hasQuota(provider, model) { if (provider === 'openai') { return this.hasOpenAIFreeQuota(model); } const modelQuota = this.getModelQuota(provider, model); if (modelQuota) { return ((modelQuota.dailyTokenLimit === 0 || modelQuota.dailyTokensUsed < modelQuota.dailyTokenLimit) && (modelQuota.dailyRequestLimit === 0 || modelQuota.dailyRequestsUsed < modelQuota.dailyRequestLimit)); } return true; } hasOpenAIFreeQuota(model) { const providerQuota = this.getProviderQuota('openai'); if (!providerQuota || !providerQuota.info?.freeQuota) return true; const freeQuota = providerQuota.info.freeQuota; if (freeQuota.gpt4Models.includes(model)) { return freeQuota.gpt4Family.used < freeQuota.gpt4Family.limit; } if (freeQuota.gptMiniModels.includes(model)) { return freeQuota.gptMiniFamily.used < freeQuota.gptMiniFamily.limit; } return true; } trackCreditUsage(provider, creditAmount) { const providerQuota = this.getProviderQuota(provider); if (!providerQuota || providerQuota.creditBalance === undefined) { return; } if (providerQuota.creditBalance !== undefined) { providerQuota.creditBalance = Math.max(0, providerQuota.creditBalance - creditAmount); } } getCreditBalance(provider) { return this.getProviderQuota(provider)?.creditBalance || 0; } getSummary() { const summary = {}; for (const [provider, providerQuota] of Object.entries(this.quotas)) { summary[provider] = { creditBalance: providerQuota.creditBalance, creditLimit: providerQuota.creditLimit, models: {}, }; for (const [modelName, modelQuota] of Object.entries(providerQuota.models)) { summary[provider].models[modelName] = { tokens: { used: modelQuota.dailyTokensUsed, limit: modelQuota.dailyTokenLimit, percent: modelQuota.dailyTokenLimit > 0 ? (modelQuota.dailyTokensUsed / modelQuota.dailyTokenLimit) * 100 : 0, }, requests: { used: modelQuota.dailyRequestsUsed, limit: modelQuota.dailyRequestLimit, percent: modelQuota.dailyRequestLimit > 0 ? (modelQuota.dailyRequestsUsed / modelQuota.dailyRequestLimit) * 100 : 0, }, lastReset: modelQuota.lastResetDate, }; } if (provider === 'openai' && providerQuota.info?.freeQuota) { const freeQuota = providerQuota.info.freeQuota; summary[provider].freeTier = { gpt4Family: { used: freeQuota.gpt4Family.used, limit: freeQuota.gpt4Family.limit, percent: (freeQuota.gpt4Family.used / freeQuota.gpt4Family.limit) * 100, }, gptMiniFamily: { used: freeQuota.gptMiniFamily.used, limit: freeQuota.gptMiniFamily.limit, percent: (freeQuota.gptMiniFamily.used / freeQuota.gptMiniFamily.limit) * 100, }, }; } } return summary; } sendQuotaUpdate() { try { if (this.updateCallback) { const quotas = this.getSummary(); this.updateCallback(quotas); } } catch (error) { console.error('Error sending quota update:', error); } } } export const quotaTracker = new QuotaTracker(); //# sourceMappingURL=quota_tracker.js.map