import time

def wait_for_completion(client, task_id, max_wait=300):
    """Wait for the task to complete"""
    start_time = time.time()

    while time.time() - start_time < max_wait:
        result = client.tasks.get(task_id)

        if result.status == "completed":
            return result
        elif result.status == "failed":
            raise Exception(f"Task failed: {result.error}")

        time.sleep(2)  # wait 2 seconds before querying again

    raise Exception("Task timeout")

try:
    response = client.images.generate(...)
except EvolinkError as e:
    if e.status_code == 401:
        print("Invalid API key")
    elif e.status_code == 402:
        print("Insufficient account balance")
    else:
        print(f"Error: {e.message}")