Function poll_batch_until_complete

scrapegraphai/utils/batch_api.py:190–247

Poll a batch job until it completes, fails, or times out. Args: client: An initialized OpenAI client. batch_id: The batch ID to poll. poll_interval: Seconds between status checks. max_wait_time: Maximum seconds to wait before giving up. Returns: Fina

(
    client: OpenAI,
    batch_id: str,
    poll_interval: int = DEFAULT_POLL_INTERVAL,
    max_wait_time: int = DEFAULT_MAX_WAIT_TIME,
)

def poll_batch_until_complete(
    client: OpenAI,
    batch_id: str,
    poll_interval: int = DEFAULT_POLL_INTERVAL,
    max_wait_time: int = DEFAULT_MAX_WAIT_TIME,
) -> BatchJobInfo:
    """Poll a batch job until it completes, fails, or times out.

    Args:
        client: An initialized OpenAI client.
        batch_id: The batch ID to poll.
        poll_interval: Seconds between status checks.
        max_wait_time: Maximum seconds to wait before giving up.

    Returns:
        Final BatchJobInfo when the batch reaches a terminal state.

    Raises:
        TimeoutError: If max_wait_time is exceeded.
        RuntimeError: If the batch fails or is cancelled.
    """
    terminal_states = {"completed", "failed", "expired", "cancelled"}
    start_time = time.time()

    logger.info(
        f"Polling batch {batch_id} every {poll_interval}s "
        f"(max wait: {max_wait_time}s)..."
    )

    while True:
        elapsed = time.time() - start_time
        if elapsed > max_wait_time:
            raise TimeoutError(
                f"Batch {batch_id} did not complete within "
                f"{max_wait_time}s (last status check at {elapsed:.0f}s)"
            )

        info = get_batch_status(client, batch_id)

        logger.info(
            f"Batch {batch_id}: {info.status} "
            f"({info.completed_requests}/{info.total_requests} done, "
            f"{info.failed_requests} failed)"
        )

        if info.status in terminal_states:
            if info.status == "failed":
                raise RuntimeError(
                    f"Batch {batch_id} failed. "
                    f"Error file: {info.error_file_id}"
                )
            if info.status in {"expired", "cancelled"}:
                raise RuntimeError(
                    f"Batch {batch_id} was {info.status}."
                )
            return info

        time.sleep(poll_interval)