mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
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
|