init: Search Hub - 统一多搜索引擎聚合服务

hub/quota.py

@@ -0,0 +1,118 @@
+"""用量查询 — 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}