44 lines
1.2 KiB
TypeScript
44 lines
1.2 KiB
TypeScript
|
import fs from "fs";
|
||
|
|
import path from "path";
|
||
|
|
|
||
|
|
const QUOTA_DIR = path.join(process.cwd(), "data", "search-quota");
|
||
|
|
|
||
|
|
function todayKey(): string {
|
||
|
|
return new Date().toISOString().slice(0, 10);
|
||
|
|
}
|
||
|
|
|
||
|
|
function quotaFile(provider: string): string {
|
||
|
|
fs.mkdirSync(QUOTA_DIR, { recursive: true });
|
||
|
|
return path.join(QUOTA_DIR, `${provider}-${todayKey()}.json`);
|
||
|
|
}
|
||
|
|
|
||
|
|
interface QuotaState {
|
||
|
|
count: number;
|
||
|
|
}
|
||
|
|
|
||
|
|
export function getDailyQuotaUsed(provider: string): number {
|
||
|
|
try {
|
||
|
|
const file = quotaFile(provider);
|
||
|
|
if (!fs.existsSync(file)) return 0;
|
||
|
|
const state = JSON.parse(fs.readFileSync(file, "utf8")) as QuotaState;
|
||
|
|
return state.count ?? 0;
|
||
|
|
} catch {
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
export function incrementDailyQuota(provider: string): number {
|
||
|
|
const used = getDailyQuotaUsed(provider) + 1;
|
||
|
|
fs.writeFileSync(quotaFile(provider), JSON.stringify({ count: used } satisfies QuotaState));
|
||
|
|
return used;
|
||
|
|
}
|
||
|
|
|
||
|
|
export function canUseDailyQuota(provider: string, limit: number): boolean {
|
||
|
|
return getDailyQuotaUsed(provider) < limit;
|
||
|
|
}
|
||
|
|
|
||
|
|
export function resetQuotaForTests(): void {
|
||
|
|
if (fs.existsSync(QUOTA_DIR)) {
|
||
|
|
fs.rmSync(QUOTA_DIR, { recursive: true, force: true });
|
||
|
|
}
|
||
|
|
}
|