gemini 1.5 pro & flash (#1308)

2024-12-02 15:59:59 -08:00
parent 57e6404881
commit 6db01f4783
2 changed files with 13 additions and 2 deletions
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -252,7 +252,17 @@ if settings.ENABLE_GEMINI:
    LLMConfigRegistry.register_config(
        "GEMINI_PRO",
        LLMConfig(
-            "gemini/gemini-pro-vision",
+            "gemini/gemini-1.5-pro",
+            ["GEMINI_API_KEY"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_output_tokens=8192,
+        ),
+    )
+    LLMConfigRegistry.register_config(
+        "GEMINI_FLUSH",
+        LLMConfig(
+            "gemini/gemini-1.5-flash",
            ["GEMINI_API_KEY"],
            supports_vision=True,
            add_assistant_prefix=False,