key_proxy.py

#!/usr/bin/env python3
"""
tchkiller API Key Round-Robin Proxy

透明反向代理，将多个 API key 轮询转发到同一上游 API。
对 tchkiller 完全透明，不丢失上下文，支持 SSE 流式透传。

用法:
    python3 key_proxy.py --keys "sk-1,sk-2,sk-3" --upstream https://open.bigmodel.cn/api/anthropic
    python3 key_proxy.py --keys-file keys.txt --upstream https://open.bigmodel.cn/api/anthropic --port 8090
"""

import argparse
import asyncio
import itertools
import sys
import time
from pathlib import Path

try:
    from aiohttp import web, ClientSession, ClientTimeout
except ImportError:
    print("❌ 需要 aiohttp: pip install aiohttp")
    sys.exit(1)


class KeyPool:
    """线程安全的 key 轮询池，支持临时禁用耗尽的 key"""

    def __init__(self, keys: list[str]):
        self._keys = keys
        self._cycle = itertools.cycle(range(len(keys)))
        self._cooldown: dict[int, float] = {}  # idx → 解禁时间
        self._lock = asyncio.Lock()
        self._req_count = 0

    @property
    def total(self) -> int:
        return len(self._keys)

    async def next(self) -> tuple[int, str]:
        """获取下一个可用 key，返回 (idx, key)"""
        async with self._lock:
            self._req_count += 1
            now = time.time()
            # 最多尝试一圈
            for _ in range(len(self._keys)):
                idx = next(self._cycle)
                cooldown_until = self._cooldown.get(idx, 0)
                if now >= cooldown_until:
                    return idx, self._keys[idx]
            # 全部冷却中，返回等待时间最短的
            earliest_idx = min(self._cooldown, key=self._cooldown.get)
            return earliest_idx, self._keys[earliest_idx]

    async def mark_exhausted(self, idx: int, cooldown_s: int = 60):
        """标记 key 额度耗尽，冷却 N 秒"""
        async with self._lock:
            self._cooldown[idx] = time.time() + cooldown_s
            active = sum(1 for i in range(len(self._keys))
                         if time.time() >= self._cooldown.get(i, 0))
            print(f"  ⚠️  key-{idx + 1} 额度耗尽，冷却 {cooldown_s}s（剩余 {active}/{len(self._keys)} 可用）")

    @property
    def req_count(self) -> int:
        return self._req_count


