This is the decorator I've been using. It doesn't have the functionality @prrao87 mentioned in the discord to pass in the failure to the model when retrying, but I should probably implement that.

def retry(
    model_class: Type[BaseModel],
    retries: int = 3,
    delay: float = 0.5
):
    """
    Retries LMP up to 'retries' times if the output fails to validate against 'model_class'.
    Args:
        model_class: The Pydantic model class used to validate the function output.
        retries: Maximum number of retries before giving up.
        delay: Delay in seconds between retries.
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for _ in range(retries):
                try:
                    result = func(*args, **kwargs)
                    model_class.model_validate(json.loads(result)) 
                    return result
                except ValidationError as validation_error:
                    logging.warning("Model validation failed:")
                    last_exception = validation_error
                    time.sleep(delay)

            if last_exception is not None:
                raise last_exception
        return wrapper
    return decorator

Retry API? #416

Description

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions