OpenAI Batch API Enqueued Token Limits, Rate Limits & Cost Optimization

import json
from openai import OpenAI

client = OpenAI()


ACTIVE_STATUSES = {"validating", "in_progress", "finalizing"}


def estimate_enqueued_tokens_for_batch(batch) -> int:
    """Estimate the token budget consumption of an active batch."""
    if batch.status not in ACTIVE_STATUSES:
        return 0

    total = batch.request_counts.total if batch.request_counts else 0
    done = (batch.request_counts.completed + batch.request_counts.failed) if batch.request_counts else 0
    remaining = max(total - done, 0)

    # This is an estimate — you know your workload better than we can
    # Typical assumption: 500 input + 500 output tokens per request average
    return remaining * 1000


def current_enqueued_load() -> dict:
    """Sum enqueued tokens across all your active batches."""
    load_by_model = {}
    cursor = None

    while True:
        page = client.batches.list(limit=100, after=cursor)
        for batch in page.data:
            if batch.status not in ACTIVE_STATUSES:
                continue
            # We don't know the model without inspecting the input file — track
            # via metadata tag at submission time
            model = batch.metadata.get("model", "unknown") if batch.metadata else "unknown"
            load_by_model[model] = load_by_model.get(model, 0) + estimate_enqueued_tokens_for_batch(batch)

        if not page.has_more:
            break
        cursor = page.data[-1].id

    return load_by_model


# ✅ Tag every submission with model, enabling load tracking
def submit_tracked(prompts, model: str, avg_input_tokens: int, avg_output_tokens: int) -> str:
    """Submit with metadata that lets us estimate load later."""
    import io
    lines = []
    for p in prompts:
        lines.append(json.dumps({
            "custom_id": p["id"],
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {"model": model, "messages": p["messages"],
                     "max_completion_tokens": avg_output_tokens},
        }))

    upload = client.files.create(
        file=("input.jsonl", io.BytesIO("\n".join(lines).encode())),
        purpose="batch",
    )

    batch = client.batches.create(
        input_file_id=upload.id,
        endpoint="/v1/chat/completions",
        completion_window="24h",
        metadata={
            "model": model,
            "expected_tokens_per_req": str(avg_input_tokens + avg_output_tokens),
        },
    )
    return batch.id


# ✅ Rate-limit-aware submission queue
class BatchScheduler:
    """Enforce a max-enqueued-tokens budget per model."""

    def __init__(self, per_model_budget: dict[str, int]):
        self.budgets = per_model_budget    # e.g. {"gpt-5.4": 100_000_000, "o3": 20_000_000}

    def can_submit(self, model: str, expected_tokens: int) -> bool:
        current = current_enqueued_load().get(model, 0)
        remaining = self.budgets[model] - current
        return expected_tokens <= remaining

    def submit_if_capacity(self, prompts, model, avg_input, avg_output):
        expected = len(prompts) * (avg_input + avg_output)
        if not self.can_submit(model, expected):
            raise RuntimeError(
                f"insufficient enqueued token budget for {model}: "
                f"need {expected:,}, budget remaining ~{self.budgets[model] - current_enqueued_load().get(model, 0):,}"
            )
        return submit_tracked(prompts, model, avg_input, avg_output)


# ✅ Split large workloads across models to distribute load
def multi_model_submit(prompts, model_ratios: dict[str, float]):
    """Split prompts across multiple compatible models.
    model_ratios: e.g. {"gpt-5.4": 0.7, "gpt-5.4-nano": 0.3}"""
    batches = {}
    idx = 0
    for model, ratio in model_ratios.items():
        count = int(len(prompts) * ratio)
        chunk = prompts[idx:idx + count]
        idx += count
        batches[model] = submit_tracked(chunk, model, avg_input_tokens=500, avg_output_tokens=500)
    return batches

import json
import hashlib


# Prompt Caching requires a stable prefix ≥ 1024 tokens.
# Structure your prompts so the FIRST 1024+ tokens are identical across items.


