feat: bundle google-cloud-aiplatform for Vertex partner models

openhands-agent-server/openhands/agent_server/agent-server.spec

@@ -8,6 +8,7 @@ import os
 import site
 import sys
 from PyInstaller.utils.hooks import (
+    collect_all,
     collect_submodules,
     collect_data_files,
     copy_metadata,
@@ -17,6 +18,58 @@ from PyInstaller.utils.hooks import (
 # and cause LoadLibrary to fail at runtime with "Invalid access to memory location".
 IS_WINDOWS = sys.platform == "win32"
 
+# Optional Vertex AI bundle. google-cloud-aiplatform is an opt-in extra
+# (`openhands-sdk[vertex]`) and is NOT bundled in the default agent-server
+# build. To produce a binary that supports `vertex_ai/*` partner models
+# (MiniMax, Qwen, Kimi MaaS endpoints):
+#
+#   - Docker:    docker build --build-arg ENABLE_VERTEX=1 ...
+#   - From src:  uv sync --frozen --dev --no-editable --extra boto3 --extra vertex
+#                uv run pyinstaller .../agent-server.spec
+#
+# When `vertexai` is importable we use collect_all(...) for the Vertex SDK
+# and its google.cloud.* namespace siblings: the imports happen inside
+# function bodies AND traverse PEP-420 google.cloud namespace packages, so
+# collect_submodules alone misses everything below the namespace root.
+# collect_all walks the actual installed dirs.
+import importlib.util as _vertex_importlib_util
+
+_VERTEX_AVAILABLE = _vertex_importlib_util.find_spec("vertexai") is not None
+
+_vertex_pkgs = (
+    "vertexai",
+    "google.cloud.aiplatform",
+    "google.cloud.aiplatform_v1",
+    "google.cloud.aiplatform_v1beta1",
+    "google.cloud.bigquery",
+    "google.cloud.storage",
+    "google.cloud.resourcemanager",
+    "google.api_core",
+    "google.auth",
+    "google.rpc",
+    "google.genai",
+    "proto",
+    "grpc_status",
+)
+_vertex_datas = []
+_vertex_binaries = []
+_vertex_hiddenimports = []
+if _VERTEX_AVAILABLE:
+    for _pkg in _vertex_pkgs:
+        _d, _b, _h = collect_all(_pkg)
+        _vertex_datas += _d
+        _vertex_binaries += _b
+        _vertex_hiddenimports += _h
+    # google.rpc.status_pb2 is a gRPC proto stub imported dynamically; only pin
+    # it when the SDK is actually present.
+    _vertex_hiddenimports.append("google.rpc.status_pb2")
+else:
+    print(
+        "[agent-server.spec] vertexai not installed; "
+        "skipping Vertex AI bundle collection. "
+        "Install openhands-sdk[vertex] before building to include it."
+    )
+
 # Get the project root directory (current working directory when running PyInstaller)
 project_root = Path.cwd()
 # Namespace roots must be in pathex so PyInstaller can find 'openhands/...'
@@ -64,7 +117,10 @@ def get_fakeredis_data():
 a = Analysis(
     [ENTRY],
     pathex=PATHEX,
-    binaries=[],
+    binaries=[
+        # Vertex AI SDK binaries (collected via collect_all above)
+        *_vertex_binaries,
+    ],
     datas=[
         # Third-party packages that ship data
         *collect_data_files("tiktoken"),
@@ -99,6 +155,9 @@ a = Analysis(
         *copy_metadata("openhands-workspace"),
         *copy_metadata("fastmcp"),
         *copy_metadata("litellm"),
+
+        # Vertex AI SDK datas (collected via collect_all above)
+        *_vertex_datas,
     ],
     hiddenimports=[
         # Pull all OpenHands modules from the namespace (PEP 420 safe once pathex is correct)
@@ -118,6 +177,10 @@ a = Analysis(
         # unicodedata.unidata_version (e.g. unicode17_0_0 on Python 3.13).
         *collect_submodules("rich"),
 
+        # Vertex AI SDK hidden imports (collected via collect_all above; empty
+        # if openhands-sdk[vertex] is not installed in the build env).
+        *_vertex_hiddenimports,
+
         # mcp subpackages used at runtime (avoid CLI)
         "mcp.types",
         "mcp.client",

openhands-agent-server/openhands/agent_server/docker/Dockerfile

@@ -8,6 +8,11 @@ ARG USERNAME=openhands
 ARG UID=10001
 ARG GID=10001
 ARG PORT=8000
+# Opt-in build flag for the Vertex AI extra (`openhands-sdk[vertex]`). Off by
+# default to keep the published image lean. Pass `--build-arg ENABLE_VERTEX=1`
+# to bundle google-cloud-aiplatform so the resulting binary supports
+# `vertex_ai/*` partner models (MiniMax, Qwen, Kimi MaaS endpoints).
+ARG ENABLE_VERTEX=0
 
 ####################################################################################
 # Builder (source mode)
@@ -26,7 +31,7 @@ ARG PORT=8000
 # See OpenHands/software-agent-sdk#2761.
 ####################################################################################
 FROM python:3.13-bookworm AS builder
-ARG USERNAME UID GID
+ARG USERNAME UID GID ENABLE_VERTEX
 ENV UV_PROJECT_ENVIRONMENT=/agent-server/.venv
 ENV UV_PYTHON_INSTALL_DIR=/agent-server/uv-managed-python
 
@@ -48,21 +53,25 @@ COPY --chown=${USERNAME}:${USERNAME} openhands-tools ./openhands-tools
 COPY --chown=${USERNAME}:${USERNAME} openhands-workspace ./openhands-workspace
 COPY --chown=${USERNAME}:${USERNAME} openhands-agent-server ./openhands-agent-server
 RUN --mount=type=cache,target=/home/${USERNAME}/.cache,uid=${UID},gid=${GID} \
+    EXTRA_FLAGS=""; \
+    if [ "$ENABLE_VERTEX" = "1" ]; then EXTRA_FLAGS="--extra vertex"; fi; \
     uv python install 3.13 && \
     uv venv --python-preference only-managed --python 3.13 .venv && \
-    uv sync --frozen --no-editable --managed-python --extra boto3 && \
+    uv sync --frozen --no-editable --managed-python --extra boto3 $EXTRA_FLAGS && \
     readlink -f .venv/bin/python | grep -q '^/agent-server/uv-managed-python/'
 
 ####################################################################################
 # Binary Builder (binary mode)
 # We run pyinstaller here to produce openhands-agent-server
 ####################################################################################
 FROM builder AS binary-builder
-ARG USERNAME UID GID
+ARG USERNAME UID GID ENABLE_VERTEX
 
 # We need --dev for pyinstaller
 RUN --mount=type=cache,target=/home/${USERNAME}/.cache,uid=${UID},gid=${GID} \
-    uv sync --frozen --dev --no-editable --extra boto3
+    EXTRA_FLAGS=""; \
+    if [ "$ENABLE_VERTEX" = "1" ]; then EXTRA_FLAGS="--extra vertex"; fi; \
+    uv sync --frozen --dev --no-editable --extra boto3 $EXTRA_FLAGS
 
 RUN --mount=type=cache,target=/home/${USERNAME}/.cache,uid=${UID},gid=${GID} \
     uv run pyinstaller openhands-agent-server/openhands/agent_server/agent-server.spec

openhands-sdk/openhands/sdk/llm/llm.py

@@ -113,6 +113,7 @@
 from openhands.sdk.llm.utils.model_features import get_features
 from openhands.sdk.llm.utils.retry_mixin import RetryMixin
 from openhands.sdk.llm.utils.telemetry import Telemetry
+from openhands.sdk.llm.utils.vertex_preflight import assert_vertex_sdk_available
 from openhands.sdk.logger import ENV_LOG_DIR, get_logger
 from openhands.sdk.utils.deprecation import warn_deprecated
 
@@ -1670,11 +1671,17 @@ def _get_litellm_api_key_value(self) -> str | None:
             assert isinstance(self.api_key, SecretStr)
             api_key_value = self.api_key.get_secret_value()
 
+        provider = self._infer_litellm_provider()
+
+        # vertex_ai partner models require the optional Vertex SDK extra; surface
+        # a friendly install hint before LiteLLM crashes with ModuleNotFoundError.
+        assert_vertex_sdk_available(provider)
+
         # LiteLLM treats api_key for Bedrock as an AWS bearer token.
         # Passing a non-Bedrock key (e.g. OpenAI/Anthropic) can cause Bedrock
         # to reject the request with an "Invalid API Key format" error.
         # For IAM/SigV4 auth (the default Bedrock path), do not forward api_key.
-        if api_key_value is not None and self._infer_litellm_provider() == "bedrock":
+        if api_key_value is not None and provider == "bedrock":
             return None
 
         return api_key_value

openhands-sdk/openhands/sdk/llm/utils/vertex_preflight.py

@@ -0,0 +1,36 @@
+"""Preflight check for Vertex AI partner-model dependencies.
+
+`google-cloud-aiplatform` is an optional extra (`openhands-sdk[vertex]`). When a
+caller targets a `vertex_ai/*` model without the extra installed, LiteLLM fails
+with a low-level `ModuleNotFoundError` from inside its provider handler. We
+catch that earlier and surface a friendly install hint instead.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+
+from openhands.sdk.llm.exceptions import LLMBadRequestError
+
+
+_INSTALL_HINT = (
+    "Vertex AI partner models require the Vertex SDK. "
+    'Install with: pip install "openhands-sdk[vertex]"'
+)
+
+
+def _vertex_sdk_available() -> bool:
+    return importlib.util.find_spec("vertexai") is not None
+
+
+def assert_vertex_sdk_available(provider: str | None) -> None:
+    """Raise a friendly error if the caller is targeting Vertex without the SDK.
+
+    No-op for any non-`vertex_ai` provider, so it's safe to call unconditionally
+    from the transport path.
+    """
+    if provider != "vertex_ai":
+        return
+    if _vertex_sdk_available():
+        return
+    raise LLMBadRequestError(_INSTALL_HINT)

openhands-sdk/pyproject.toml

@@ -32,6 +32,7 @@ Documentation = "https://docs.openhands.dev/sdk"
 
 [project.optional-dependencies]
 boto3 = ["boto3>=1.35.0"]
+vertex = ["google-cloud-aiplatform>=1.38"]
 
 [build-system]
 requires = ["setuptools>=61.0", "wheel"]

tests/sdk/llm/test_vertex_preflight.py

@@ -0,0 +1,31 @@
+"""Tests for the Vertex AI optional-extra preflight check."""
+
+from __future__ import annotations
+
+import pytest
+
+from openhands.sdk.llm.exceptions import LLMBadRequestError
+from openhands.sdk.llm.utils import vertex_preflight
+from openhands.sdk.llm.utils.vertex_preflight import assert_vertex_sdk_available
+
+
+def test_noop_for_non_vertex_providers(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Even with the SDK absent, non-vertex providers must not raise.
+    monkeypatch.setattr(vertex_preflight, "_vertex_sdk_available", lambda: False)
+    assert_vertex_sdk_available(None)
+    assert_vertex_sdk_available("openai")
+    assert_vertex_sdk_available("bedrock")
+
+
+def test_passes_when_sdk_installed(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(vertex_preflight, "_vertex_sdk_available", lambda: True)
+    assert_vertex_sdk_available("vertex_ai")
+
+
+def test_raises_with_install_hint_when_sdk_missing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(vertex_preflight, "_vertex_sdk_available", lambda: False)
+    with pytest.raises(LLMBadRequestError) as excinfo:
+        assert_vertex_sdk_available("vertex_ai")
+    assert "openhands-sdk[vertex]" in str(excinfo.value)