def create_app(pool: KeyPool, upstream: str, proxy_token: str, openai_compat: bool = False) -> web.Application:
    """创建 aiohttp 应用。openai_compat=True 时自动将 Anthropic Messages API 转为 OpenAI Chat Completions。"""
    import json as _json

    # 共享 session，避免每次请求创建/销毁连接导致 "closing transport" 错误
    _session: ClientSession | None = None

    async def _get_session() -> ClientSession:
        nonlocal _session
        if _session is None or _session.closed:
            timeout = ClientTimeout(total=300, sock_read=180)
            _session = ClientSession(timeout=timeout)
        return _session

    async def _cleanup(app):
        if _session and not _session.closed:
            await _session.close()

    def _filter_headers(headers) -> dict:
        """过滤掉 hop-by-hop headers"""
        skip = {"transfer-encoding", "connection", "keep-alive", "content-encoding", "content-length"}
        return {k: v for k, v in headers.items() if k.lower() not in skip}

    async def _do_request(session, method, url, headers, data):
        """发送请求，遇到连接错误自动重试一次"""
        try:
            return await session.request(method=method, url=url, headers=headers, data=data)
        except (ConnectionError, OSError) as e:
            if "closing transport" in str(e).lower() or "connection" in str(e).lower():
                # 连接被上游关闭，重试一次
                return await session.request(method=method, url=url, headers=headers, data=data)
            raise

    def _anthropic_to_openai(body: bytes) -> tuple[bytes, bool]:
        """将 Anthropic Messages API 请求体转为 OpenAI Chat Completions 格式。返回 (新body, 是否stream)"""
        try:
            req = _json.loads(body)
        except Exception:
            return body, False
        messages = req.get("messages", [])
        # Anthropic system 字段 → OpenAI system message
        oai_messages = []
        if req.get("system"):
            sys_content = req["system"]
            if isinstance(sys_content, list):
                sys_content = "\n".join(b.get("text", "") for b in sys_content if b.get("type") == "text")
            oai_messages.append({"role": "system", "content": sys_content})

        for m in messages:
            role = m.get("role", "user")
            content = m.get("content", "")

            if isinstance(content, list):
                # 检查是否包含 tool_use / tool_result 块
                text_parts = []
                tool_calls = []
                tool_results = []
                for block in content:
                    btype = block.get("type", "")
                    if btype == "text":
                        text_parts.append(block.get("text", ""))
                    elif btype == "tool_use":
                        tool_calls.append({
                            "id": block.get("id", ""),
                            "type": "function",
                            "function": {
                                "name": block.get("name", ""),
                                "arguments": _json.dumps(block.get("input", {})),
                            }
                        })
                    elif btype == "tool_result":
                        tr_content = block.get("content", "")
                        if isinstance(tr_content, list):
                            tr_content = "\n".join(b.get("text", "") for b in tr_content if b.get("type") == "text")
                        tool_results.append({
                            "tool_call_id": block.get("tool_use_id", ""),
                            "content": str(tr_content),
                        })

                if role == "assistant" and tool_calls:
                    msg = {"role": "assistant", "content": "\n".join(text_parts) if text_parts else None, "tool_calls": tool_calls}
                    oai_messages.append(msg)
                elif role == "user" and tool_results:
                    # Anthropic: user message with tool_result blocks → OpenAI: separate tool messages
                    if text_parts:
                        oai_messages.append({"role": "user", "content": "\n".join(text_parts)})
                    for tr in tool_results:
                        oai_messages.append({"role": "tool", "tool_call_id": tr["tool_call_id"], "content": tr["content"]})
                else:
                    oai_messages.append({"role": role, "content": "\n".join(text_parts) if text_parts else str(content)})
            else:
                oai_messages.append({"role": role, "content": content})

        is_stream = req.get("stream", False)
        oai_req = {
            "model": req.get("model", ""),
            "messages": oai_messages,
            "max_tokens": req.get("max_tokens", 4096),
            "stream": is_stream,
        }
        if req.get("temperature") is not None:
            oai_req["temperature"] = req["temperature"]
        if req.get("top_p") is not None:
            oai_req["top_p"] = req["top_p"]

        # Anthropic tools → OpenAI tools (function calling)
        if req.get("tools"):
            oai_tools = []
            for t in req["tools"]:
                oai_tools.append({
                    "type": "function",
                    "function": {
                        "name": t.get("name", ""),
                        "description": t.get("description", ""),
                        "parameters": t.get("input_schema", {}),
                    }
                })
            oai_req["tools"] = oai_tools
            # tool_choice 映射
            tc = req.get("tool_choice")
            if tc:
                if isinstance(tc, dict):
                    tc_type = tc.get("type", "")
                    if tc_type == "auto":
                        oai_req["tool_choice"] = "auto"
                    elif tc_type == "any":
                        oai_req["tool_choice"] = "required"
                    elif tc_type == "tool":
                        oai_req["tool_choice"] = {"type": "function", "function": {"name": tc.get("name", "")}}
                elif isinstance(tc, str):
                    oai_req["tool_choice"] = tc

        return _json.dumps(oai_req).encode(), is_stream

    def _openai_to_anthropic_resp(body: bytes) -> bytes:
        """将 OpenAI Chat Completions 非流式响应转为 Anthropic Messages API 格式"""
        try:
            resp = _json.loads(body)
        except Exception:
            return body
        choices = resp.get("choices", [])
        msg = choices[0].get("message", {}) if choices else {}
        content_blocks = []
        # 文本内容
        if msg.get("content"):
            content_blocks.append({"type": "text", "text": msg["content"]})
        # 工具调用 → Anthropic tool_use blocks
        for tc in msg.get("tool_calls", []):
            fn = tc.get("function", {})
            try:
                args = _json.loads(fn.get("arguments", "{}"))
            except Exception:
                args = {"raw": fn.get("arguments", "")}
            content_blocks.append({
                "type": "tool_use",
                "id": tc.get("id", ""),
                "name": fn.get("name", ""),
                "input": args,
            })
        if not content_blocks:
            content_blocks.append({"type": "text", "text": ""})

        # stop_reason 映射
        finish = choices[0].get("finish_reason", "stop") if choices else "stop"
        stop_reason = "tool_use" if finish == "tool_calls" else "end_turn"

        usage = resp.get("usage", {})
        anthropic_resp = {
            "id": resp.get("id", "msg_proxy"),
            "type": "message",
            "role": "assistant",
            "content": content_blocks,
            "model": resp.get("model", ""),
            "stop_reason": stop_reason,
            "usage": {
                "input_tokens": usage.get("prompt_tokens", 0),
                "output_tokens": usage.get("completion_tokens", 0),
            },
        }
        return _json.dumps(anthropic_resp).encode()

    def _openai_sse_to_anthropic_sse(chunk_data: str) -> str | None:
        """将单条 OpenAI SSE data 转为 Anthropic SSE 事件序列"""
        if chunk_data.strip() == "[DONE]":
            return "event: message_stop\ndata: {}\n\n"
        try:
            c = _json.loads(chunk_data)
        except Exception:
            return None
        delta = c.get("choices", [{}])[0].get("delta", {})
        events = []
        # 文本 delta
        text = delta.get("content", "")
        if text:
            evt = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}
            events.append(f"event: content_block_delta\ndata: {_json.dumps(evt)}\n\n")
        # 工具调用 delta
        for tc in delta.get("tool_calls", []):
            fn = tc.get("function", {})
            idx = tc.get("index", 0)
            if fn.get("name"):
                # 工具调用开始
                evt_start = {"type": "content_block_start", "index": idx + 1, "content_block": {
                    "type": "tool_use", "id": tc.get("id", f"call_{idx}"), "name": fn["name"], "input": {}}}
                events.append(f"event: content_block_start\ndata: {_json.dumps(evt_start)}\n\n")
            if fn.get("arguments", ""):
                evt_delta = {"type": "content_block_delta", "index": idx + 1, "delta": {
                    "type": "input_json_delta", "partial_json": fn["arguments"]}}
                events.append(f"event: content_block_delta\ndata: {_json.dumps(evt_delta)}\n\n")
        return "".join(events) if events else None

    async def _stream_openai_to_anthropic(request, upstream_resp, status):
        """将 OpenAI SSE 流式响应转为 Anthropic SSE 格式"""
        resp = web.StreamResponse(status=status)
        resp.content_type = "text/event-stream"
        await resp.prepare(request)
        # Anthropic 流式响应开头事件
        await resp.write(b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_proxy","type":"message","role":"assistant","content":[],"model":"proxy","stop_reason":null,"usage":{"input_tokens":0,"output_tokens":0}}}\n\n')
        await resp.write(b'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n')
        buf = ""
        async for raw_chunk in upstream_resp.content.iter_any():
            buf += raw_chunk.decode("utf-8", errors="replace")
            while "\n" in buf:
                line, buf = buf.split("\n", 1)
                line = line.strip()
                if line.startswith("data: "):
                    data_str = line[6:]
                    converted = _openai_sse_to_anthropic_sse(data_str)
                    if converted:
                        await resp.write(converted.encode())
        await resp.write_eof()
        return resp

    async def proxy_handler(request: web.Request) -> web.StreamResponse:
        auth = request.headers.get("x-api-key", "") or request.headers.get("Authorization", "").removeprefix("Bearer ")
        if proxy_token and auth != proxy_token:
            return web.json_response({"error": "unauthorized"}, status=401)

        idx, key = await pool.next()
        path = request.path_qs
        body = await request.read()

        # OpenAI 兼容模式: /v1/messages → 转换为 /v1/chat/completions
        need_translate = openai_compat and "/v1/messages" in path
        is_stream = False
        if need_translate:
            body, is_stream = _anthropic_to_openai(body)
            path = path.replace("/v1/messages", "/v1/chat/completions")

        target_url = f"{upstream.rstrip('/')}{path}"

        fwd_headers = {}
        for h in ("content-type", "accept", "anthropic-version", "anthropic-beta",
                   "x-api-key", "authorization"):
            if h in request.headers:
                fwd_headers[h] = request.headers[h]
        fwd_headers["x-api-key"] = key
        fwd_headers["authorization"] = f"Bearer {key}"

        try:
            session = await _get_session()
            upstream_resp = await _do_request(session, request.method, target_url, fwd_headers, body)
            try:
                    status = upstream_resp.status

                    if status in (402, 429):
                        await pool.mark_exhausted(idx)
                        idx2, key2 = await pool.next()
                        if idx2 != idx:
                            fwd_headers["x-api-key"] = key2
                            fwd_headers["authorization"] = f"Bearer {key2}"
                            retry_resp = await _do_request(session, request.method, target_url, fwd_headers, body)
                            try:
                                status = retry_resp.status
                                print(f"  [{pool.req_count:>4}] {request.method} {path[:50]} → key-{idx2 + 1} (retry, {status})")
                                if retry_resp.headers.get("content-type", "").startswith("text/event-stream"):
                                    if need_translate:
                                        return await _stream_openai_to_anthropic(request, retry_resp, status)
                                    resp = web.StreamResponse(status=status, headers=_filter_headers(retry_resp.headers))
                                    resp.content_type = "text/event-stream"
                                    await resp.prepare(request)
                                    async for chunk in retry_resp.content.iter_any():
                                        await resp.write(chunk)
                                    await resp.write_eof()
                                    return resp
                                resp_body = await retry_resp.read()
                                if need_translate and status == 200:
                                    resp_body = _openai_to_anthropic_resp(resp_body)
                                return web.Response(status=status, headers=_filter_headers(retry_resp.headers), body=resp_body)
                            finally:
                                retry_resp.release()

                    print(f"  [{pool.req_count:>4}] {request.method} {path[:50]} → key-{idx + 1} ({status})")

                    if upstream_resp.headers.get("content-type", "").startswith("text/event-stream"):
                        if need_translate:
                            return await _stream_openai_to_anthropic(request, upstream_resp, status)
                        resp = web.StreamResponse(status=status, headers=_filter_headers(upstream_resp.headers))
                        resp.content_type = "text/event-stream"
                        await resp.prepare(request)
                        async for chunk in upstream_resp.content.iter_any():
                            await resp.write(chunk)
                        await resp.write_eof()
                        return resp

                    resp_body = await upstream_resp.read()
                    if need_translate and status == 200:
                        resp_body = _openai_to_anthropic_resp(resp_body)
                    return web.Response(status=status, headers=_filter_headers(upstream_resp.headers), body=resp_body)
            finally:
                upstream_resp.release()
        except Exception as e:
            print(f"  [{pool.req_count:>4}] ❌ key-{idx + 1} 连接失败: {e}")
            return web.json_response({"error": str(e)}, status=502)

    app = web.Application()
    app.on_shutdown.append(_cleanup)
    app.router.add_route("*", "/{path:.*}", proxy_handler)
    return app


def main():
    parser = argparse.ArgumentParser(
        description="API Key Round-Robin Proxy for tchkiller",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python3 key_proxy.py --keys "sk-1,sk-2,sk-3" --upstream https://open.bigmodel.cn/api/anthropic
  python3 key_proxy.py --keys-file keys.txt --upstream https://api.minimaxi.com/anthropic --port 8090
  python3 key_proxy.py --keys-file keys.txt --upstream https://open.bigmodel.cn/api/anthropic --token mytoken
        """,
    )
    parser.add_argument("--keys", help="逗号分隔的 API keys")
    parser.add_argument("--keys-file", help="API keys 文件（每行一个 key）")
    parser.add_argument("--upstream", required=True, help="上游 API base URL")
    parser.add_argument("--port", type=int, default=8090, help="监听端口 (默认 8090)")
    parser.add_argument("--host", default="127.0.0.1", help="监听地址 (默认 127.0.0.1)")
    parser.add_argument("--token", default="proxy-local-token", help="Proxy 鉴权 token (默认 proxy-local-token)")
    parser.add_argument("--openai", action="store_true", help="OpenAI 兼容模式：自动将 Anthropic Messages API 转为 OpenAI Chat Completions")
    parser.add_argument("--cooldown", type=int, default=60, help="key 额度耗尽冷却秒数 (默认 60)")
    args = parser.parse_args()

    # 加载 keys
    keys = []
    if args.keys:
        keys = [k.strip() for k in args.keys.split(",") if k.strip()]
    elif args.keys_file:
        p = Path(args.keys_file)
        if not p.exists():
            print(f"❌ keys 文件不存在: {p}")
            sys.exit(1)
        keys = [line.strip() for line in p.read_text().splitlines() if line.strip() and not line.startswith("#")]
    else:
        print("❌ 必须指定 --keys 或 --keys-file")
        sys.exit(1)

    if not keys:
        print("❌ 没有可用的 API key")
        sys.exit(1)

    pool = KeyPool(keys)

    print(f"""
{'═' * 55}
  🔄 tchkiller API Key Round-Robin Proxy
{'═' * 55}
  上游: {args.upstream}
  Keys: {pool.total} 个
  监听: http://{args.host}:{args.port}
  鉴权: {args.token[:12]}...
  冷却: {args.cooldown}s (额度耗尽后)
  模式: {'OpenAI→Anthropic 协议转换' if args.openai else '透明转发'}
{'═' * 55}
""")

    app = create_app(pool, args.upstream, args.token, openai_compat=args.openai)
    web.run_app(app, host=args.host, port=args.port, print=lambda _: None)


if __name__ == "__main__":
    main()