|
from typing import Optional |
|
|
|
from smolagents import LiteLLMModel |
|
from tenacity import retry, stop_after_attempt, before_sleep_log, retry_if_exception_type, wait_exponential, wait_random |
|
import litellm |
|
import logging |
|
|
|
logging.basicConfig(level=logging.WARNING) |
|
logger = logging.getLogger(__name__) |
|
|
|
class LiteLLMModelWithBackOff(LiteLLMModel): |
|
def __init__(self, max_tokens: Optional[int] = 1500, *args, **kwargs): |
|
super().__init__(*args, **kwargs) |
|
self.max_tokens = max_tokens |
|
|
|
@retry( |
|
stop=stop_after_attempt(450), |
|
wait=wait_exponential(min=1, max=120, exp_base=2, multiplier=1) + wait_random(0, 5), |
|
before_sleep=before_sleep_log(logger, logging.WARNING), |
|
retry=retry_if_exception_type(( |
|
litellm.Timeout, |
|
litellm.RateLimitError, |
|
litellm.APIConnectionError, |
|
litellm.InternalServerError |
|
)) |
|
) |
|
def __call__(self, *args, **kwargs): |
|
return super().__call__(max_tokens=self.max_tokens, *args, **kwargs) |
|
|
|
|
