Retry API?

[alex-dixon]

Users have asked whether there will be a retry decorator or api of some kind. Just opening this issue to track it.

Edit. Changed title to reflect discussion of retries in general vs a particular decorator

[prrao87]

In my case, a simple naive retry logic (max number of retries customizable) could work. As an example, I ask ell to prompt an LLM to generate a Cypher query (that can query a graph database), but in certain cases, the LLM can misinterpret the schema, and use the wrong property name in the query, which results in an empty response from the database. However, a simple retry can work and generate the right query in the second attempt, preventing the empty response. Something as simple as @retry(max_tries=3) could work in such a scenario.

[alex-dixon]

Some questions I had, some rhetorical just to get me thinking:

• What business is retry of ell’s? Isn’t this more generic functionality than ell itself? Would clear examples help more than adding a retry api in ell?
• If retry is part of ell, should ell track them? Ell can know “this was retry attempt 1” and log that along with the failure. This would require an api that lets ell know a retry is taking place.
• Are retry-related parameters part of an lmp version?
• Add a separate decorator and stack them or add a parameter to the existing one or X?
• What constitutes a failure that should be retried according to a retry policy vs a failure that should not be retried? How best to api-ify this? Anything short of an arbitrary predicate function is restrictive. Should the retry predicate be part of versioning?
• Do we serialize errors? Would logging retries be the first instance of ell recording failures as part of tracing?

[apandy02]

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 

[prrao87]

@apandy02 I think this retry logic is exactly what @MadcowD took issue with in his taking apart of Instructor's code base in his livestream at this exact time stamp.

The comment in the chat was:

If all instructor did was retry till we get the right schema that would be bad lol

From a functionality standpoint, I think I agree, it would be great if the retry logic would append the error message to the subsequent prompt so that the LLM now makes a better attempt the second time. If the LLM truly sucks, it can fail after 3 attempts and the next fallback can be passed downstream.

[apandy02]

Yeah, I agree with you @prrao87. I think to actually implement this well we'd have to write it within ell instead of as a wrapper around it. In the meantime since I need this functionality for my use case, I made the following addition to my decorator to provide error context to the LMP:

def retry(
    model_class: Type[BaseModel],
    retries: int = 3,
    delay: float = 0.5
):
    """
    A basic decorator that retries the function if Pydantic validation fails.
    Must be applied AFTER the ell.simple or ell.complex decorator.
    
    :param model_class: The Pydantic model class used to validate the function output
    :param retries: Maximum number of retries before giving up
    :param delay: Delay in seconds between retries
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(retries):
                try:
                    result = func(*args, **kwargs, error_context=str(last_exception))
                    model_class.model_validate(json.loads(result)) 
                    return result
                except ValidationError as validation_error:
                    last_exception = validation_error
                    if attempt < retries - 1:
                        time.sleep(delay)
            
            if last_exception is not None:
                raise last_exception
                
        return wrapper
    return decorator

Which can be used as:

@retry(PydanticClass)
@ell.simple(model="gpt-4o-mini-2024-07-18", max_tokens=2048)
def analysis_simple(sys_message: str, user_data: str, error_context: str | None = None) -> str:
    sys_message += f"You must absolutely respond in this format as a json string with no exceptions: {PydanticClass.model_json_schema()}"
    if error_context is not None:
        sys_message += f"Make sure not to trigger the following error: {error_context}"
    
    return [
        ell.system(sys_message),
        ell.user(user_data)
    ]

This is very hacky, but I plan on adding a check to the decorator to ensure the function being called has this extra error context param and using it until we have a better solution here.