import threading

# 放在 build_proxy_handler 外层
sem = threading.BoundedSemaphore(2)   # 同时最多 2 个在途上游请求

# 计费请求：先过速率闸门，再过并发闸门
if is_billing:
    bucket.acquire()   # 速率闸门：3s 补 1，排队式
    sem.acquire()      # 并发闸门：最多 2 个在途
try:
    self._send_upstream(url, body, max_retries)
finally:
    if is_billing:
        sem.release()  # 无论成败都要释放槽位

def _send_upstream(self, url, body, max_retries=3):
    for attempt in range(max_retries + 1):
        # ... 构造并发送请求 ...
        except urllib.error.HTTPError as e:
            if e.code == 429 and attempt < max_retries:
                ra = float(e.headers.get("Retry-After", 1))
                time.sleep(ra)      # 就地退避，不抛回客户端
                continue            # 重试
            # 非 429 或重试耗尽：如实回给客户端

import threading
import time

class TokenBucket:
    """固定速率令牌桶: capacity 容量, rate 每秒补充量。线程安全。"""
    def __init__(self, capacity: float, rate: float):
        self.capacity = capacity
        self.rate = rate
        self.tokens = float(capacity)
        self.last = time.monotonic()
        self.lock = threading.Lock()

    def acquire(self):
        """阻塞直到取到 1 个令牌。"""
        while True:
            with self.lock:
                now = time.monotonic()
                self.tokens = min(self.capacity, self.tokens + (now - self.last) * self.rate)
                self.last = now
                if self.tokens >= 1.0:
                    self.tokens -= 1.0
                    return
                wait = (1.0 - self.tokens) / self.rate
            time.sleep(wait)

import argparse
import json
import os
import sys
import threading
import time
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

DEFAULT_UPSTREAM = "https://developer.amd.com.cn/radeon/api/v1"

