litellm cooldown config (#711)

This commit is contained in:
LawyZheng 2024-08-19 16:49:26 +08:00 committed by GitHub
parent aff5f442f9
commit 46a139bd27
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 10 additions and 0 deletions

View file

@ -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,
)

View file

@ -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):