RULE/model_integration/utils.py

import functools
import time
from typing import Any, Callable


def retry(max_attempts: int = 3, delay: float = 2.0, backoff: float = 2.0):
    """
    Retry decorator with exponential backoff for flaky API calls.
    """
    def decorator(func: Callable):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            current_delay = delay
            for attempt in range(1, max_attempts + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if attempt == max_attempts:
                        raise e
                    time.sleep(current_delay)
                    current_delay *= backoff
            return None
        return wrapper
    return decorator

def wait_for_completion(
    check_status_callable: Callable[[], Any],
    is_completed_callable: Callable[[Any], bool],
    poll_interval: int = 5,
    max_wait_seconds: int = 600
) -> Any:
    """
    Polls a status check function until a completion condition is met.
    The prompt enforces waiting for AT LEAST 60 seconds (but usually we wait until completed).
    Here, the requirement was stated as "WAIT for at least 60 seconds (blocking or polling)".
    We implement it by polling and checking the condition.
    """
    start_time = time.time()
    
    while True:
        elapsed = time.time() - start_time
        
        result = check_status_callable()
        
        if is_completed_callable(result):
            # If the API marks it as complete, return immediately.
            return result
            
        if elapsed > max_wait_seconds:
            raise TimeoutError(f"Operation timed out after {max_wait_seconds} seconds")
            
        time.sleep(poll_interval)