def build_proxy_handler(bucket, sem, upstream, api_key, max_retries=3):
    class ProxyHandler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"

        def log_message(self, fmt, *args):
            sys.stderr.write("[proxy %s] %s\n" % (time.strftime("%H:%M:%S"), fmt % args))

        def _send_upstream(self, url, body, max_retries=max_retries):
            """向上游发送请求，遇到 429 自动就地退避重试，切断客户端重试风暴"""
            for attempt in range(max_retries + 1):
                req = urllib.request.Request(url, data=body, method=self.command)
                for h, v in self.headers.items():
                    low = h.lower()
                    if low in ("host", "content-length", "connection", "accept-encoding"):
                        continue
                    req.add_header(h, v)
                if api_key:
                    req.add_header("Authorization", f"Bearer {api_key}")

                try:
                    with urllib.request.urlopen(req, timeout=300) as resp:
                        out = resp.read()
                        self.send_response(resp.status)
                        for h, v in resp.headers.items():
                            if h.lower() in ("content-length", "connection", "transfer-encoding"):
                                continue
                            self.send_header(h, v)
                        self.send_header("Content-Length", str(len(out)))
                        self.end_headers()
                        self.wfile.write(out)
                        return
                except urllib.error.HTTPError as e:
                    if e.code == 429 and attempt < max_retries:
                        ra = float(e.headers.get("Retry-After", 1) or 1)
                        self.log_message("收到 429 process_concurrency，触发就地退避重试 (attempt %d/%d, sleep %.1fs)",
                                         attempt + 1, max_retries, ra)
                        time.sleep(ra)
                        continue
                    out = e.read()
                    self.send_response(e.code)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(out)))
                    self.end_headers()
                    self.wfile.write(out)
                    return
                except Exception as e:
                    payload = json.dumps({"error": {"message": f"proxy upstream error: {e}",
                                                    "type": "proxy_error"}}).encode()
                    self.send_response(502)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(payload)))
                    self.end_headers()
                    self.wfile.write(payload)
                    return

        def _forward(self):
            is_billing = self.command in ("POST", "PUT", "PATCH", "DELETE")

            # 计费写请求：先过速率闸门（令牌桶），再过并发闸门（信号量）
            if is_billing:
                bucket.acquire()  # 速率闸门：3s 补 1，平滑排队
                sem.acquire()     # 并发闸门：限制在途上游连接数

            try:
                length = int(self.headers.get("Content-Length", 0) or 0)
                body = self.rfile.read(length) if length else None

                # 路径归一化：upstream 已包含 /api/v1，剥离请求路径多余的 /v1 前缀
                path = self.path.split("?", 1)[0]
                if path.startswith("/v1/") and not path.startswith("/v1.1/"):
                    path = path[3:]  # /v1/chat/completions -> /chat/completions

                url = upstream.rstrip("/") + path
                self._send_upstream(url, body, max_retries)
            finally:
                if is_billing:
                    sem.release()  # 无论成功失败必释放并发槽位

        def do_GET(self):
            path = self.path.split("?", 1)[0]
            # 健康检查接口
            if path == "/healthz":
                body = b'{"status":"ok"}\n'
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                return

            # 模型探测拦截：本地响应，不透传 AMD —— 解决 429 的核心所在！
            # AMD 对 404 的探测请求也计入 RPM，探测频繁会迅速吃光 20 次额度
            if path in ("/v1/models", "/models", "/api/v1/models", "/api/tags",
                        "/v1/props", "/props", "/version",
                        "/v1/models/DeepSeek-V4-Flash", "/api/show"):
                body = json.dumps({"object": "list", "data": [
                    {"id": "DeepSeek-V4-Flash", "object": "model"},
                    {"id": "GLM-5.2", "object": "model"},
                    {"id": "MinerU2.5-Pro", "object": "model"},
                ]}).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                return
            self._forward()

        do_POST = do_PUT = do_DELETE = do_PATCH = _forward

    return ProxyHandler


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8077)
    ap.add_argument("--upstream", default=DEFAULT_UPSTREAM)
    ap.add_argument("--rpm", type=float, default=20.0, help="每分钟请求数限额 (默认 20)")
    ap.add_argument("--concurrency", type=int, default=2, help="同时在途最大并发请求数 (默认 2)")
    ap.add_argument("--max-retries", type=int, default=3, help="上游 429 就地退避最大重试次数 (默认 3)")
    ap.add_argument("--api-key", default=None, help="缺省时读取环境变量 CUSTOM_API_KEY")
    args = ap.parse_args()

    api_key = args.api_key or os.environ.get("CUSTOM_API_KEY", "")
    bucket = TokenBucket(capacity=args.rpm, rate=args.rpm / 60.0)
    sem = threading.BoundedSemaphore(args.concurrency)
    handler = build_proxy_handler(bucket, sem, args.upstream, api_key, max_retries=args.max_retries)
    server = ThreadingHTTPServer(("127.0.0.1", args.port), handler)
    print(f"[proxy] listening on http://127.0.0.1:{args.port} -> {args.upstream} (rpm={args.rpm:g}, conc={args.concurrency}, retry={args.max_retries})")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        server.shutdown()

if __name__ == "__main__":
    main()

import time, random, threading, requests

class RateLimiter:
    def __init__(self, rpm=15, max_concurrency=3):  # 留余量，设置 15 RPM 避免触顶
        self.interval = 60.0 / rpm
        self.semaphore = threading.BoundedSemaphore(max_concurrency)  # 控制并发槽
        self.lock = threading.Lock()
        self.next_time = 0.0

    def wait(self):
        with self.lock:
            now = time.monotonic()
            if now < self.next_time:
                time.sleep(self.next_time - now)
            self.next_time = max(time.monotonic(), self.next_time) + self.interval
        self.semaphore.acquire()  # 占用并发槽

    def release(self):
        self.semaphore.release()

limiter = RateLimiter(rpm=15, max_concurrency=3)

def call_api(payload, max_retries=6):
    last_err = None
    for attempt in range(max_retries):
        limiter.wait()
        try:
            resp = requests.post(API_URL, json=payload, headers=HEADERS, timeout=120)
            if resp.status_code == 429:
                retry_after = resp.headers.get("Retry-After")
                delay = float(retry_after) if retry_after else min(2 ** attempt, 60) + random.uniform(0, 1)
                last_err = f"429: {resp.text[:200]}"
                time.sleep(delay)
                continue
            resp.raise_for_status()
            return resp
        except requests.exceptions.Timeout:
            last_err = "timeout"
            time.sleep(min(2 ** attempt, 30))
        finally:
            limiter.release()  # 确保释放并发槽
    raise RuntimeError(f"重试多次后仍失败: {last_err}")