From b536f1d27f46e9a2e4cb0c772084618953e40079 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Thu, 10 Oct 2024 22:57:54 -0700 Subject: [PATCH] gpt4o mini pricing for massive (#957) --- skyvern/forge/sdk/api/llm/config_registry.py | 1 + skyvern/forge/sdk/api/llm/models.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py index dda1b4581..b9682e76c 100644 --- a/skyvern/forge/sdk/api/llm/config_registry.py +++ b/skyvern/forge/sdk/api/llm/config_registry.py @@ -210,6 +210,7 @@ def get_config(cls, llm_key: str) -> LLMRouterConfig | LLMConfig: api_base=SettingsManager.get_settings().AZURE_GPT4O_MINI_API_BASE, api_key=SettingsManager.get_settings().AZURE_GPT4O_MINI_API_KEY, api_version=SettingsManager.get_settings().AZURE_GPT4O_MINI_API_VERSION, + model_info={"model_name": "azure/gpt-4o-mini"}, ), supports_vision=True, add_assistant_prefix=False, diff --git a/skyvern/forge/sdk/api/llm/models.py b/skyvern/forge/sdk/api/llm/models.py index bc6de7f01..0d8ca4732 100644 --- a/skyvern/forge/sdk/api/llm/models.py +++ b/skyvern/forge/sdk/api/llm/models.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Any, Awaitable, Literal, Optional, Protocol, TypedDict +from typing import Any, Awaitable, Literal, Optional, Protocol from litellm import AllowedFailsPolicy @@ -7,10 +7,12 @@ from skyvern.forge.sdk.settings_manager import SettingsManager -class LiteLLMParams(TypedDict): +@dataclass +class LiteLLMParams: api_key: str | None api_version: str | None api_base: str | None + model_info: dict[str, Any] = field(default_factory=dict) @dataclass(frozen=True)