gate.py

import requests
import threading
import time

from bucket import Bucket
from syncer import sync_bucket_usage

class AIBudgetGate:
    def __init__(self, config: dict):
        self.budget_threshold = config["budget_threshold"]
        self.sync_interval = config["sync_interval"]
        self.buckets = [
            Bucket(
                admin_key = bucket_config["admin_key"],
                base_uri = bucket_config["base_uri"],
                models = config["model_groups"][bucket_config["model_group"]],
                name = name,
                token_budget = bucket_config["token_budget"],
                token_used = 0,
                upstream_url = bucket_config["upstream_url"]
            )
            for name, bucket_config in config["buckets"].items()
        ]

    def start_background_sync(self) -> None:
        for bucket in self.buckets:
            session = requests.Session()

            def sync_loop():
                while True:
                    sync_bucket_usage(bucket, session)
                    time.sleep(self.sync_interval)

            t = threading.Thread(target = sync_loop, daemon = True)
            t.start()