OpenAI Responses API 429 레이트리밋 토큰버킷으로 끝내기

import asyncio
import time
from dataclasses import dataclass

@dataclass
class TokenBucket:
    capacity: float
    refill_per_sec: float
    tokens: float = 0.0
    updated_at: float = 0.0

    def __post_init__(self):
        self.tokens = self.capacity
        self.updated_at = time.monotonic()

    def _refill(self):
        now = time.monotonic()
        elapsed = now - self.updated_at
        if elapsed > 0:
            self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_per_sec)
            self.updated_at = now

    async def acquire(self, amount: float = 1.0):
        while True:
            self._refill()
            if self.tokens >= amount:
                self.tokens -= amount
                return
            # 부족한 토큰을 채우는 데 필요한 시간만큼 sleep
            deficit = amount - self.tokens
            wait_sec = deficit / self.refill_per_sec
            await asyncio.sleep(min(wait_sec, 0.25))


class DualLimiter:
    def __init__(self, rpm: int, tpm: int, burst_seconds: float = 3.0):
        # RPM -> 초당 리필
        rps = rpm / 60.0
        tps = tpm / 60.0
        self.req_bucket = TokenBucket(capacity=rps * burst_seconds, refill_per_sec=rps)
        self.tok_bucket = TokenBucket(capacity=tps * burst_seconds, refill_per_sec=tps)

    async def acquire(self, estimated_tokens: int):
        # 두 버킷 모두 확보해야 요청 전송
        await self.req_bucket.acquire(1)
        await self.tok_bucket.acquire(estimated_tokens)

from openai import AsyncOpenAI

client = AsyncOpenAI()
limiter = DualLimiter(rpm=600, tpm=120_000, burst_seconds=5)

def estimate_tokens(prompt: str, max_output_tokens: int) -> int:
    # 실전에서는 tiktoken 등으로 입력 토큰을 계산하고,
    # 출력은 max_output_tokens 또는 p95 추정치를 사용
    return int(len(prompt) / 4) + max_output_tokens

async def call_responses(prompt: str):
    est = estimate_tokens(prompt, max_output_tokens=800)
    await limiter.acquire(est)

    return await client.responses.create(
        model="gpt-4.1-mini",
        input=prompt,
        max_output_tokens=800,
    )

-- KEYS[1] = bucket key
-- ARGV[1] = capacity
-- ARGV[2] = refill_per_ms
-- ARGV[3] = amount
-- ARGV[4] = now_ms

local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refill_per_ms = tonumber(ARGV[2])
local amount = tonumber(ARGV[3])
local now = tonumber(ARGV[4])

local data = redis.call('HMGET', key, 'tokens', 'updated_at')
local tokens = tonumber(data[1])
local updated_at = tonumber(data[2])

if tokens == nil then
  tokens = capacity
  updated_at = now
end

local elapsed = now - updated_at
if elapsed > 0 then
  tokens = math.min(capacity, tokens + elapsed * refill_per_ms)
  updated_at = now
end

if tokens >= amount then
  tokens = tokens - amount
  redis.call('HMSET', key, 'tokens', tokens, 'updated_at', updated_at)
  redis.call('PEXPIRE', key, 600000)
  return {1, 0}
else
  local deficit = amount - tokens
  local wait_ms = math.ceil(deficit / refill_per_ms)
  redis.call('HMSET', key, 'tokens', tokens, 'updated_at', updated_at)
  redis.call('PEXPIRE', key, 600000)
  return {0, wait_ms}
end

import asyncio
import time
import redis.asyncio as redis

LUA = open("bucket.lua", "r", encoding="utf-8").read()

class RedisTokenBucket:
    def __init__(self, r: redis.Redis, key: str, capacity: float, refill_per_sec: float):
        self.r = r
        self.key = key
        self.capacity = capacity
        self.refill_per_ms = refill_per_sec / 1000.0
        self.sha = None

    async def load(self):
        self.sha = await self.r.script_load(LUA)

    async def acquire(self, amount: float):
        assert self.sha is not None
        while True:
            now_ms = int(time.time() * 1000)
            ok, wait_ms = await self.r.evalsha(
                self.sha,
                1,
                self.key,
                self.capacity,
                self.refill_per_ms,
                amount,
                now_ms,
            )
            if int(ok) == 1:
                return
            await asyncio.sleep(min(float(wait_ms) / 1000.0, 0.25))


class RedisDualLimiter:
    def __init__(self, r: redis.Redis, prefix: str, rpm: int, tpm: int, burst_seconds: float = 5.0):
        rps = rpm / 60.0
        tps = tpm / 60.0
        self.req = RedisTokenBucket(r, f"{prefix}:req", capacity=rps * burst_seconds, refill_per_sec=rps)
        self.tok = RedisTokenBucket(r, f"{prefix}:tok", capacity=tps * burst_seconds, refill_per_sec=tps)

    async def load(self):
        await self.req.load()
        await self.tok.load()

    async def acquire(self, estimated_tokens: int):
        # 순서 중요: req 먼저 잡고 tok 잡는 방식은 tok에서 대기할 때 req 토큰을 점유할 수 있음
        # 실전에서는 둘 다 동시에 확보하거나, 실패 시 롤백이 가능한 설계를 고려
        await self.req.acquire(1)
        await self.tok.acquire(estimated_tokens)