litellm cooldown config (#711)

skyvern/forge/sdk/api/llm/api_handler_factory.py

@@ -46,6 +46,10 @@ class LLMAPIHandlerFactory:
             ),
             num_retries=llm_config.num_retries,
             retry_after=llm_config.retry_delay_seconds,
+            disable_cooldowns=llm_config.disable_cooldowns,
+            allowed_fails=llm_config.allowed_fails,
+            allowed_fails_policy=llm_config.allowed_fails_policy,
+            cooldown_time=llm_config.cooldown_time,
             set_verbose=(False if SettingsManager.get_settings().is_cloud_environment() else llm_config.set_verbose),
             enable_pre_call_checks=True,
         )

skyvern/forge/sdk/api/llm/models.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass, field
 from typing import Any, Awaitable, Literal, Optional, Protocol, TypedDict
 
+from litellm import AllowedFailsPolicy
+
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.settings_manager import SettingsManager
 
@@ -62,6 +64,10 @@ class LLMRouterConfig(LLMConfigBase):
     num_retries: int = 1
     retry_delay_seconds: int = 15
     set_verbose: bool = False
+    disable_cooldowns: bool | None = None
+    allowed_fails: int | None = None
+    allowed_fails_policy: AllowedFailsPolicy | None = None
+    cooldown_time: float | None = None
 
 
 class LLMAPIHandler(Protocol):