mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2025-09-02 18:50:24 +00:00
63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, Awaitable, Literal, Protocol
|
|
|
|
from skyvern.forge.sdk.models import Step
|
|
from skyvern.forge.sdk.settings_manager import SettingsManager
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LLMConfig:
|
|
model_name: str
|
|
required_env_vars: list[str]
|
|
supports_vision: bool
|
|
add_assistant_prefix: bool
|
|
|
|
def get_missing_env_vars(self) -> list[str]:
|
|
missing_env_vars = []
|
|
for env_var in self.required_env_vars:
|
|
env_var_value = getattr(SettingsManager.get_settings(), env_var, None)
|
|
if not env_var_value:
|
|
missing_env_vars.append(env_var)
|
|
|
|
return missing_env_vars
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LLMRouterModelConfig:
|
|
model_name: str
|
|
# https://litellm.vercel.app/docs/routing
|
|
litellm_params: dict[str, Any]
|
|
tpm: int | None = None
|
|
rpm: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LLMRouterConfig(LLMConfig):
|
|
model_list: list[LLMRouterModelConfig]
|
|
# All three redis parameters are required. Even if there isn't a password, it should be an empty string.
|
|
main_model_group: str
|
|
redis_host: str | None = None
|
|
redis_port: int | None = None
|
|
redis_password: str | None = None
|
|
fallback_model_group: str | None = None
|
|
routing_strategy: Literal[
|
|
"simple-shuffle",
|
|
"least-busy",
|
|
"usage-based-routing",
|
|
"usage-based-routing-v2",
|
|
"latency-based-routing",
|
|
] = "usage-based-routing"
|
|
num_retries: int = 1
|
|
retry_delay_seconds: int = 15
|
|
set_verbose: bool = False
|
|
|
|
|
|
class LLMAPIHandler(Protocol):
|
|
def __call__(
|
|
self,
|
|
prompt: str,
|
|
step: Step | None = None,
|
|
screenshots: list[bytes] | None = None,
|
|
parameters: dict[str, Any] | None = None,
|
|
) -> Awaitable[dict[str, Any]]:
|
|
...
|