From 197cc6069d4493199a3e5d0b66fa36582bf95606 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <shu@skyvern.com>
Date: Tue, 4 Feb 2025 09:03:48 +0000
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=84=20synced=20local=20'skyvern/'=20wi?=
 =?UTF-8?q?th=20remote=20'skyvern/'?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

<!-- ELLIPSIS_HIDDEN -->

> [!IMPORTANT]
> Add Azure O3 Mini support with new settings and configuration registration in `config.py` and `config_registry.py`.
>
>   - **Settings**:
>     - Add `ENABLE_AZURE_O3_MINI` flag in `Settings` class in `config.py`.
>     - Add Azure O3 Mini related settings: `AZURE_O3_MINI_DEPLOYMENT`, `AZURE_O3_MINI_API_KEY`, `AZURE_O3_MINI_API_BASE`, `AZURE_O3_MINI_API_VERSION`.
>   - **Configuration**:
>     - Register `AZURE_OPENAI_O3_MINI` configuration in `LLMConfigRegistry` in `config_registry.py`.
>     - Set `litellm_params` with `api_base`, `api_key`, `api_version`, and model info for Azure O3 Mini.
>     - Configure with `supports_vision=False`, `add_assistant_prefix=False`, `max_completion_tokens=16384`, `temperature=None`, `reasoning_effort="low"`.
>
> <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=Skyvern-AI%2Fskyvern-cloud&utm_source=github&utm_medium=referral)<sup> for e65e77213e2e9494722ea7341cec204e01f0f699. It will automatically update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->
---
 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",