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, num_retries=llm_config.num_retries,
retry_after=llm_config.retry_delay_seconds, 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), set_verbose=(False if SettingsManager.get_settings().is_cloud_environment() else llm_config.set_verbose),
enable_pre_call_checks=True, enable_pre_call_checks=True,
) )

View file

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