# ✅ GOOD — long stable system prompt (cache-eligible)
SYSTEM_PROMPT = """You are an expert classifier for legal documents.

You will receive a document and must classify it into exactly one of these categories:
- Contract: Legal agreements between parties
- Court Filing: Documents filed with a court
- Correspondence: Letters, emails, official notifications
- Regulation: Government rules, statutes, administrative rules
- Opinion: Judicial opinions, expert opinions, legal analysis
- Other: Anything that doesn't fit above

For each classification, respond with a JSON object:
{
  "category": "one of the above exactly",
  "confidence": 0.0-1.0,
  "reasoning": "brief 1-2 sentence explanation"
}

Consider these features:
- Language formality
- Presence of party names, dates, signatures
- Reference to case numbers, citations, statutes
- Structure and formatting

Additional guidance:
- When multiple categories apply, pick the primary function
- If ambiguous, use confidence < 0.7 and explain in reasoning
- Never invent categories not in the list

Examples:

Input: "This Agreement ('Agreement') is entered into on..."
Output: {"category": "Contract", "confidence": 0.98, "reasoning": "Standard contract opening language with defined parties."}

Input: "IN THE UNITED STATES DISTRICT COURT..."
Output: {"category": "Court Filing", "confidence": 0.99, "reasoning": "Court header identifies this as filed litigation."}

Input: "Dear Ms. Chen, Following our discussion on..."
Output: {"category": "Correspondence", "confidence": 0.9, "reasoning": "Business letter format with salutation."}

(... more examples ...)

Now classify the following document:
"""


def build_batch_with_caching(documents: list[dict]) -> bytes:
    """All items share the same SYSTEM_PROMPT prefix — maximizes cache hits."""
    lines = []
    for doc in documents:
        req = {
            "custom_id": doc["id"],
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {
                "model": "gpt-5.4",
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},   # SHARED — cached after first use
                    {"role": "user",   "content": doc["text"]},     # varies per item
                ],
                "max_completion_tokens": 300,
            },
        }
        lines.append(json.dumps(req))
    return "\n".join(lines).encode()


# ❌ ANTI-PATTERN — unique system prompt per item breaks caching
def bad_no_caching(documents):
    lines = []
    for doc in documents:
        # System prompt includes item-specific info → cache miss every time
        system = f"You are a classifier. Working on document {doc['id']}. Categories: ..."
        req = {
            "custom_id": doc["id"], "method": "POST", "url": "/v1/chat/completions",
            "body": {"model": "gpt-5.4",
                     "messages": [{"role": "system", "content": system},
                                  {"role": "user", "content": doc["text"]}]},
        }
        lines.append(json.dumps(req))
    return "\n".join(lines).encode()


# ✅ Even better — sort items by similarity so cache stays warm
# Not strictly required for Batch (cache works within 24h regardless), but helps
# ensure adjacent items in the batch benefit from prior cache warming.
def sort_for_cache_locality(documents: list[dict], key_fn=None) -> list[dict]:
    key_fn = key_fn or (lambda d: d.get("source", ""))
    return sorted(documents, key=key_fn)


# ✅ Cost estimator — factor cache hits into your budget projection
def estimate_batch_cost(
    n_requests: int,
    avg_input_tokens: int,
    avg_output_tokens: int,
    cache_hit_ratio: float,               # fraction of input tokens likely cached
    input_price_per_1m_sync: float,
    output_price_per_1m_sync: float,
) -> dict:
    """Estimate cost with Batch discount + Prompt Cache stacking."""

    BATCH_DISCOUNT = 0.5                  # 50% off
    CACHE_DISCOUNT = 0.5                  # cached input another 50% off (of the batched rate)

    input_tokens = n_requests * avg_input_tokens
    output_tokens = n_requests * avg_output_tokens

    cached_input = int(input_tokens * cache_hit_ratio)
    uncached_input = input_tokens - cached_input

    # Batch input pricing
    input_batch_rate = input_price_per_1m_sync * BATCH_DISCOUNT
    output_batch_rate = output_price_per_1m_sync * BATCH_DISCOUNT

    # Cached input: batch rate × cache discount
    cached_cost = (cached_input / 1_000_000) * input_batch_rate * CACHE_DISCOUNT
    uncached_cost = (uncached_input / 1_000_000) * input_batch_rate
    output_cost = (output_tokens / 1_000_000) * output_batch_rate

    total = cached_cost + uncached_cost + output_cost
    sync_baseline = (
        (input_tokens / 1_000_000) * input_price_per_1m_sync +
        (output_tokens / 1_000_000) * output_price_per_1m_sync
    )

    return {
        "batch_cached_cost_usd": cached_cost,
        "batch_uncached_cost_usd": uncached_cost,
        "batch_output_cost_usd": output_cost,
        "batch_total_usd": total,
        "sync_baseline_usd": sync_baseline,
        "savings_vs_sync_pct": (1 - total / sync_baseline) * 100 if sync_baseline else 0,
    }


# Example
result = estimate_batch_cost(
    n_requests=100_000,
    avg_input_tokens=1500,
    avg_output_tokens=200,
    cache_hit_ratio=0.85,                 # 85% of input is the shared system prompt
    input_price_per_1m_sync=2.50,
    output_price_per_1m_sync=10.00,
)
# {'batch_cached_cost_usd': ~80, 'batch_uncached_cost_usd': ~28, 'batch_output_cost_usd': ~100,
#  'batch_total_usd': ~208, 'sync_baseline_usd': ~575, 'savings_vs_sync_pct': ~64%}

