From 4bdc32ce2ccc803b4ed8390fcb333ca623774a55 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Tue, 4 Feb 2025 17:11:29 +0800
Subject: [PATCH] add azure o3 mini (#1711)

---
 skyvern/config.py                            |  7 ++++++
 skyvern/forge/sdk/api/llm/config_registry.py | 25 ++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/skyvern/config.py b/skyvern/config.py
index 0d564eb77f..215428b18b 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -103,6 +103,7 @@ class Settings(BaseSettings):
     ENABLE_ANTHROPIC: bool = False
     ENABLE_AZURE: bool = False
     ENABLE_AZURE_GPT4O_MINI: bool = False
+    ENABLE_AZURE_O3_MINI: bool = False
     ENABLE_BEDROCK: bool = False
     ENABLE_GEMINI: bool = False
     # OPENAI
@@ -121,6 +122,12 @@ class Settings(BaseSettings):
     AZURE_GPT4O_MINI_API_BASE: str | None = None
     AZURE_GPT4O_MINI_API_VERSION: str | None = None
 
+    # AZURE o3 mini
+    AZURE_O3_MINI_DEPLOYMENT: str | None = None
+    AZURE_O3_MINI_API_KEY: str | None = None
+    AZURE_O3_MINI_API_BASE: str | None = None
+    AZURE_O3_MINI_API_VERSION: str | None = None
+
     # GEMINI
     GEMINI_API_KEY: str | None = None
 
diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py
index 7811f3a96b..f9e31a8081 100644
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -279,6 +279,31 @@ def get_config(cls, llm_key: str) -> LLMRouterConfig | LLMConfig:
         ),
     )
 
+if settings.ENABLE_AZURE_O3_MINI:
+    LLMConfigRegistry.register_config(
+        "AZURE_OPENAI_O3_MINI",
+        LLMConfig(
+            f"azure/{settings.AZURE_O3_MINI_DEPLOYMENT}",
+            [
+                "AZURE_O3_MINI_DEPLOYMENT",
+                "AZURE_O3_MINI_API_KEY",
+                "AZURE_O3_MINI_API_BASE",
+                "AZURE_O3_MINI_API_VERSION",
+            ],
+            litellm_params=LiteLLMParams(
+                api_base=settings.AZURE_O3_MINI_API_BASE,
+                api_key=settings.AZURE_O3_MINI_API_KEY,
+                api_version=settings.AZURE_O3_MINI_API_VERSION,
+                model_info={"model_name": "azure/o3-mini"},
+            ),
+            supports_vision=False,
+            add_assistant_prefix=False,
+            max_completion_tokens=16384,
+            temperature=None,  # Temperature isn't supported in the O-model series
+            reasoning_effort="low",
+        ),
+    )
+
 if settings.ENABLE_GEMINI:
     LLMConfigRegistry.register_config(
         "GEMINI_PRO",