mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10: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)
263 lines
11 KiB
Python
263 lines
11 KiB
Python
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
from src.core.api_format.signature import normalize_signature_key
|
||
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
|
||
from src.services.usage._recording_helpers import (
|
||
build_usage_params,
|
||
deserialize_body_if_json,
|
||
sanitize_request_metadata,
|
||
)
|
||
from src.services.usage._types import UsageCostInfo, UsageRecordParams
|
||
|
||
|
||
class UsageBillingIntegrationMixin:
|
||
"""计费集成方法 -- 准备用量记录的共享逻辑"""
|
||
|
||
@classmethod
|
||
async def _prepare_usage_record(
|
||
cls,
|
||
params: UsageRecordParams,
|
||
) -> tuple[dict[str, Any], float]:
|
||
"""准备用量记录的共享逻辑
|
||
|
||
此方法提取了 record_usage 和 record_usage_async 的公共处理逻辑：
|
||
- 获取费率倍数
|
||
- 计算成本
|
||
- 构建 Usage 参数
|
||
|
||
Args:
|
||
params: 用量记录参数数据类
|
||
|
||
Returns:
|
||
(usage_params 字典, total_cost 总成本)
|
||
"""
|
||
# 计费口径以 Provider 为准（优先 endpoint_api_format）
|
||
billing_api_format: str | None = None
|
||
if params.endpoint_api_format:
|
||
try:
|
||
billing_api_format = normalize_signature_key(str(params.endpoint_api_format))
|
||
except Exception:
|
||
billing_api_format = None
|
||
if billing_api_format is None and params.api_format:
|
||
try:
|
||
billing_api_format = normalize_signature_key(str(params.api_format))
|
||
except Exception:
|
||
billing_api_format = None
|
||
|
||
input_tokens_for_billing = normalize_input_tokens_for_billing(
|
||
billing_api_format,
|
||
params.input_tokens,
|
||
params.cache_read_input_tokens,
|
||
)
|
||
|
||
# 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
|
||
actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
|
||
params.db, params.provider_api_key_id, params.provider_id, billing_api_format
|
||
)
|
||
|
||
metadata = dict(params.metadata or {})
|
||
is_failed_request = params.status_code >= 400 or params.error_message is not None
|
||
|
||
# Helper: compute billing task_type (billing domain)
|
||
billing_task_type = (params.request_type or "").lower()
|
||
if billing_task_type not in {"chat", "cli", "video", "image", "audio"}:
|
||
billing_task_type = "chat"
|
||
|
||
# 使用新计费系统计算费用
|
||
from src.services.billing.service import BillingService
|
||
|
||
request_count = 0 if is_failed_request else 1
|
||
has_cache_tokens = bool(
|
||
params.cache_creation_input_tokens > 0 or params.cache_read_input_tokens > 0
|
||
)
|
||
effective_cache_ttl_minutes = params.cache_ttl_minutes
|
||
|
||
# 主链路很多场景不会显式传 cache_ttl_minutes，这里补全以确保 1h/5m TTL 差异化计价生效。
|
||
if effective_cache_ttl_minutes is None and has_cache_tokens and params.provider_api_key_id:
|
||
try:
|
||
from src.models.database import ProviderAPIKey
|
||
|
||
key_ttl = (
|
||
params.db.query(ProviderAPIKey.cache_ttl_minutes)
|
||
.filter(ProviderAPIKey.id == params.provider_api_key_id)
|
||
.scalar()
|
||
)
|
||
if key_ttl is not None:
|
||
key_ttl_int = int(key_ttl)
|
||
if key_ttl_int >= 0:
|
||
effective_cache_ttl_minutes = key_ttl_int
|
||
except Exception:
|
||
# Best-effort fallback below.
|
||
pass
|
||
|
||
# 无法从 key 获取时，尽量从 5m/1h 细分回推（主要覆盖 Claude cache_creation）。
|
||
if effective_cache_ttl_minutes is None and has_cache_tokens:
|
||
t5m = int(params.cache_creation_input_tokens_5m or 0)
|
||
t1h = int(params.cache_creation_input_tokens_1h or 0)
|
||
if t1h > 0 and t5m == 0:
|
||
effective_cache_ttl_minutes = 60
|
||
elif t5m > 0 and t1h == 0:
|
||
effective_cache_ttl_minutes = 5
|
||
elif t1h > 0:
|
||
# 混合场景优先按长 TTL 计，避免 1h 缓存被按 5m 误计。
|
||
effective_cache_ttl_minutes = 60
|
||
|
||
dims: dict[str, Any] = {
|
||
"input_tokens": input_tokens_for_billing,
|
||
"output_tokens": params.output_tokens,
|
||
"cache_creation_input_tokens": params.cache_creation_input_tokens,
|
||
"cache_read_input_tokens": params.cache_read_input_tokens,
|
||
"request_count": request_count,
|
||
}
|
||
if effective_cache_ttl_minutes is not None:
|
||
dims["cache_ttl_minutes"] = effective_cache_ttl_minutes
|
||
# If tiered pricing is disabled, force first tier by using tier-key=0.
|
||
if not params.use_tiered_pricing:
|
||
dims["total_input_context"] = 0
|
||
|
||
billing = BillingService(params.db)
|
||
result = billing.calculate(
|
||
task_type=billing_task_type,
|
||
model=params.model,
|
||
provider_id=params.provider_id or "",
|
||
dimensions=dims,
|
||
strict_mode=None,
|
||
)
|
||
snap = result.snapshot
|
||
|
||
breakdown = snap.cost_breakdown or {}
|
||
input_cost = float(breakdown.get("input_cost", 0.0))
|
||
output_cost = float(breakdown.get("output_cost", 0.0))
|
||
cache_creation_cost = float(breakdown.get("cache_creation_cost", 0.0))
|
||
cache_read_cost = float(breakdown.get("cache_read_cost", 0.0))
|
||
request_cost = float(breakdown.get("request_cost", 0.0))
|
||
cache_cost = cache_creation_cost + cache_read_cost
|
||
total_cost = float(snap.total_cost or 0.0)
|
||
|
||
rv = snap.resolved_variables or {}
|
||
|
||
def _as_float(v: Any, d: float | None) -> float | None:
|
||
try:
|
||
if v is None:
|
||
return d
|
||
return float(v)
|
||
except Exception:
|
||
return d
|
||
|
||
input_price = _as_float(rv.get("input_price_per_1m"), 0.0) or 0.0
|
||
output_price = _as_float(rv.get("output_price_per_1m"), 0.0) or 0.0
|
||
cache_creation_price = _as_float(rv.get("cache_creation_price_per_1m"), None)
|
||
cache_read_price = _as_float(rv.get("cache_read_price_per_1m"), None)
|
||
request_price = _as_float(rv.get("price_per_request"), None)
|
||
|
||
# Audit snapshot (pruned later by sanitize_request_metadata)
|
||
metadata["billing_snapshot"] = snap.to_dict()
|
||
|
||
# Best-effort prune metadata to reduce DB/memory pressure.
|
||
metadata = sanitize_request_metadata(metadata)
|
||
|
||
# 构建 Usage 参数
|
||
request_body = deserialize_body_if_json(params.request_body)
|
||
provider_request_body = deserialize_body_if_json(params.provider_request_body)
|
||
response_body = deserialize_body_if_json(params.response_body)
|
||
client_response_body = deserialize_body_if_json(params.client_response_body)
|
||
usage_params = build_usage_params(
|
||
db=params.db,
|
||
user=params.user,
|
||
api_key=params.api_key,
|
||
provider=params.provider,
|
||
model=params.model,
|
||
input_tokens=input_tokens_for_billing,
|
||
output_tokens=params.output_tokens,
|
||
cache_creation_input_tokens=params.cache_creation_input_tokens,
|
||
cache_read_input_tokens=params.cache_read_input_tokens,
|
||
cache_creation_input_tokens_5m=params.cache_creation_input_tokens_5m,
|
||
cache_creation_input_tokens_1h=params.cache_creation_input_tokens_1h,
|
||
request_type=params.request_type,
|
||
api_format=params.api_format,
|
||
api_family=params.api_family,
|
||
endpoint_kind=params.endpoint_kind,
|
||
endpoint_api_format=params.endpoint_api_format,
|
||
has_format_conversion=params.has_format_conversion,
|
||
is_stream=params.is_stream,
|
||
response_time_ms=params.response_time_ms,
|
||
first_byte_time_ms=params.first_byte_time_ms,
|
||
status_code=params.status_code,
|
||
error_message=params.error_message,
|
||
metadata=metadata,
|
||
request_headers=params.request_headers,
|
||
request_body=request_body,
|
||
provider_request_headers=params.provider_request_headers,
|
||
provider_request_body=provider_request_body,
|
||
response_headers=params.response_headers,
|
||
client_response_headers=params.client_response_headers,
|
||
response_body=response_body,
|
||
client_response_body=client_response_body,
|
||
request_id=params.request_id,
|
||
provider_id=params.provider_id,
|
||
provider_endpoint_id=params.provider_endpoint_id,
|
||
provider_api_key_id=params.provider_api_key_id,
|
||
status=params.status,
|
||
target_model=params.target_model,
|
||
cost=UsageCostInfo(
|
||
input_cost=input_cost,
|
||
output_cost=output_cost,
|
||
cache_creation_cost=cache_creation_cost,
|
||
cache_read_cost=cache_read_cost,
|
||
cache_cost=cache_cost,
|
||
request_cost=request_cost,
|
||
total_cost=total_cost,
|
||
input_price=input_price,
|
||
output_price=output_price,
|
||
cache_creation_price=cache_creation_price,
|
||
cache_read_price=cache_read_price,
|
||
request_price=request_price,
|
||
actual_rate_multiplier=actual_rate_multiplier,
|
||
is_free_tier=is_free_tier,
|
||
),
|
||
)
|
||
|
||
return usage_params, total_cost
|
||
|
||
@classmethod
|
||
async def _prepare_usage_records_batch(
|
||
cls,
|
||
params_list: list[UsageRecordParams],
|
||
) -> list[tuple[dict[str, Any], float, Exception | None]]:
|
||
"""批量并行准备用量记录（性能优化）
|
||
|
||
并行调用 _prepare_usage_record，提高批量处理效率。
|
||
|
||
Args:
|
||
params_list: 用量记录参数列表
|
||
|
||
Returns:
|
||
列表，每项为 (usage_params, total_cost, exception)
|
||
如果处理成功，exception 为 None
|
||
"""
|
||
import asyncio
|
||
|
||
async def prepare_single(
|
||
params: UsageRecordParams,
|
||
) -> tuple[dict[str, Any], float, Exception | None]:
|
||
try:
|
||
usage_params, total_cost = await cls._prepare_usage_record(params)
|
||
return (usage_params, total_cost, None)
|
||
except Exception as e:
|
||
return ({}, 0.0, e)
|
||
|
||
if not params_list:
|
||
return []
|
||
|
||
# 避免一次性创建过多 task（并且 _prepare_usage_record 内部也可能包含并行调用）
|
||
# 这里采用分批 gather 来限制并发量。
|
||
chunk_size = 50
|
||
results: list[tuple[dict[str, Any], float, Exception | None]] = []
|
||
for i in range(0, len(params_list), chunk_size):
|
||
chunk = params_list[i : i + chunk_size]
|
||
chunk_results = await asyncio.gather(*(prepare_single(p) for p in chunk))
|
||
results.extend(chunk_results)
|
||
return results
|