add gemini 2 5 support (#2850)

2025-07-01 14:38:17 +09:00
parent cb17dbbb6f
commit eb0e8a21ee
4 changed files with 64 additions and 35 deletions
--- a/poetry.lock
+++ b/poetry.lock
@@ -2714,22 +2714,6 @@ http2 = ["h2 (>=3,<5)"]
 socks = ["socksio (==1.*)"]
 zstd = ["zstandard (>=0.18.0)"]

-[[package]]
-name = "httpx-aiohttp"
-version = "0.1.4"
-description = "Aiohttp transport for HTTPX"
-optional = false
-python-versions = ">=3.9"
-groups = ["main"]
-files = [
-    {file = "httpx_aiohttp-0.1.4-py3-none-any.whl", hash = "sha256:740a8725af7b7a03d12f21ccd48a83457baa037304646589b87595746c05c87e"},
-    {file = "httpx_aiohttp-0.1.4.tar.gz", hash = "sha256:61030eed28deeac26286d2e872b7c167f5450b7b0eec5a617ae7d3f7da9c8684"},
-]
-
-[package.dependencies]
-aiohttp = ">=3,<4"
-httpx = ">=0.28.1,<1"
-
 [[package]]
 name = "httpx-sse"
 version = "0.4.0"
@@ -3774,21 +3758,20 @@ pyyaml-ft = {version = ">=8.0.0", markers = "python_version >= \"3.13\""}

 [[package]]
 name = "litellm"
-version = "1.71.1"
+version = "1.73.6"
 description = "Library to easily interface with LLM API providers"
 optional = false
 python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
 groups = ["main"]
 files = [
-    {file = "litellm-1.71.1-py3-none-any.whl", hash = "sha256:9b94e250c58fba3c87c6ebb77e33c1cc8aa9110cee99dfdc37b368a11cec57c7"},
-    {file = "litellm-1.71.1.tar.gz", hash = "sha256:c20e5917fdbe771ba4b6d1862b3d38d6e89cfba53e85bb337013f848256566eb"},
+    {file = "litellm-1.73.6-py3-none-any.whl", hash = "sha256:98b3c7f436e6521e280f98faf9bad06c4c76d6a1678db2b370ffa175c206d288"},
+    {file = "litellm-1.73.6.tar.gz", hash = "sha256:072ff9225aaa9caafe0fa9782df0ba4c93f893791cc55dd5ed218a401e788e88"},
 ]

 [package.dependencies]
-aiohttp = "*"
+aiohttp = ">=3.10"
 click = "*"
 httpx = ">=0.23.0"
-httpx-aiohttp = {version = ">=0.1.4", markers = "python_version >= \"3.9\""}
 importlib-metadata = ">=6.8.0"
 jinja2 = ">=3.1.2,<4.0.0"
 jsonschema = ">=4.22.0,<5.0.0"
@@ -3799,8 +3782,9 @@ tiktoken = ">=0.7.0"
 tokenizers = "*"

 [package.extras]
+caching = ["diskcache (>=5.6.1,<6.0.0)"]
 extra-proxy = ["azure-identity (>=1.15.0,<2.0.0)", "azure-keyvault-secrets (>=4.8.0,<5.0.0)", "google-cloud-kms (>=2.21.3,<3.0.0)", "prisma (==0.11.0)", "redisvl (>=0.4.1,<0.5.0) ; python_version >= \"3.9\" and python_version < \"3.14\"", "resend (>=0.8.0,<0.9.0)"]
-proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "boto3 (==1.34.34)", "cryptography (>=43.0.1,<44.0.0)", "fastapi (>=0.115.5,<0.116.0)", "fastapi-sso (>=0.16.0,<0.17.0)", "gunicorn (>=23.0.0,<24.0.0)", "litellm-enterprise (==0.1.6)", "litellm-proxy-extras (==0.2.0)", "mcp (==1.5.0) ; python_version >= \"3.10\"", "orjson (>=3.9.7,<4.0.0)", "pynacl (>=1.5.0,<2.0.0)", "python-multipart (>=0.0.18,<0.0.19)", "pyyaml (>=6.0.1,<7.0.0)", "rich (==13.7.1)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0) ; sys_platform != \"win32\"", "websockets (>=13.1.0,<14.0.0)"]
+proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "boto3 (==1.34.34)", "cryptography (>=43.0.1,<44.0.0)", "fastapi (>=0.115.5,<0.116.0)", "fastapi-sso (>=0.16.0,<0.17.0)", "gunicorn (>=23.0.0,<24.0.0)", "litellm-enterprise (==0.1.9)", "litellm-proxy-extras (==0.2.6)", "mcp (==1.9.3) ; python_version >= \"3.10\"", "orjson (>=3.9.7,<4.0.0)", "pynacl (>=1.5.0,<2.0.0)", "python-multipart (>=0.0.18,<0.0.19)", "pyyaml (>=6.0.1,<7.0.0)", "rich (==13.7.1)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0) ; sys_platform != \"win32\"", "websockets (>=13.1.0,<14.0.0)"]
 utils = ["numpydoc"]

 [[package]]
@@ -8862,4 +8846,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.11,<3.14"
-content-hash = "905f8c850623caf9a845defaa3c1592a5229e96b395b2632298db18e40453c98"
+content-hash = "807cb0c70058604538b422763a89b6f75869e570049d746a4e39a67194ead7d0"
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@ python-multipart = "^0.0.18"
 toml = "^0.10.2"
 jinja2 = "^3.1.2"
 uvicorn = {extras = ["standard"], version = "^0.24.0.post1"}
