From dbcc48833ef029512744ce8d119307cb74a9a782 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Sun, 4 May 2025 01:02:01 -0700 Subject: [PATCH] fix LiteLLMParams typing (#2285) --- skyvern/forge/sdk/api/llm/config_registry.py | 4 ++-- skyvern/forge/sdk/api/llm/models.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py index e82ca844..a8f3cfd9 100644 --- a/skyvern/forge/sdk/api/llm/config_registry.py +++ b/skyvern/forge/sdk/api/llm/config_registry.py @@ -747,7 +747,7 @@ if settings.ENABLE_VERTEX_AI: add_assistant_prefix=False, max_completion_tokens=65535, litellm_params=LiteLLMParams( - vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + vertex_credentials=settings.VERTEX_CREDENTIALS, ), ), ) @@ -760,7 +760,7 @@ if settings.ENABLE_VERTEX_AI: add_assistant_prefix=False, max_completion_tokens=65535, litellm_params=LiteLLMParams( - vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + vertex_credentials=settings.VERTEX_CREDENTIALS, ), ), ) diff --git a/skyvern/forge/sdk/api/llm/models.py b/skyvern/forge/sdk/api/llm/models.py index c7813055..2c54b5a4 100644 --- a/skyvern/forge/sdk/api/llm/models.py +++ b/skyvern/forge/sdk/api/llm/models.py @@ -9,11 +9,12 @@ from skyvern.forge.sdk.schemas.task_v2 import TaskV2, Thought from skyvern.forge.sdk.settings_manager import SettingsManager -class LiteLLMParams(TypedDict): +class LiteLLMParams(TypedDict, total=False): api_key: str | None api_version: str | None api_base: str | None model_info: dict[str, Any] | None + vertex_credentials: str | None @dataclass(frozen=True)