chaos
b98010591e
- 新增 token 估算函数 (_estimate_tokens, _estimate_messages_tokens) - 新增滑动窗口截断 (_trim_messages): 保留 system/tool 消息, 从最旧对话消息开始丢弃 - 新增单条消息中间截断 (_shorten_message_content): 保留头尾, 迭代逼近目标 token 数 - 新增 _maybe_trim_context: 集成到代理逻辑, 请求体超限时自动截断 - 截断目标: 上游 192K token 限制的 73.5% (~141K), 留出响应空间 - 测试验证: 202K token 请求自动截断至 125K, 上游返回 200
778 lines
29 KiB
Python
Executable File
778 lines
29 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
华为云 Token 动态网关
|
||
- 6小时缓存机制
|
||
- 支持内存扫描自动刷新
|
||
- HUAWEI_TOKEN 环境变量最高优先级
|
||
- Token 持久化到 /etc/huawei-gateway.env
|
||
- SSE 流式转发
|
||
- 连接池 (20 连接)
|
||
- 401 自动重试
|
||
- 兼容生产环境 (Waitress 32线程 / Gunicorn)
|
||
"""
|
||
import os
|
||
import re
|
||
import sys
|
||
import json
|
||
import time
|
||
import logging
|
||
import threading
|
||
import traceback
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
from flask import Flask, request, Response
|
||
|
||
# 尝试导入 requests，失败则给出明确提示
|
||
try:
|
||
import requests
|
||
except ImportError:
|
||
print("错误：缺少 requests 模块。请运行: pip install requests")
|
||
sys.exit(1)
|
||
|
||
# ================= 配置 =================
|
||
CACHE_TTL = 19800 # 5.5 小时（安全线）
|
||
MAX_WORKERS = 8 # 内存扫描线程数
|
||
MAX_MEM_SEGMENT = 200 * 1024 * 1024 # 单段最大扫描 200MB
|
||
TOKEN_PATTERN = re.compile(b'Bearer ([A-Za-z0-9+/=_-]{100,})')
|
||
TARGET_HOST = 'tokenhub.developer.huaweicloud.com'
|
||
|
||
# ================= 请求体限制配置 =================
|
||
UPSTREAM_BODY_LIMIT = 1200 * 1024 # 上游 APIG 请求体限制 ~1.2MB (实测边界1260KB)
|
||
UPSTREAM_TIMEOUT_MIN = 60 # 小请求超时 60s
|
||
UPSTREAM_TIMEOUT_MAX = 300 # 大请求超时 300s
|
||
RETRY_ON_429 = 2 # 429限流重试次数
|
||
RETRY_ON_504 = 1 # 504超时重试次数
|
||
|
||
# ================= 上下文自动截断配置 =================
|
||
# 参考 LiteLLM trim_messages 方案: 滑动窗口 + system/tool 保留 + 中间截断
|
||
MAX_CONTEXT_TOKENS = 196608 # 上游实测上限 192K tokens
|
||
RESPONSE_BUDGET = 8192 # 预留 8K tokens 给回复
|
||
TRIM_RATIO = 0.75 # 截断到可用空间的 75%
|
||
MAX_TRIM_ATTEMPTS = 5 # 单条消息最大截断尝试次数
|
||
ENABLE_AUTO_TRIM = True # 是否启用自动截断
|
||
|
||
# ================= 并发配置 =================
|
||
POOL_CONNECTIONS = 64 # 连接池大小（须 ≥ Waitress 线程数）
|
||
POOL_MAXSIZE = 64 # 单主机最大连接数
|
||
WAITRESS_THREADS = 64 # Waitress 工作线程数
|
||
SSE_CHUNK_SIZE = 4096 # SSE 流式转发块大小（越小首 token 延迟越低）
|
||
|
||
# ================= 日志 =================
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s [%(levelname)s] %(message)s',
|
||
handlers=[
|
||
logging.StreamHandler(sys.stdout)
|
||
]
|
||
)
|
||
logger = logging.getLogger('huawei-gateway')
|
||
|
||
# ================= 缓存 =================
|
||
class TokenCache:
|
||
def __init__(self):
|
||
self._token = None
|
||
self._expires_at = 0
|
||
self._lock = threading.RLock()
|
||
self._last_scan = 0
|
||
self._scan_interval = 60 # 扫描间隔最小 60 秒
|
||
self._blacklist = set() # 已失效的 token 指纹
|
||
|
||
def _fingerprint(self, token):
|
||
"""取 token 前 16 + 后 16 字符做指纹"""
|
||
if len(token) <= 32:
|
||
return token
|
||
return token[:16] + token[-16:]
|
||
|
||
def get(self):
|
||
with self._lock:
|
||
now = time.time()
|
||
if self._token and now < self._expires_at:
|
||
return self._token
|
||
return None
|
||
|
||
def set(self, token, ttl=CACHE_TTL):
|
||
with self._lock:
|
||
self._token = token
|
||
self._expires_at = time.time() + ttl
|
||
self._last_scan = time.time()
|
||
self._blacklist.discard(self._fingerprint(token))
|
||
|
||
def blacklist_current(self):
|
||
"""将当前 token 加入黑名单"""
|
||
with self._lock:
|
||
if self._token:
|
||
self._blacklist.add(self._fingerprint(self._token))
|
||
self._token = None
|
||
self._expires_at = 0
|
||
|
||
def is_blacklisted(self, token):
|
||
with self._lock:
|
||
return self._fingerprint(token) in self._blacklist
|
||
|
||
def is_scan_cooldown(self):
|
||
with self._lock:
|
||
return (time.time() - self._last_scan) < self._scan_interval
|
||
|
||
def clear_blacklist(self):
|
||
"""清空黑名单"""
|
||
with self._lock:
|
||
self._blacklist.clear()
|
||
|
||
def get_expires_in(self):
|
||
"""返回 token 剩余有效期（秒）"""
|
||
with self._lock:
|
||
return max(0, self._expires_at - time.time())
|
||
|
||
def get_blacklist_count(self):
|
||
"""返回黑名单大小"""
|
||
with self._lock:
|
||
return len(self._blacklist)
|
||
|
||
def fingerprint(self, token):
|
||
"""公开方法：获取 token 指纹"""
|
||
with self._lock:
|
||
return self._fingerprint(token)
|
||
|
||
def clear(self):
|
||
with self._lock:
|
||
self._token = None
|
||
self._expires_at = 0
|
||
|
||
|
||
cache = TokenCache()
|
||
|
||
# ================= 内存扫描 =================
|
||
def scan_pid_mem(pid):
|
||
"""扫描单个进程的内存寻找 Token"""
|
||
maps_path = f'/proc/{pid}/maps'
|
||
mem_path = f'/proc/{pid}/mem'
|
||
|
||
if not os.path.exists(maps_path) or not os.path.exists(mem_path):
|
||
return None
|
||
|
||
try:
|
||
with open(maps_path, 'r') as f:
|
||
for line in f:
|
||
parts = line.split()
|
||
if len(parts) < 2:
|
||
continue
|
||
perms = parts[1]
|
||
if 'r' not in perms or 'w' not in perms:
|
||
continue
|
||
|
||
addrs = parts[0].split('-')
|
||
if len(addrs) != 2:
|
||
continue
|
||
|
||
start = int(addrs[0], 16)
|
||
end = int(addrs[1], 16)
|
||
size = end - start
|
||
|
||
if size > MAX_MEM_SEGMENT or size < 1024:
|
||
continue
|
||
|
||
try:
|
||
with open(mem_path, 'rb') as mem:
|
||
mem.seek(start)
|
||
chunk_size = 64 * 1024
|
||
remaining = size
|
||
while remaining > 0:
|
||
to_read = min(chunk_size, remaining)
|
||
data = mem.read(to_read)
|
||
if not data:
|
||
break
|
||
|
||
for match in TOKEN_PATTERN.finditer(data):
|
||
token = match.group(1).decode('ascii', errors='replace')
|
||
if len(token) > 200:
|
||
return token
|
||
|
||
remaining -= len(data)
|
||
except (PermissionError, OSError, ValueError):
|
||
continue
|
||
|
||
except (PermissionError, OSError, ProcessLookupError):
|
||
pass
|
||
return None
|
||
|
||
|
||
def find_token_in_memory():
|
||
"""在所有进程中扫描 Token"""
|
||
# 快速路径：缓存有效直接返回（避免每次请求都查环境变量/文件）
|
||
cached = cache.get()
|
||
if cached:
|
||
return cached
|
||
|
||
# HUAWEI_TOKEN 环境变量优先级最高
|
||
env_token = os.environ.get('HUAWEI_TOKEN', '').strip()
|
||
if env_token and len(env_token) > 200 and not cache.is_blacklisted(env_token):
|
||
cache.set(env_token)
|
||
logger.info("Token 从 HUAWEI_TOKEN 环境变量加载")
|
||
return env_token
|
||
|
||
# 从持久化文件加载
|
||
env_file = '/etc/huawei-gateway.env'
|
||
if os.path.isfile(env_file):
|
||
try:
|
||
with open(env_file, 'r') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line.startswith('HUAWEI_TOKEN='):
|
||
file_token = line.split('=', 1)[1].strip().strip('"').strip("'")
|
||
if file_token and len(file_token) > 200 and not cache.is_blacklisted(file_token):
|
||
cache.set(file_token)
|
||
logger.info("Token 从持久化文件加载")
|
||
return file_token
|
||
break
|
||
except (OSError, IOError):
|
||
pass
|
||
|
||
if cache.is_scan_cooldown():
|
||
return cache.get()
|
||
|
||
try:
|
||
pids = [pid for pid in os.listdir('/proc') if pid.isdigit()]
|
||
except OSError:
|
||
logger.error("无法访问 /proc 目录")
|
||
return None
|
||
|
||
# 优先扫描常见进程
|
||
priority_pids = []
|
||
other_pids = []
|
||
|
||
for pid in pids:
|
||
try:
|
||
exe_path = os.readlink(f'/proc/{pid}/exe')
|
||
if any(x in exe_path for x in ['python', 'node', 'java', 'chrome', 'electron']):
|
||
priority_pids.append(pid)
|
||
else:
|
||
other_pids.append(pid)
|
||
except (OSError, PermissionError):
|
||
other_pids.append(pid)
|
||
|
||
all_pids = priority_pids + other_pids
|
||
|
||
found_tokens = []
|
||
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||
futures = {executor.submit(scan_pid_mem, pid): pid for pid in all_pids}
|
||
for future in as_completed(futures):
|
||
try:
|
||
token = future.result(timeout=5)
|
||
if token and not cache.is_blacklisted(token):
|
||
found_tokens.append((token, futures[future]))
|
||
except Exception:
|
||
continue
|
||
|
||
for token, pid in found_tokens:
|
||
cache.set(token)
|
||
logger.info(f"Token 已刷新 (来源 PID: {pid})")
|
||
return token
|
||
|
||
return cache.get()
|
||
|
||
|
||
# ================= HTTP 会话池 =================
|
||
# max_retries=0：禁用 urllib3 自动重试，由 proxy 手动控制重试逻辑（避免双重重试）
|
||
http_session = requests.Session()
|
||
adapter = requests.adapters.HTTPAdapter(
|
||
pool_connections=POOL_CONNECTIONS,
|
||
pool_maxsize=POOL_MAXSIZE,
|
||
max_retries=0
|
||
)
|
||
http_session.mount('https://', adapter)
|
||
http_session.mount('http://', adapter)
|
||
|
||
|
||
# ================= 上下文自动截断 (参考 LiteLLM trim_messages) =================
|
||
def _estimate_tokens(text):
|
||
"""粗略估算文本的 token 数。
|
||
英文约 4 字符/token，中文约 1.5 字符/token，混合取 ~3 字符/token。
|
||
"""
|
||
if not text:
|
||
return 0
|
||
# 统计中文字符比例
|
||
chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
|
||
total_chars = len(text)
|
||
if total_chars == 0:
|
||
return 0
|
||
chinese_ratio = chinese_chars / total_chars
|
||
# 中文多的文本 token 密度更高
|
||
chars_per_token = 4.0 - 2.5 * chinese_ratio # 纯英文=4, 纯中文=1.5
|
||
return max(1, int(total_chars / chars_per_token))
|
||
|
||
|
||
def _estimate_message_tokens(msg):
|
||
"""估算单条消息的 token 数 (含 role 开销 ~4 tokens)"""
|
||
content = msg.get('content', '')
|
||
if isinstance(content, list):
|
||
# 多模态消息: 提取文本部分
|
||
text_parts = []
|
||
for part in content:
|
||
if isinstance(part, dict):
|
||
if part.get('type') == 'text':
|
||
text_parts.append(part.get('text', ''))
|
||
elif part.get('type') == 'image_url':
|
||
text_parts.append('') # 图片按 ~512 token 估算
|
||
elif isinstance(part, str):
|
||
text_parts.append(part)
|
||
text = ' '.join(text_parts)
|
||
tokens = _estimate_tokens(text) + 512 * sum(
|
||
1 for p in content if isinstance(p, dict) and p.get('type') == 'image_url'
|
||
)
|
||
elif isinstance(content, str):
|
||
tokens = _estimate_tokens(content)
|
||
else:
|
||
tokens = _estimate_tokens(str(content))
|
||
# function_call / tool_calls 额外 token
|
||
if 'function_call' in msg or 'tool_calls' in msg:
|
||
tokens += _estimate_tokens(json.dumps(msg.get('tool_calls', msg.get('function_call', ''))))
|
||
return tokens + 4 # role 开销
|
||
|
||
|
||
def _estimate_messages_tokens(messages):
|
||
"""估算消息列表的总 token 数"""
|
||
return sum(_estimate_message_tokens(m) for m in messages)
|
||
|
||
|
||
def _shorten_message_content(content, target_tokens):
|
||
"""从中间截断消息内容，保留头尾。迭代逼近目标 token 数。"""
|
||
if not isinstance(content, str):
|
||
return content
|
||
if not content:
|
||
return content
|
||
|
||
current_tokens = _estimate_tokens(content)
|
||
if current_tokens <= target_tokens:
|
||
return content
|
||
|
||
marker = "\n...[truncated]...\n"
|
||
for _ in range(MAX_TRIM_ATTEMPTS):
|
||
current_tokens = _estimate_tokens(content)
|
||
if current_tokens <= target_tokens:
|
||
break
|
||
# 保守比例: 留 90% 空间避免截断标记导致超限
|
||
ratio = (target_tokens * 0.9) / current_tokens
|
||
new_length = max(10, int(len(content) * ratio))
|
||
half = new_length // 2
|
||
content = content[:half] + marker + content[-half:]
|
||
|
||
return content
|
||
|
||
|
||
def _trim_messages(messages, max_tokens):
|
||
"""滑动窗口截断消息列表。
|
||
|
||
策略 (参考 LiteLLM):
|
||
1. 分离 system 消息 (始终保留，超限从中间截断)
|
||
2. 分离末尾 tool 消息 (始终保留)
|
||
3. 对话消息从最新向最旧遍历，超限时丢弃最旧
|
||
4. 单条消息超限时从中间截断
|
||
"""
|
||
if not messages:
|
||
return messages
|
||
|
||
# 分离 system 消息
|
||
system_messages = [m for m in messages if m.get('role') == 'system']
|
||
non_system = [m for m in messages if m.get('role') != 'system']
|
||
|
||
# 分离末尾连续的 tool 消息
|
||
tool_messages = []
|
||
for m in reversed(non_system):
|
||
if m.get('role') != 'tool':
|
||
break
|
||
tool_messages.append(m)
|
||
tool_messages.reverse()
|
||
conversation = non_system[:len(non_system) - len(tool_messages)] if tool_messages else non_system
|
||
|
||
# 计算 system + tool 的 token 开销
|
||
system_tokens = _estimate_messages_tokens(system_messages)
|
||
tool_tokens = _estimate_messages_tokens(tool_messages)
|
||
overhead = system_tokens + tool_tokens
|
||
|
||
# 如果 system 消息本身就超限，截断 system
|
||
available = max_tokens - tool_tokens
|
||
if available <= 0:
|
||
# tool 消息本身就超限了，只能尽力返回
|
||
return system_messages[:1] + tool_messages if system_messages else tool_messages
|
||
|
||
if system_tokens > available * 0.5:
|
||
# system 消息占太多，从中间截断每条 system 消息
|
||
target_system_tokens = int(available * 0.3)
|
||
for m in system_messages:
|
||
current = _estimate_message_tokens(m)
|
||
if current > target_system_tokens // len(system_messages):
|
||
m['content'] = _shorten_message_content(
|
||
m.get('content', ''), target_system_tokens // max(1, len(system_messages))
|
||
)
|
||
system_tokens = _estimate_messages_tokens(system_messages)
|
||
|
||
# 剩余给对话消息的空间
|
||
conv_budget = max_tokens - system_tokens - tool_tokens
|
||
if conv_budget <= 0:
|
||
return system_messages + tool_messages
|
||
|
||
# 从最新向最旧遍历，滑动窗口
|
||
final_conv = []
|
||
used = 0
|
||
for msg in reversed(conversation):
|
||
msg_tokens = _estimate_message_tokens(msg)
|
||
if used + msg_tokens <= conv_budget:
|
||
final_conv.insert(0, msg)
|
||
used += msg_tokens
|
||
else:
|
||
# 尝试截断这条消息
|
||
remaining = conv_budget - used
|
||
if remaining > 50 and 'function_call' not in msg and 'tool_calls' not in msg:
|
||
trimmed = dict(msg)
|
||
trimmed['content'] = _shorten_message_content(msg.get('content', ''), remaining - 4)
|
||
if _estimate_message_tokens(trimmed) <= remaining:
|
||
final_conv.insert(0, trimmed)
|
||
used += _estimate_message_tokens(trimmed)
|
||
# 空间不够，停止加入更旧的消息
|
||
break
|
||
|
||
return system_messages + final_conv + tool_messages
|
||
|
||
|
||
def _maybe_trim_context(raw_body):
|
||
"""检查并截断请求体中的消息列表。返回 (new_body, trimmed_info)"""
|
||
if not ENABLE_AUTO_TRIM:
|
||
return raw_body, None
|
||
|
||
try:
|
||
body = json.loads(raw_body)
|
||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||
return raw_body, None
|
||
|
||
messages = body.get('messages')
|
||
if not messages or not isinstance(messages, list):
|
||
return raw_body, None
|
||
|
||
total_tokens = _estimate_messages_tokens(messages)
|
||
target_limit = int((MAX_CONTEXT_TOKENS - RESPONSE_BUDGET) * TRIM_RATIO)
|
||
|
||
if total_tokens <= target_limit:
|
||
return raw_body, None # 无需截断
|
||
|
||
original_count = len(messages)
|
||
trimmed_messages = _trim_messages(messages, target_limit)
|
||
new_tokens = _estimate_messages_tokens(trimmed_messages)
|
||
|
||
if len(trimmed_messages) == original_count and new_tokens >= total_tokens:
|
||
return raw_body, None # 截断没效果
|
||
|
||
body['messages'] = trimmed_messages
|
||
new_body = json.dumps(body, ensure_ascii=False).encode('utf-8')
|
||
|
||
info = {
|
||
'original_messages': original_count,
|
||
'trimmed_messages': len(trimmed_messages),
|
||
'original_tokens_est': total_tokens,
|
||
'trimmed_tokens_est': new_tokens,
|
||
'target_limit': target_limit,
|
||
}
|
||
logger.info(
|
||
f"上下文截断: {original_count}→{len(trimmed_messages)} 条消息, "
|
||
f"~{total_tokens}→~{new_tokens} tokens (目标≤{target_limit})"
|
||
)
|
||
return new_body, info
|
||
|
||
|
||
# ================= Flask 应用 =================
|
||
app = Flask(__name__)
|
||
|
||
|
||
# ================= 全局请求日志（捕获所有请求，包括404） =================
|
||
@app.before_request
|
||
def log_every_request():
|
||
body_preview = ""
|
||
if request.method in ('POST', 'PUT', 'PATCH') and request.content_length and request.content_length < 2048:
|
||
body_preview = request.get_data()[:200].decode('utf-8', errors='replace')
|
||
logger.info(f">>> {request.method} {request.full_path} | body={request.content_length or 0}bytes | from={request.remote_addr} | {body_preview}")
|
||
|
||
|
||
@app.route('/health')
|
||
def health():
|
||
"""健康检查端点"""
|
||
token = cache.get()
|
||
return {
|
||
"status": "healthy",
|
||
"token_cached": token is not None,
|
||
"token_expires_in": cache.get_expires_in(),
|
||
"blacklisted": cache.get_blacklist_count()
|
||
}, 200
|
||
|
||
|
||
@app.route('/set_token', methods=['POST'])
|
||
def set_token():
|
||
"""手动注入有效 Token"""
|
||
data = request.get_json(force=True, silent=True) if request.is_json else {}
|
||
token = data.get('token', '')
|
||
if not token or len(token) < 100:
|
||
return {"error": "请提供有效的 token"}, 400
|
||
cache.set(token)
|
||
# 持久化到文件以便重启后恢复
|
||
env_file = '/etc/huawei-gateway.env'
|
||
try:
|
||
with open(env_file, 'w') as f:
|
||
f.write(f'HUAWEI_TOKEN={token}\n')
|
||
os.chmod(env_file, 0o600)
|
||
except (OSError, IOError):
|
||
pass
|
||
logger.info("Token 已手动注入并持久化")
|
||
return {"status": "ok", "token_fingerprint": cache.fingerprint(token)}, 200
|
||
|
||
|
||
# ================= 通用上游请求 =================
|
||
def _forward_upstream(method, target_url, headers, raw_body, cookies, timeout):
|
||
"""向上游发起请求并返回 response 对象（stream=True）"""
|
||
return http_session.request(
|
||
method=method,
|
||
url=target_url,
|
||
headers=headers,
|
||
data=raw_body,
|
||
cookies=cookies,
|
||
allow_redirects=False,
|
||
timeout=timeout,
|
||
stream=True
|
||
)
|
||
|
||
|
||
def _build_response(resp):
|
||
"""根据上游 response 构建转发给客户端的 Flask Response"""
|
||
# 过滤 hop-by-hop 头和压缩编码头
|
||
skip_headers = {'transfer-encoding', 'content-encoding', 'content-length',
|
||
'connection', 'keep-alive', 'upgrade'}
|
||
response_headers = [(k, v) for k, v in resp.headers.items() if k.lower() not in skip_headers]
|
||
|
||
# SSE 流式转发
|
||
content_type = resp.headers.get('Content-Type', '')
|
||
if 'text/event-stream' in content_type or resp.headers.get('Transfer-Encoding', '') == 'chunked':
|
||
def sse_stream():
|
||
try:
|
||
for chunk in resp.iter_content(chunk_size=SSE_CHUNK_SIZE):
|
||
if chunk:
|
||
yield chunk
|
||
finally:
|
||
resp.close()
|
||
|
||
return Response(
|
||
sse_stream(),
|
||
status=resp.status_code,
|
||
headers=response_headers,
|
||
direct_passthrough=True
|
||
)
|
||
else:
|
||
# 非流式响应：读取完整内容
|
||
content = resp.content
|
||
resp.close()
|
||
|
||
# 记录上游非 200 响应（在读取 content 之后，避免提前消费流）
|
||
if resp.status_code != 200:
|
||
try:
|
||
logger.warning(f"上游返回 {resp.status_code}: {content[:500].decode('utf-8', errors='replace')}")
|
||
except Exception:
|
||
logger.warning(f"上游返回 {resp.status_code}")
|
||
|
||
# 华为云 ModelArts 错误 → OpenAI 标准格式
|
||
if resp.status_code >= 400:
|
||
try:
|
||
err = json.loads(content)
|
||
if 'error_code' in err and 'error_msg' in err:
|
||
openai_err = {
|
||
"error": {
|
||
"message": err.get('error_msg', ''),
|
||
"type": err.get('error', {}).get('type', 'server_error'),
|
||
"code": err.get('error_code', ''),
|
||
"param": None
|
||
}
|
||
}
|
||
content = json.dumps(openai_err).encode('utf-8')
|
||
response_headers = [(k, v) for k, v in response_headers if k.lower() != 'content-length']
|
||
except Exception:
|
||
pass
|
||
|
||
return Response(
|
||
content,
|
||
status=resp.status_code,
|
||
headers=response_headers
|
||
)
|
||
|
||
|
||
@app.route('/v1/<path:subpath>', methods=['POST', 'GET', 'OPTIONS', 'PUT', 'DELETE'])
|
||
@app.route('/v2/<path:subpath>', methods=['POST', 'GET', 'OPTIONS', 'PUT', 'DELETE'])
|
||
def proxy(subpath):
|
||
if request.method == 'OPTIONS':
|
||
return Response(status=200, headers={
|
||
'Access-Control-Allow-Origin': '*',
|
||
'Access-Control-Allow-Methods': 'GET, POST, PUT, DELETE, OPTIONS',
|
||
'Access-Control-Allow-Headers': 'Content-Type, Authorization'
|
||
})
|
||
|
||
real_token = find_token_in_memory()
|
||
if not real_token:
|
||
logger.error("未在内存中找到华为云 Token")
|
||
return {"error": "未在内存中找到华为云Token，请确保华为云相关应用正在运行"}, 500
|
||
|
||
# 构建请求头
|
||
headers = {}
|
||
for k, v in request.headers:
|
||
kl = k.lower()
|
||
if kl not in ('host', 'content-length', 'connection', 'accept-encoding', 'transfer-encoding'):
|
||
headers[k] = v
|
||
|
||
headers['Authorization'] = f'Bearer {real_token}'
|
||
headers['Host'] = TARGET_HOST
|
||
headers['Connection'] = 'keep-alive'
|
||
|
||
target_url = f'https://{TARGET_HOST}/v2/{subpath}'
|
||
|
||
# ============ /models 请求直接返回，不转发上游 ============
|
||
if subpath in ('models', 'models/'):
|
||
return {
|
||
"object": "list",
|
||
"data": [
|
||
{"id": "glm-5.1", "object": "model", "owned": "zhipu"},
|
||
{"id": "glm-5.2", "object": "model", "owned": "zhipu"},
|
||
]
|
||
}
|
||
|
||
# ============ 请求体处理：大小校验 + 动态超时 ============
|
||
raw_body = request.get_data()
|
||
body_size = len(raw_body)
|
||
|
||
# 请求体超过上游限制，直接返回清晰错误
|
||
if body_size > UPSTREAM_BODY_LIMIT:
|
||
logger.warning(f"请求体超限: {body_size // 1024}KB > {UPSTREAM_BODY_LIMIT // 1024}KB")
|
||
return {
|
||
"error": {
|
||
"message": f"请求体过大({body_size // 1024}KB)，超过API网关限制({UPSTREAM_BODY_LIMIT // 1024}KB)。请减少请求内容长度。",
|
||
"type": "invalid_request_error",
|
||
"code": "content_too_large",
|
||
"param": None
|
||
}
|
||
}, 413
|
||
|
||
# ============ 上下文自动截断 ============
|
||
# 在请求体大小校验之后、转发之前，对 messages 做滑动窗口截断
|
||
if 'chat/completions' in subpath or 'messages' in subpath:
|
||
raw_body, trim_info = _maybe_trim_context(raw_body)
|
||
if trim_info:
|
||
body_size = len(raw_body) # 更新截断后的 body 大小
|
||
# 截断后重新检查是否仍超限
|
||
if body_size > UPSTREAM_BODY_LIMIT:
|
||
logger.warning(f"截断后请求体仍超限: {body_size // 1024}KB")
|
||
return {
|
||
"error": {
|
||
"message": f"上下文截断后请求体仍过大({body_size // 1024}KB)，请减少请求内容长度。",
|
||
"type": "invalid_request_error",
|
||
"code": "content_too_large",
|
||
"param": None
|
||
}
|
||
}, 413
|
||
|
||
# 动态超时：根据请求体大小自动调整
|
||
if body_size > 800 * 1024:
|
||
upstream_timeout = UPSTREAM_TIMEOUT_MAX
|
||
elif body_size > 200 * 1024:
|
||
upstream_timeout = 180
|
||
else:
|
||
upstream_timeout = UPSTREAM_TIMEOUT_MIN
|
||
|
||
try:
|
||
resp = _forward_upstream(request.method, target_url, headers, raw_body, request.cookies, upstream_timeout)
|
||
|
||
# 401 兜底：Token 可能提前过期，加入黑名单后强制刷新重试
|
||
if resp.status_code == 401:
|
||
resp.close()
|
||
logger.warning("收到 401，将当前 Token 加入黑名单并强制刷新...")
|
||
cache.blacklist_current()
|
||
new_token = find_token_in_memory()
|
||
if new_token and new_token != real_token:
|
||
headers['Authorization'] = f'Bearer {new_token}'
|
||
resp = _forward_upstream(request.method, target_url, headers, raw_body, request.cookies, upstream_timeout)
|
||
# 如果新 Token 也 401，清空黑名单避免锁死
|
||
if resp.status_code == 401:
|
||
resp.close()
|
||
logger.warning("新 Token 也 401，清空黑名单避免锁死")
|
||
cache.clear_blacklist()
|
||
new_token2 = find_token_in_memory()
|
||
if new_token2:
|
||
headers['Authorization'] = f'Bearer {new_token2}'
|
||
resp = _forward_upstream(request.method, target_url, headers, raw_body, request.cookies, upstream_timeout)
|
||
|
||
# ============ 429 限流重试 ============
|
||
if resp.status_code == 429:
|
||
for retry_i in range(1, RETRY_ON_429 + 1):
|
||
resp.close()
|
||
wait = retry_i * 5 # 5s, 10s
|
||
logger.warning(f"上游 429 限流，等待{wait}s后重试({retry_i}/{RETRY_ON_429})...")
|
||
time.sleep(wait)
|
||
resp = _forward_upstream(request.method, target_url, headers, raw_body, request.cookies, upstream_timeout)
|
||
if resp.status_code != 429:
|
||
break
|
||
logger.warning("重试仍返回 429")
|
||
|
||
# ============ 504 超时重试 ============
|
||
if resp.status_code == 504:
|
||
for retry_i in range(1, RETRY_ON_504 + 1):
|
||
resp.close()
|
||
wait = retry_i * 3
|
||
logger.warning(f"上游 504 超时，等待{wait}s后重试({retry_i}/{RETRY_ON_504})...")
|
||
time.sleep(wait)
|
||
resp = _forward_upstream(request.method, target_url, headers, raw_body, request.cookies, upstream_timeout)
|
||
if resp.status_code != 504:
|
||
break
|
||
|
||
return _build_response(resp)
|
||
|
||
except requests.exceptions.Timeout:
|
||
logger.error("请求华为云 API 超时")
|
||
return {"error": "网关超时，请稍后重试"}, 504
|
||
except requests.exceptions.ConnectionError:
|
||
logger.error("无法连接到华为云 API")
|
||
return {"error": "无法连接到华为云服务"}, 502
|
||
except Exception as e:
|
||
logger.error(f"网关转发失败: {traceback.format_exc()}")
|
||
return {"error": f"网关转发失败: {str(e)}"}, 500
|
||
|
||
|
||
def main():
|
||
port = int(sys.argv[1]) if len(sys.argv) > 1 else 8080
|
||
host = sys.argv[2] if len(sys.argv) > 2 else '127.0.0.1'
|
||
|
||
# 从持久化文件加载 token
|
||
env_file = '/etc/huawei-gateway.env'
|
||
if os.path.isfile(env_file):
|
||
try:
|
||
with open(env_file, 'r') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line.startswith('HUAWEI_TOKEN=') and 'HUAWEI_TOKEN' not in os.environ:
|
||
val = line.split('=', 1)[1].strip().strip('"').strip("'")
|
||
if val and len(val) > 200:
|
||
os.environ['HUAWEI_TOKEN'] = val
|
||
logger.info("从持久化文件恢复 Token")
|
||
break
|
||
except (OSError, IOError):
|
||
pass
|
||
|
||
# 尝试使用生产级 WSGI 服务器
|
||
try:
|
||
import waitress
|
||
logger.info(f"使用 Waitress 启动网关 ({host}:{port})")
|
||
waitress.serve(app, host=host, port=port, threads=WAITRESS_THREADS)
|
||
except ImportError:
|
||
try:
|
||
import gunicorn.app.wsgiapp
|
||
logger.info(f"使用 Gunicorn 启动网关 ({host}:{port})")
|
||
os.execlp('gunicorn', 'gunicorn', '-w', '4', '-b', f'{host}:{port}', '--access-logfile', '-', 'huawei_gateway:app')
|
||
except (ImportError, OSError):
|
||
logger.warning("未安装 Waitress/Gunicorn，使用 Flask 开发服务器（建议生产环境安装 waitress）")
|
||
logger.info(f"启动网关 ({host}:{port})")
|
||
app.run(host=host, port=port, debug=False, threaded=True)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|