add gemini 2.0 flash (#1745)

2025-02-07 14:39:31 +08:00
parent f9296ba1d6
commit 22cefee5f9
1 changed files with 20 additions and 0 deletions
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -305,6 +305,26 @@ if settings.ENABLE_AZURE_O3_MINI:
    )

 if settings.ENABLE_GEMINI:
+    LLMConfigRegistry.register_config(
+        "GEMINI_FLASH_2_0",
+        LLMConfig(
+            "gemini/gemini-2.0-flash-001",
+            ["GEMINI_API_KEY"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_completion_tokens=8192,
+        ),
+    )
+    LLMConfigRegistry.register_config(
+        "GEMINI_FLASH_2_0_LITE",
+        LLMConfig(
+            "gemini/gemini-2.0-flash-lite-preview-02-05",
+            ["GEMINI_API_KEY"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_completion_tokens=8192,
+        ),
+    )
    LLMConfigRegistry.register_config(
        "GEMINI_PRO",
        LLMConfig(