add gemeni pro support (#1307)

2024-12-02 15:51:18 -08:00
parent 7f6b2c0929
commit 57e6404881
4 changed files with 37 additions and 0 deletions
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -51,6 +51,7 @@ if not any(
        settings.ENABLE_AZURE,
        settings.ENABLE_AZURE_GPT4O_MINI,
        settings.ENABLE_BEDROCK,
+        settings.ENABLE_GEMINI,
    ]
 ):
    raise NoProviderEnabledError()
@@ -246,3 +247,15 @@ if settings.ENABLE_AZURE_GPT4O_MINI:
            add_assistant_prefix=False,
        ),
    )
+
+if settings.ENABLE_GEMINI:
+    LLMConfigRegistry.register_config(
+        "GEMINI_PRO",
+        LLMConfig(
+            "gemini/gemini-pro-vision",
+            ["GEMINI_API_KEY"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_output_tokens=8192,
+        ),
+    )