-litellm = ">=1.71.1"
+litellm = ">=1.73.6"
 playwright = [
    {version = ">1.46.0", python = ">=3.12,<3.14"},
    {version = "=1.46.0", python = ">=3.11,<3.12"}
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -292,9 +292,9 @@ class Settings(BaseSettings):

        if self.is_cloud_environment():
            return {
-                "gemini-2.5-pro-preview-05-06": {"llm_key": "VERTEX_GEMINI_2.5_PRO_PREVIEW", "label": "Gemini 2.5 Pro"},
+                "gemini-2.5-pro-preview-05-06": {"llm_key": "VERTEX_GEMINI_2.5_PRO", "label": "Gemini 2.5 Pro"},
                "gemini-2.5-flash-preview-05-20": {
-                    "llm_key": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
+                    "llm_key": "VERTEX_GEMINI_2.5_FLASH",
                    "label": "Gemini 2.5 Flash",
                },
                "azure/gpt-4.1": {"llm_key": "AZURE_OPENAI_GPT4_1", "label": "GPT 4.1"},
@@ -311,9 +311,9 @@ class Settings(BaseSettings):
        else:
            # TODO: apparently the list for OSS is to be much larger
            return {
-                "gemini-2.5-pro-preview-05-06": {"llm_key": "VERTEX_GEMINI_2.5_PRO_PREVIEW", "label": "Gemini 2.5 Pro"},
+                "gemini-2.5-pro-preview-05-06": {"llm_key": "VERTEX_GEMINI_2.5_PRO", "label": "Gemini 2.5 Pro"},
                "gemini-2.5-flash-preview-05-20": {
-                    "llm_key": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
+                    "llm_key": "VERTEX_GEMINI_2.5_FLASH",
                    "label": "Gemini 2.5 Flash",
                },
                "azure/gpt-4.1": {"llm_key": "AZURE_OPENAI_GPT4_1", "label": "GPT 4.1"},
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -652,6 +652,16 @@ if settings.ENABLE_GEMINI:
            max_completion_tokens=8192,
        ),
    )
+    LLMConfigRegistry.register_config(
+        "GEMINI_2.5_PRO",
+        LLMConfig(
+            "gemini/gemini-2.5-pro",
+            ["GEMINI_API_KEY"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_completion_tokens=65536,
+        ),
+    )
    LLMConfigRegistry.register_config(
        "GEMINI_2.5_PRO_PREVIEW",
        LLMConfig(
@@ -672,6 +682,16 @@ if settings.ENABLE_GEMINI:
            max_completion_tokens=65536,
        ),
    )
+    LLMConfigRegistry.register_config(
+        "GEMINI_2.5_FLASH",
+        LLMConfig(
+            "gemini/gemini-2.5-flash",
+            ["GEMINI_API_KEY"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_completion_tokens=65536,
+        ),
+    )
    LLMConfigRegistry.register_config(
        "GEMINI_2.5_FLASH_PREVIEW",
        LLMConfig(
@@ -856,18 +876,28 @@ if settings.ENABLE_NOVITA:
 # my_vertex_credentials = json.dumps(json.load(open("my_credentials_file.json")))
 # Set the value of my_vertex_credentials as the environment variable VERTEX_CREDENTIALS
 # NOTE: If you want to specify a location, make sure the model is available in the target location.
+# If you want to use the global location, you must set the VERTEX_PROJECT_ID environment variable.
 # See documentation: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#united-states
 if settings.ENABLE_VERTEX_AI and settings.VERTEX_CREDENTIALS:
-    if not settings.VERTEX_PROJECT_ID:
-        import json
-
-        credentials = json.loads(settings.VERTEX_CREDENTIALS)
-        settings.VERTEX_PROJECT_ID = credentials.get("project_id")
-
    api_base: str | None = None
-    if settings.VERTEX_LOCATION == "global":
+    if settings.VERTEX_LOCATION == "global" and settings.VERTEX_PROJECT_ID:
        api_base = f"https://aiplatform.googleapis.com/v1/projects/{settings.VERTEX_PROJECT_ID}/locations/global/publishers/google/models"

+    LLMConfigRegistry.register_config(
+        "VERTEX_GEMINI_2.5_PRO",
+        LLMConfig(
+            "vertex_ai/gemini-2.5-pro",
+            ["VERTEX_CREDENTIALS"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_completion_tokens=65535,
+            litellm_params=LiteLLMParams(
+                vertex_credentials=settings.VERTEX_CREDENTIALS,
+                api_base=f"{api_base}/gemini-2.5-pro" if api_base else None,
+                vertex_location=settings.VERTEX_LOCATION,
+            ),
+        ),
+    )
    LLMConfigRegistry.register_config(
        "VERTEX_GEMINI_2.5_PRO_PREVIEW",
        LLMConfig(
@@ -883,6 +913,21 @@ if settings.ENABLE_VERTEX_AI and settings.VERTEX_CREDENTIALS:
            ),
        ),
    )
+    LLMConfigRegistry.register_config(
+        "VERTEX_GEMINI_2.5_FLASH",
+        LLMConfig(
+            "vertex_ai/gemini-2.5-flash",
+            ["VERTEX_CREDENTIALS"],
+            supports_vision=True,
+            add_assistant_prefix=False,
+            max_completion_tokens=65535,
+            litellm_params=LiteLLMParams(
+                vertex_credentials=settings.VERTEX_CREDENTIALS,
+                api_base=f"{api_base}/gemini-2.5-flash" if api_base else None,
+                vertex_location=settings.VERTEX_LOCATION,
+            ),
+        ),
+    )
    LLMConfigRegistry.register_config(
        "VERTEX_GEMINI_2.5_FLASH_PREVIEW",
        LLMConfig(