mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
|
"""Rolling-window cost tracking for the Account Pool.
|
||
|
|
|
||
|
|
Each key has a configurable token budget per rolling window (e.g. 5 hours).
|
||
|
|
When the budget is exhausted the key is marked as unschedulable by the pool
|
||
|
|
manager. A "soft threshold" (default 80 %) causes the pool to *prefer*
|
||
|
|
other keys but still allows traffic if no alternatives exist.
|
||
|
|
|
||
|
|
All state is stored in Redis sorted sets via :mod:`redis_ops`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import TYPE_CHECKING
|
||
|
|
|
||
|
|
from src.services.provider.pool import redis_ops
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from src.services.provider.pool.config import PoolConfig
|
||
|
|
|
||
|
|
|
||
|
|
async def record_usage(
|
||
|
|
provider_id: str,
|
||
|
|
key_id: str,
|
||
|
|
tokens: int,
|
||
|
|
config: PoolConfig,
|
||
|
|
) -> None:
|
||
|
|
"""Record *tokens* used by *key_id* in the rolling cost window."""
|
||
|
|
if tokens <= 0:
|
||
|
|
return
|
||
|
|
if config.cost_limit_per_key_tokens is None:
|
||
|
|
return # cost tracking disabled
|
||
|
|
await redis_ops.add_cost_entry(provider_id, key_id, tokens, config.cost_window_seconds)
|
||
|
|
|
||
|
|
|
||
|
|
async def get_window_usage(
|
||
|
|
provider_id: str,
|
||
|
|
key_id: str,
|
||
|
|
config: PoolConfig,
|
||
|
|
) -> int:
|
||
|
|
"""Return total tokens used by *key_id* within the current window."""
|
||
|
|
return await redis_ops.get_cost_window_total(provider_id, key_id, config.cost_window_seconds)
|
||
|
|
|
||
|
|
|
||
|
|
async def is_at_limit(
|
||
|
|
provider_id: str,
|
||
|
|
key_id: str,
|
||
|
|
config: PoolConfig,
|
||
|
|
) -> bool:
|
||
|
|
"""Return ``True`` if the key has exhausted its budget."""
|
||
|
|
if config.cost_limit_per_key_tokens is None:
|
||
|
|
return False
|
||
|
|
total = await get_window_usage(provider_id, key_id, config)
|
||
|
|
return total >= config.cost_limit_per_key_tokens
|
||
|
|
|
||
|
|
|
||
|
|
async def is_approaching_limit(
|
||
|
|
provider_id: str,
|
||
|
|
key_id: str,
|
||
|
|
config: PoolConfig,
|
||
|
|
) -> bool:
|
||
|
|
"""Return ``True`` if the key is above the soft threshold."""
|
||
|
|
if config.cost_limit_per_key_tokens is None:
|
||
|
|
return False
|
||
|
|
total = await get_window_usage(provider_id, key_id, config)
|
||
|
|
threshold = config.cost_limit_per_key_tokens * config.cost_soft_threshold_percent / 100
|
||
|
|
return total >= threshold
|