Flux

"""
Flux API Rate Limiter: In-Memory Sliding-Window Token Bucket
"""
import time
from typing import Dict, Tuple
from pydantic import BaseModel

class RateLimitDecision(BaseModel):
    allowed: bool
    remaining_tokens: float
    retry_after_sec: float
    tenant_id: str

class FluxRateLimiter:
    def __init__(self, capacity: float = 100.0, refill_rate_per_sec: float = 10.0):
        self.capacity = capacity
        self.refill_rate = refill_rate_per_sec
        self.tenants: Dict[str, Tuple[float, float]] = {}

    def acquire(self, tenant_id: str, cost: float = 1.0) -> RateLimitDecision:
        now = time.time()
        tokens, last_refill = self.tenants.get(tenant_id, (self.capacity, now))

        # Calculate accrued tokens
        delta = now - last_refill
        tokens = min(self.capacity, tokens + delta * self.refill_rate)

        if tokens >= cost:
            tokens -= cost
            self.tenants[tenant_id] = (tokens, now)
            return RateLimitDecision(allowed=True, remaining_tokens=tokens, retry_after_sec=0.0, tenant_id=tenant_id)
        else:
            deficit = cost - tokens
            retry_after = deficit / self.refill_rate
            self.tenants[tenant_id] = (tokens, now)
            return RateLimitDecision(allowed=False, remaining_tokens=tokens, retry_after_sec=round(retry_after, 2), tenant_id=tenant_id)