# /// script
# requires-python = ">=3.10"
# dependencies = ["openai>=1", "httpx>=0.28"]
# ///
"""Stream an ACS completion while allowing a long cold start.

ACS_API_KEY=... uv run examples/python_long_requests.py --model llama-8b
Optional ACS_API_BASE must include /v1. No automatic retries are performed.
"""

import argparse
import os
import sys

import httpx
from openai import APIError, OpenAI


def make_client(*, api_key: str, base_url: str) -> OpenAI:
    return OpenAI(
        api_key=api_key,
        base_url=base_url,
        # Read timeout is the maximum gap between received bytes, not total
        # request duration. SSE heartbeats keep a waiting request active.
        timeout=httpx.Timeout(connect=20, read=900, write=60, pool=20),
        max_retries=0,
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="llama-8b")
    parser.add_argument("--prompt", default="Once upon a time")
    parser.add_argument("--max-tokens", type=int, default=64)
    args = parser.parse_args()
    key = os.environ.get("ACS_API_KEY")
    if not key:
        parser.error("set ACS_API_KEY in your environment")
    with make_client(
        api_key=key,
        base_url=os.environ.get("ACS_API_BASE", "https://infra.acsresearch.org/v1"),
    ) as client:
        try:
            with client.completions.create(
                model=args.model,
                prompt=args.prompt,
                max_tokens=args.max_tokens,
                stream=True,
            ) as stream:
                for chunk in stream:
                    for choice in chunk.choices:
                        print(choice.text, end="", flush=True)
            print()
        except APIError as exc:
            # The SDK raises for a structured SSE error even after heartbeats
            # have committed HTTP 200. Previously printed text is incomplete.
            print(f"\nRequest failed; output above is incomplete: {exc}", file=sys.stderr)
            return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())