119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
"""用量查询 — Tavily 剩余额度 + 百度 VDB 免费额度"""
|
||
|
||
import time
|
||
import hmac
|
||
import hashlib
|
||
import datetime
|
||
import requests
|
||
|
||
|
||
def check_tavily_usage(api_key: str) -> dict:
|
||
if not api_key:
|
||
return {'error': '未配置 API Key', 'available': False}
|
||
|
||
try:
|
||
resp = requests.get(
|
||
'https://api.tavily.com/usage',
|
||
headers={'Authorization': f'Bearer {api_key}'},
|
||
timeout=10,
|
||
)
|
||
if resp.status_code == 429:
|
||
return {'error': '查询过于频繁（10分钟内限10次）', 'available': False}
|
||
if resp.status_code != 200:
|
||
return {'error': f'API 返回 {resp.status_code}', 'available': False}
|
||
|
||
data = resp.json()
|
||
key_data = data.get('key', {})
|
||
acct = data.get('account', {})
|
||
|
||
usage = key_data.get('usage', 0) or 0
|
||
limit = key_data.get('limit', 0) or 0
|
||
remaining = max(limit - usage, 0) if limit else 0
|
||
|
||
return {
|
||
'available': True,
|
||
'key': {'usage': usage, 'limit': limit, 'remaining': remaining},
|
||
'account': {
|
||
'plan_limit': acct.get('plan_limit', 0) or 0,
|
||
'search_usage': acct.get('search_usage', 0) or 0,
|
||
'paygo_usage': acct.get('paygo_usage', 0) or 0,
|
||
},
|
||
}
|
||
except requests.exceptions.Timeout:
|
||
return {'error': '请求超时', 'available': False}
|
||
except requests.exceptions.RequestException as e:
|
||
return {'error': f'网络错误: {e}', 'available': False}
|
||
|
||
|
||
def _bce_sign(access_key: str, secret_key: str, method: str, path: str,
|
||
headers: dict, params: dict = None) -> str:
|
||
timestamp = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ')
|
||
expiration = 1800
|
||
|
||
auth_string = f'bce-auth-v1/{access_key}/{timestamp}/{expiration}'
|
||
signing_key = hmac.new(
|
||
secret_key.encode('utf-8'),
|
||
auth_string.encode('utf-8'),
|
||
hashlib.sha256,
|
||
).hexdigest()
|
||
|
||
signed_headers = sorted(headers.keys(), key=lambda k: k.lower())
|
||
canonical_headers = ''.join(
|
||
f'{k.lower()}:{headers[k].strip()}\n' for k in signed_headers
|
||
)
|
||
canonical_uri = path
|
||
|
||
if params:
|
||
canonical_query = '&'.join(
|
||
f'{k}={v}' for k, v in sorted(params.items())
|
||
)
|
||
else:
|
||
canonical_query = ''
|
||
|
||
canonical_request = f'{method}\n{canonical_uri}\n{canonical_query}\n{canonical_headers}'
|
||
|
||
signature = hmac.new(
|
||
signing_key.encode('utf-8'),
|
||
canonical_request.encode('utf-8'),
|
||
hashlib.sha256,
|
||
).hexdigest()
|
||
|
||
return (f'bce-auth-v1/{access_key}/{timestamp}/{expiration}/'
|
||
f'{";".join(k.lower() for k in signed_headers)}/{signature}')
|
||
|
||
|
||
def check_baidu_vdb_quota(access_key: str, secret_key: str) -> dict:
|
||
if not access_key or not secret_key:
|
||
return {'error': '未配置 VDB Access Key / Secret Key', 'available': False}
|
||
|
||
host = 'vdb.bj.baidubce.com'
|
||
path = '/v1/vdb/instance/freeQuota'
|
||
method = 'GET'
|
||
|
||
headers = {
|
||
'host': host,
|
||
}
|
||
auth = _bce_sign(access_key, secret_key, method, path, headers)
|
||
headers['Authorization'] = auth
|
||
headers['Content-Type'] = 'application/json'
|
||
|
||
try:
|
||
resp = requests.get(
|
||
f'https://{host}{path}',
|
||
headers=headers,
|
||
timeout=10,
|
||
)
|
||
if resp.status_code != 200:
|
||
return {'error': f'API 返回 {resp.status_code}', 'available': False}
|
||
|
||
data = resp.json()
|
||
free_quota = data.get('freeQuota', 0)
|
||
return {
|
||
'available': True,
|
||
'freeQuota': free_quota,
|
||
}
|
||
except requests.exceptions.Timeout:
|
||
return {'error': '请求超时', 'available': False}
|
||
except requests.exceptions.RequestException as e:
|
||
return {'error': f'网络错误: {e}', 'available': False}
|