From 6db01f4783b076e0b4d4cd7e32e85e27279b5328 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Mon, 2 Dec 2024 15:59:59 -0800 Subject: [PATCH] gemini 1.5 pro & flash (#1308) --- README.md | 3 ++- skyvern/forge/sdk/api/llm/config_registry.py | 12 +++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 42c68977..9e047fb6 100644 --- a/README.md +++ b/README.md @@ -285,7 +285,8 @@ More extensive documentation can be found on our [documentation website](https:/ | `ENABLE_AZURE` | Register Azure OpenAI models | Boolean | `true`, `false` | | `ENABLE_BEDROCK` | Register AWS Bedrock models| Boolean | `true`, `false` | | `ENABLE_GEMINI` | Register Gemini models| Boolean | `true`, `false` | -| `LLM_KEY` | The name of the model you want to use | String | Currently supported llm keys: `OPENAI_GPT4_TURBO`, `OPENAI_GPT4V`, `OPENAI_GPT4O`, `OPENAI_GPT4O_MINI`, `ANTHROPIC_CLAUDE3`, `ANTHROPIC_CLAUDE3_OPUS`, `ANTHROPIC_CLAUDE3_SONNET`, `ANTHROPIC_CLAUDE3_HAIKU`, `ANTHROPIC_CLAUDE3.5_SONNET`, `BEDROCK_ANTHROPIC_CLAUDE3_OPUS`, `BEDROCK_ANTHROPIC_CLAUDE3_SONNET`, `BEDROCK_ANTHROPIC_CLAUDE3_HAIKU`, `BEDROCK_ANTHROPIC_CLAUDE3.5_SONNET`, `AZURE_OPENAI` | +| `LLM_KEY` | The name of the model you want to use | String | Currently supported llm keys: `OPENAI_GPT4_TURBO`, `OPENAI_GPT4V`, `OPENAI_GPT4O`, `OPENAI_GPT4O_MINI`, `ANTHROPIC_CLAUDE3`, `ANTHROPIC_CLAUDE3_OPUS`, `ANTHROPIC_CLAUDE3_SONNET`, `ANTHROPIC_CLAUDE3_HAIKU`, `ANTHROPIC_CLAUDE3.5_SONNET`, `BEDROCK_ANTHROPIC_CLAUDE3_OPUS`, `BEDROCK_ANTHROPIC_CLAUDE3_SONNET`, `BEDROCK_ANTHROPIC_CLAUDE3_HAIKU`, `BEDROCK_ANTHROPIC_CLAUDE3.5_SONNET`, `AZURE_OPENAI`, `GEMINI_PRO`, `GEMINI_FLUSH`| +| `SECONDARY_LLM_KEY` | The name of the model for mini agents skyvern runs with | String | Currently supported llm keys: `OPENAI_GPT4_TURBO`, `OPENAI_GPT4V`, `OPENAI_GPT4O`, `OPENAI_GPT4O_MINI`, `ANTHROPIC_CLAUDE3`, `ANTHROPIC_CLAUDE3_OPUS`, `ANTHROPIC_CLAUDE3_SONNET`, `ANTHROPIC_CLAUDE3_HAIKU`, `ANTHROPIC_CLAUDE3.5_SONNET`, `BEDROCK_ANTHROPIC_CLAUDE3_OPUS`, `BEDROCK_ANTHROPIC_CLAUDE3_SONNET`, `BEDROCK_ANTHROPIC_CLAUDE3_HAIKU`, `BEDROCK_ANTHROPIC_CLAUDE3.5_SONNET`, `AZURE_OPENAI`, `GEMINI_PRO`, `GEMINI_FLUSH`| | `OPENAI_API_KEY` | OpenAI API Key | String | `sk-1234567890` | | `OPENAI_API_BASE` | OpenAI API Base, optional | String | `https://openai.api.base` | | `OPENAI_ORGANIZATION` | OpenAI Organization ID, optional | String | `your-org-id` | diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py index 73ff61f1..4e4defdd 100644 --- 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,