96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
API 成本统计 - 记录调用成本，超预算告警
|
|
"""
|
|
import os
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
COST_FILE = ROOT / "config" / ".api_cost.json"
|
|
BUDGET_DAILY = float(os.getenv("API_COST_BUDGET_DAILY", "0")) # 0=不限制
|
|
BUDGET_MONTHLY = float(os.getenv("API_COST_BUDGET_MONTHLY", "0"))
|
|
|
|
# 单次调用预估成本（元，可按实际调整）
|
|
COST_PER_CALL = {
|
|
"gemini_extract": 0.02,
|
|
"gemini_qa": 0.01,
|
|
"gemini_vision": 0.015,
|
|
"openai_chat": 0.02,
|
|
"openai_embedding": 0.001,
|
|
"qwen_enhance": 0.05,
|
|
}
|
|
|
|
|
|
def _load() -> dict:
|
|
if not COST_FILE.exists():
|
|
return {"daily": {}, "monthly": {}}
|
|
try:
|
|
with open(COST_FILE, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return {"daily": {}, "monthly": {}}
|
|
|
|
|
|
def _save(data: dict):
|
|
COST_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(COST_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def record(service: str, count: int = 1, cost: Optional[float] = None):
|
|
"""记录一次 API 调用"""
|
|
c = cost if cost is not None else COST_PER_CALL.get(service, 0.01) * count
|
|
today = datetime.now().strftime("%Y-%m-%d")
|
|
month = datetime.now().strftime("%Y-%m")
|
|
data = _load()
|
|
data["daily"][today] = data["daily"].get(today, 0) + c
|
|
data["monthly"][month] = data["monthly"].get(month, 0) + c
|
|
_save(data)
|
|
return c
|
|
|
|
|
|
def get_today_cost() -> float:
|
|
today = datetime.now().strftime("%Y-%m-%d")
|
|
data = _load()
|
|
return data["daily"].get(today, 0)
|
|
|
|
|
|
def get_month_cost() -> float:
|
|
month = datetime.now().strftime("%Y-%m")
|
|
data = _load()
|
|
return data["monthly"].get(month, 0)
|
|
|
|
|
|
async def check_budget_alert():
|
|
"""检查是否超预算，超则企微告警"""
|
|
if BUDGET_DAILY <= 0 and BUDGET_MONTHLY <= 0:
|
|
return
|
|
try:
|
|
from config.config import WECHAT_WEBHOOK
|
|
if not WECHAT_WEBHOOK:
|
|
return
|
|
import httpx
|
|
today = get_today_cost()
|
|
month = get_month_cost()
|
|
msg_parts = []
|
|
if BUDGET_DAILY > 0 and today >= BUDGET_DAILY:
|
|
msg_parts.append(f"⚠️ 今日 API 成本 {today:.2f} 元 ≥ 预算 {BUDGET_DAILY} 元")
|
|
if BUDGET_MONTHLY > 0 and month >= BUDGET_MONTHLY:
|
|
msg_parts.append(f"⚠️ 本月 API 成本 {month:.2f} 元 ≥ 预算 {BUDGET_MONTHLY} 元")
|
|
if not msg_parts:
|
|
return
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
await client.post(WECHAT_WEBHOOK, json={
|
|
"msgtype": "markdown",
|
|
"markdown": {"content": "\n".join(msg_parts)}
|
|
})
|
|
logger.info(f"[成本] 告警已发送: {msg_parts}")
|
|
except Exception as e:
|
|
logger.warning(f"[成本] 告警失败: {e}")
|