Sending requests at scale

pip install requests

import time
import requests

BASE_URL = "https://api.sailresearch.com/v1"
API_KEY = "YOUR_KEY_HERE"
HEADERS = {"Authorization": f"Bearer {API_KEY}"}

# 1. Submit a batch
batch_requests = [
    {
        "custom_id": f"request-{i}",
        "params": {
            "model": "zai-org/GLM-5.2-FP8",
            "max_output_tokens": 2048,
            "input": [{"role": "user", "content": f"What is a fun fact about the number {i}?"}],
            "metadata": {"completion_window": "balanced"},
        },
    }
    for i in range(100)
]

resp = requests.post(
    f"{BASE_URL}/batches",
    headers=HEADERS,
    json={
        "endpoint": "/v1/responses",
        "label": "my-batch-job",
        "requests": batch_requests,
    },
)
batch = resp.json()
batch_id = batch["id"]
print(f"Created batch {batch_id} with {batch['request_counts']} requests")

# 2. Poll until all requests are completed
while True:
    status_resp = requests.get(f"{BASE_URL}/batches/{batch_id}", headers=HEADERS)
    status = status_resp.json()
    request_status = status["request_status"]

    done = sum(1 for r in request_status.values() if r["status"] in ("COMPLETED", "FAILED", "CANCELLED"))
    print(f"Progress: {done}/{len(request_status)}")

    if done == len(request_status):
        break
    time.sleep(5)

# 3. Retrieve results
for custom_id, info in request_status.items():
    if info["status"] == "COMPLETED":
        result = requests.get(
            f"{BASE_URL}/batches/{batch_id}/{custom_id}", headers=HEADERS
        ).json()
        output_text = "".join(
            c["text"]
            for item in result.get("output", [])
            for c in item.get("content", [])
            if c.get("type") == "output_text"
        )
        print(f"{custom_id}: {output_text[:100]}...")

pip install tqdm
pip install 'openai[aiohttp]'

import argparse
import asyncio

from tqdm import tqdm
from openai import DefaultAioHttpClient
from openai import AsyncOpenAI

async def main(
    num_requests: int,
    input_tokens: int,
    max_output_tokens: int,
    model: str,
):
    base_url = "https://api.sailresearch.com/v1"
    api_key = "YOUR_KEY_HERE"

    async with AsyncOpenAI(
        base_url=base_url,
        api_key=api_key,
        http_client=DefaultAioHttpClient(),
    ) as client:
        # List supported models
        models = await client.models.list()
        supported_models = [m.id for m in models.data]
        print("Supported Models:")
        print(supported_models)

        # Submit requests concurrently
        sem = asyncio.Semaphore(200)
        pbar_submit = tqdm(total=num_requests, desc="Submitting requests")

        async def submit(i):
            filler = ""
            if input_tokens > 0:
                filler = " " + ("word " * input_tokens)
            content = f"TASK {i}: What is a fun fact about the number {i}? Then, find the word at index {i} in the following sequence of words: {filler}"

            async with sem:
                response = await client.responses.create(
                    model=model,
                    input=[{"role": "user", "content": content}],
                    max_output_tokens=max_output_tokens,
                    background=True,
                )
            pbar_submit.update(1)
            return response.id

        response_ids = await asyncio.gather(*[submit(i) for i in range(num_requests)])
        pbar_submit.close()
        response_ids = list(response_ids)
        print(f"Created {len(response_ids)} response IDs")

        # Poll for completions
        usable = {}
        poll_sem = asyncio.Semaphore(200)
        pbar = tqdm(total=len(response_ids), desc="Polling responses")

        async def poll(response_id):
            for _ in range(3600):  # 1 hour timeout
                async with poll_sem:
                    response = await client.responses.retrieve(response_id)
                if response.status == "completed":
                    usable[response_id] = response
                    pbar.update(1)
                    return
                if response.status == "incomplete":
                    reason = getattr(response.incomplete_details, "reason", None)
                    if reason == "max_output_tokens":
                        usable[response_id] = response
                    else:
                        print(f"{response_id} is incomplete because of {reason or 'an unknown reason'}")
                    pbar.update(1)
                    return
                if response.status in {"failed", "cancelled"}:
                    print(f"{response_id} ended with status {response.status}")
                    pbar.update(1)
                    return
                await asyncio.sleep(1)
            print(f"Timed out waiting for {response_id}")

        await asyncio.gather(*[poll(rid) for rid in response_ids])
        pbar.close()

        for response_id in response_ids:
            resp = usable.get(response_id)
            if resp is None:
                continue
            output_text = "".join(
                c.text
                for item in resp.output
                for c in getattr(item, "content", []) or []
                if getattr(c, "type", None) == "output_text"
            )
            print(f"\n\n{response_id} {resp.status} with output:\n{output_text}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num-requests", type=int, default=20000)
    parser.add_argument("--input-tokens", type=int, default=50)
    parser.add_argument("--max-output-tokens", type=int, default=4000)
    parser.add_argument("--model", type=str, default="zai-org/GLM-5.2-FP8")
    args = parser.parse_args()

    asyncio.run(
        main(
            num_requests=args.num_requests,
            input_tokens=args.input_tokens,
            max_output_tokens=args.max_output_tokens,
            model=args.model,
        )
    )