From 46a139bd27e7dd2d77e09473a3bea8ca633735bc Mon Sep 17 00:00:00 2001 From: LawyZheng Date: Mon, 19 Aug 2024 16:49:26 +0800 Subject: [PATCH] litellm cooldown config (#711) --- skyvern/forge/sdk/api/llm/api_handler_factory.py | 4 ++++ skyvern/forge/sdk/api/llm/models.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/skyvern/forge/sdk/api/llm/api_handler_factory.py b/skyvern/forge/sdk/api/llm/api_handler_factory.py index 0d24b165..494cdff2 100644 --- a/skyvern/forge/sdk/api/llm/api_handler_factory.py +++ b/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, ) diff --git a/skyvern/forge/sdk/api/llm/models.py b/skyvern/forge/sdk/api/llm/models.py index ae4f4986..bc6de7f0 100644 --- a/skyvern/forge/sdk/api/llm/models.py +++ b/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):