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}")