scripts/call_v1_stream.py

#!/usr/bin/env python3
from __future__ import annotations

import os
import sys

import requests

base_url = os.environ.get("OPENAI_BASE_URL") or os.environ.get("SPACE_URL")
if not base_url:
    raise SystemExit("Set OPENAI_BASE_URL or SPACE_URL")
base_url = base_url.rstrip("/")
if not base_url.endswith("/v1"):
    base_url += "/v1"

hf_token = os.environ.get("HF_TOKEN", "")
model = os.environ.get("MODEL", "unsloth/diffusiongemma-26B-A4B-it-GGUF:Q4_K_M")
prompt = " ".join(sys.argv[1:]) or "Write a short haiku about GPUs."

headers = {"Content-Type": "application/json"}
if hf_token:
    headers["Authorization"] = f"Bearer {hf_token}"

with requests.post(
    f"{base_url}/chat/completions",
    headers=headers,
    json={
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": int(os.environ.get("MAX_TOKENS", "256")),
        "stream": True,
    },
    stream=True,
    timeout=int(os.environ.get("REQUEST_TIMEOUT_SECONDS", "1200")),
) as resp:
    print(resp.status_code, resp.reason, file=sys.stderr)
    for line in resp.iter_lines(decode_unicode=True):
        if line:
            print(line)