Spaces:
Running
Running
File size: 1,731 Bytes
3b31030 799013a 3b31030 c4604aa 3b31030 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
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 |