rate_limiters.py

import time
import datetime
from google.api_core.exceptions import ResourceExhausted

def safe_invoke_with_retry_gemini(llm_with_tools, messages, max_retries=3, wait_seconds=60):
    """
    Safely invokes a Gemini LLM with automatic retries on rate limit errors.

    This function attempts to call the provided LLM with the given messages. If a ResourceExhausted
    (rate limit) error occurs, it waits for a specified number of seconds and retries, up to a maximum
    number of retries. Other exceptions are raised immediately.

    Args:
        llm_with_tools: The Gemini LLM instance with tools bound.
        messages (list): List of messages to send to the LLM.
        max_retries (int): Maximum number of retry attempts on rate limit errors.
        wait_seconds (int): Seconds to wait between retries.

    Returns:
        The result of llm_with_tools.invoke(messages) if successful.

    Raises:
        ResourceExhausted: If the maximum number of retries is reached due to rate limiting.
        Exception: Any other exception encountered during invocation.
    """
    for attempt in range(1, max_retries + 1):
        try:
            return llm_with_tools.invoke(messages)
        except ResourceExhausted as e:
            print(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
                  f"[Attempt {attempt}] Rate limit hit. Retrying in {wait_seconds} seconds...")
            if attempt == max_retries:
                print("Max retries reached. Raising exception.")
                raise e
            time.sleep(wait_seconds)
        except Exception as e:
            print(f"[Attempt {attempt}] Unexpected error: {e}")
            raise e  # Raise immediately for other types of errors