gpt4o mini pricing for massive (#957)

Skyvern-AI · Oct 11, 2024 · b536f1d · b536f1d
1 parent c239927
commit b536f1d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git 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
@@ -1,16 +1,18 @@
 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
 
 from skyvern.forge.sdk.models import Step
 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)