From 47ad15a8646cf03c201471046c0c808ba9d154b8 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Tue, 29 Apr 2025 03:11:56 +0800 Subject: [PATCH] add vertex ai (#2243) --- poetry.lock | 11 ++- pyproject.toml | 2 +- skyvern/config.py | 4 ++ skyvern/forge/sdk/api/llm/config_registry.py | 72 ++++++++++++++++++++ 4 files changed, 82 insertions(+), 7 deletions(-) diff --git a/poetry.lock b/poetry.lock index 33fa1c19..8e4f7bcf 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand. [[package]] name = "aioboto3" @@ -2839,14 +2839,13 @@ files = [ [[package]] name = "litellm" -version = "1.67.0" +version = "1.67.4.post1" 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.67.0-py3-none-any.whl", hash = "sha256:d297126f45eea8d8a3df9c0de1d9491ff20e78dab5d1aa3820602082501ba89e"}, - {file = "litellm-1.67.0.tar.gz", hash = "sha256:18439db292d85b1d886bfa35de9d999600ecc6b4fc1137f12e6810d2133c8cec"}, + {file = "litellm-1.67.4.post1.tar.gz", hash = "sha256:057f2505f82d8c3f83d705c375b0d1931de998b13e239a6b06e16ee351fda648"}, ] [package.dependencies] @@ -2864,7 +2863,7 @@ tokenizers = "*" [package.extras] 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-proxy-extras (==0.1.11)", "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)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0)", "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-proxy-extras (==0.1.13)", "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)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0)", "websockets (>=13.1.0,<14.0.0)"] [[package]] name = "mako" @@ -6829,4 +6828,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = "^3.11,<3.12" -content-hash = "926815050df2b2d2fbdb96ac5084cb0e19a628a04d29cbc78ea63936e11b213c" +content-hash = "48a40666bd3b6efbfe9b48fda58b3751d249152223a409c9afbfbfe4aa639a20" diff --git a/pyproject.toml b/pyproject.toml index fb9550f3..ccb2fa42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ python-multipart = "^0.0.6" toml = "^0.10.2" jinja2 = "^3.1.2" uvicorn = {extras = ["standard"], version = "^0.24.0.post1"} -litellm = ">=1.67.0" +litellm = ">=1.67.4" playwright = "1.46.0" pillow = "^10.1.0" starlette-context = "^0.3.6" diff --git a/skyvern/config.py b/skyvern/config.py index 9cfc637e..fbf18753 100644 --- a/skyvern/config.py +++ b/skyvern/config.py @@ -125,6 +125,7 @@ class Settings(BaseSettings): ENABLE_AZURE_O3_MINI: bool = False ENABLE_BEDROCK: bool = False ENABLE_GEMINI: bool = False + ENABLE_VERTEX_AI: bool = False ENABLE_AZURE_CUA: bool = False ENABLE_OPENAI_COMPATIBLE: bool = False # OPENAI @@ -205,6 +206,9 @@ class Settings(BaseSettings): # GEMINI GEMINI_API_KEY: str | None = None + # VERTEX_AI + VERTEX_CREDENTIALS: str | None = None + # NOVITA AI ENABLE_NOVITA: bool = False NOVITA_API_KEY: str | None = None diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py index b9a51c85..9d1bad69 100644 --- a/skyvern/forge/sdk/api/llm/config_registry.py +++ b/skyvern/forge/sdk/api/llm/config_registry.py @@ -733,6 +733,78 @@ if settings.ENABLE_NOVITA: ), ) +# Create a GCP service account WITH the Vertex AI API access enabled +# Get the credentials json file. See documentation: https://support.google.com/a/answer/7378726?hl=en +# 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 +if settings.ENABLE_VERTEX_AI: + LLMConfigRegistry.register_config( + "VERTEX_GEMINI_2.5_PRO_EXP_03_25", + LLMConfig( + "vertex_ai/gemini-2.5-pro-preview-03-25", + ["VERTEX_CREDENTIALS"], + supports_vision=True, + add_assistant_prefix=False, + max_completion_tokens=65535, + litellm_params=LiteLLMParams( + vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + ), + ), + ) + LLMConfigRegistry.register_config( + "VERTEX_GEMINI_2.5_PRO_PREVIEW_03_25", + LLMConfig( + "vertex_ai/gemini-2.5-pro-preview-03-25", + ["VERTEX_CREDENTIALS"], + supports_vision=True, + add_assistant_prefix=False, + max_completion_tokens=65535, + litellm_params=LiteLLMParams( + vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + ), + ), + ) + LLMConfigRegistry.register_config( + "VERTEX_GEMINI_FLASH_2_0", + LLMConfig( + "vertex_ai/gemini-2.0-flash-001", + ["VERTEX_CREDENTIALS"], + supports_vision=True, + add_assistant_prefix=False, + max_completion_tokens=8192, + litellm_params=LiteLLMParams( + vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + ), + ), + ) + LLMConfigRegistry.register_config( + "VERTEX_GEMINI_PRO", + LLMConfig( + "vertex_ai/gemini-1.5-pro", + ["VERTEX_CREDENTIALS"], + supports_vision=True, + add_assistant_prefix=False, + max_completion_tokens=8192, + litellm_params=LiteLLMParams( + vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + ), + ), + ) + LLMConfigRegistry.register_config( + "VERTEX_GEMINI_FLASH", + LLMConfig( + "vertex_ai/gemini-1.5-flash", + ["VERTEX_CREDENTIALS"], + supports_vision=True, + add_assistant_prefix=False, + max_completion_tokens=8192, + litellm_params=LiteLLMParams( + vertex_credentials=settings.VERTEX_CREDENTIALS, # type: ignore + ), + ), + ) + + # Add support for dynamically configuring OpenAI-compatible LLM models # Based on liteLLM's support for OpenAI-compatible APIs # See documentation: https://docs.litellm.ai/docs/providers/openai_compatible