# Rule of thumb for reasoning tokens per request:
#   low effort:      500-2,000 tokens
#   medium effort:   3,000-15,000 tokens
#   high effort:     15,000-100,000 tokens

REASONING_TOKENS_ESTIMATE = {
    "low":      2_000,
    "medium":   8_000,
    "high":    30_000,
    "minimal":    500,     # GPT-5 only
    "none":         0,     # GPT-5 only
}


def estimate_reasoning_batch_cost(
    n_requests: int,
    effort: str,
    avg_input_tokens: int,
    avg_visible_output: int,
    model_output_rate_batch: float,       # $ per 1M output tokens at batch rate
    model_input_rate_batch: float,
) -> dict:
    reasoning_per_req = REASONING_TOKENS_ESTIMATE[effort]

    total_input = n_requests * avg_input_tokens
    total_reasoning = n_requests * reasoning_per_req
    total_visible = n_requests * avg_visible_output
    total_output = total_reasoning + total_visible

    input_cost = (total_input / 1_000_000) * model_input_rate_batch
    output_cost = (total_output / 1_000_000) * model_output_rate_batch

    return {
        "n_requests": n_requests,
        "effort": effort,
        "input_tokens": total_input,
        "reasoning_tokens": total_reasoning,
        "visible_output_tokens": total_visible,
        "input_cost_usd": input_cost,
        "output_cost_usd": output_cost,
        "total_cost_usd": input_cost + output_cost,
        "reasoning_share_of_cost_pct": (total_reasoning / total_output * 100) if total_output else 0,
    }


# Example — o3 batch, 10K requests, medium effort
# Assume o3 batch rates: $1 input, $4 output per 1M tokens (hypothetical)
result = estimate_reasoning_batch_cost(
    n_requests=10_000,
    effort="medium",
    avg_input_tokens=1_000,
    avg_visible_output=500,
    model_output_rate_batch=4.00,
    model_input_rate_batch=1.00,
)
# reasoning_tokens: 80,000,000; output_cost: $340
# vs. if you naively budgeted only visible output:
#   just_visible = 10_000 * 500 = 5M tokens = $20
#   you'd be off by 17×


# ✅ Cap reasoning per request via max_completion_tokens
def batch_body_with_reasoning_cap(prompt: str, effort: str = "medium") -> dict:
    max_reasoning_expected = REASONING_TOKENS_ESTIMATE[effort]
    max_visible = 500
    return {
        "model": "o3",
        "messages": [{"role": "user", "content": prompt}],
        "reasoning_effort": effort,
        # Total budget = expected reasoning + visible + safety margin
        "max_completion_tokens": max_reasoning_expected + max_visible + 2_000,
    }


# ✅ Cost tracker — after batch completes, verify actual vs estimated
def audit_batch_cost(batch_id: str, sync_input_rate: float, sync_output_rate: float):
    from openai import OpenAI
    client = OpenAI()
    batch = client.batches.retrieve(batch_id)

    # Sum actual tokens from the output file
    import json
    total_input = 0
    total_output = 0
    total_reasoning = 0

    for line in stream_lines(batch.output_file_id):
        row = json.loads(line)
        body = row.get("response", {}).get("body", {})
        usage = body.get("usage", {})
        total_input += usage.get("prompt_tokens", 0)
        total_output += usage.get("completion_tokens", 0)

        details = usage.get("completion_tokens_details", {})
        total_reasoning += details.get("reasoning_tokens", 0)

    # Actual costs at batch rate (50% discount)
    input_cost = (total_input / 1_000_000) * sync_input_rate * 0.5
    output_cost = (total_output / 1_000_000) * sync_output_rate * 0.5

    print(f"Batch {batch_id} actual usage:")
    print(f"  Input tokens:     {total_input:,}")
    print(f"  Output tokens:    {total_output:,} (of which {total_reasoning:,} reasoning)")
    print(f"  Input cost:       ${input_cost:,.2f}")
    print(f"  Output cost:      ${output_cost:,.2f}")
    print(f"  Reasoning share:  {total_reasoning / total_output * 100:.1f}%")


def stream_lines(file_id):
    from openai import OpenAI
    text = OpenAI().files.content(file_id).read().decode()
    for line in text.split("\n"):
        if line.strip():
            yield line


# ✅ Route effort per item to save on batch costs
def route_effort_by_length(prompt: str) -> str:
    """Simple heuristic — long prompts probably need more thinking."""
    if len(prompt) < 500:
        return "low"
    if len(prompt) < 3000:
        return "medium"
    return "high"