feat(vertex_ai): add native Vertex AI Interactions API support

2026-08-28 05:25:59 +00:00 · 2026-08-25 09:55:13 -07:00 · 2026-08-25 09:55:13 -07:00 · 530dab32b9
commit 530dab32b9
parent bb27bfd9a7
6 changed files with 395 additions and 0 deletions
--- a/litellm/init.py
+++ b/litellm/init.py
@ -1801,6 +1801,9 @@ if TYPE_CHECKING:
    from .llms.gemini.interactions.transformation import (
        GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
    )
+    from .llms.vertex_ai.interactions.transformation import (
+        VertexAIInteractionsConfig as VertexAIInteractionsConfig,
+    )
    from .llms.openai.chat.o_series_transformation import (
        OpenAIOSeriesConfig as OpenAIOSeriesConfig,
        OpenAIOSeriesConfig as OpenAIO1Config,
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@ -242,6 +242,7 @@ LLM_CONFIG_NAMES: Final = (
    "OpenRouterResponsesAPIConfig",
    "BedrockMantleResponsesAPIConfig",
    "GoogleAIStudioInteractionsConfig",
+    "VertexAIInteractionsConfig",
    "OpenAIOSeriesConfig",
    "AnthropicSkillsConfig",
    "BaseSkillsAPIConfig",
@ -977,6 +978,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
        ".llms.gemini.interactions.transformation",
        "GoogleAIStudioInteractionsConfig",
    ),
+    "VertexAIInteractionsConfig": (
+        ".llms.vertex_ai.interactions.transformation",
+        "VertexAIInteractionsConfig",
+    ),
    "OpenAIOSeriesConfig": (
        ".llms.openai.chat.o_series_transformation",
        "OpenAIOSeriesConfig",
--- a/litellm/interactions/utils.py
+++ b/litellm/interactions/utils.py
@ -47,6 +47,13 @@ def get_provider_interactions_api_config(

        return GoogleAIStudioInteractionsConfig()

+    if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value):
+        from litellm.llms.vertex_ai.interactions.transformation import (
+            VertexAIInteractionsConfig,
+        )
+
+        return VertexAIInteractionsConfig()
+
    return None


--- a/litellm/llms/vertex_ai/interactions/init.py
+++ b/litellm/llms/vertex_ai/interactions/init.py
--- a/litellm/llms/vertex_ai/interactions/transformation.py
+++ b/litellm/llms/vertex_ai/interactions/transformation.py
@ -0,0 +1,149 @@
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from litellm.litellm_core_utils.url_utils import encode_url_path_segment
+from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig
+from litellm.llms.vertex_ai.common_utils import validate_vertex_location
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1"
+VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global"
+
+
+@dataclass(frozen=True, slots=True)
+class VertexInteractionsTarget:
+    base_url: str
+    project_id: str
+    location: str
+
+    @property
+    def collection_url(self) -> str:
+        return (
+            f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}"
+            f"/projects/{self.project_id}/locations/{self.location}/interactions"
+        )
+
+    def interaction_url(self, interaction_id: str) -> str:
+        encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id")
+        return f"{self.collection_url}/{encoded_interaction_id}"
+
+
+class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig):
+    def __init__(
+        self,
+        mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None,
+    ) -> None:
+        super().__init__()
+        self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = (
+            mint_access_token or self._mint_access_token_with_vertex_base
+        )
+
+    def _mint_access_token_with_vertex_base(
+        self,
+        credentials: VERTEX_CREDENTIALS_TYPES | None,
+        project_id: str | None,
+    ) -> tuple[str, str]:
+        return self._ensure_access_token(
+            credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai"
+        )
+
+    @property
+    def custom_llm_provider(self) -> LlmProviders:
+        return LlmProviders.VERTEX_AI
+
+    @property
+    def api_version(self) -> str:
+        return VERTEX_INTERACTIONS_API_VERSION
+
+    def get_default_vertex_location(self) -> str:
+        return VERTEX_INTERACTIONS_DEFAULT_LOCATION
+
+    def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]:
+        raw_params: Final = litellm_params.model_dump()
+        return self._mint_access_token(
+            self.safe_get_vertex_ai_credentials(raw_params),
+            self.safe_get_vertex_ai_project(raw_params),
+        )
+
+    def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget:
+        _, project_id = self._mint(litellm_params)
+        if not project_id:
+            raise ValueError(
+                "Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT"
+            )
+        location: Final = validate_vertex_location(
+            self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION
+        )
+        return VertexInteractionsTarget(
+            base_url=self.get_api_base(api_base or None, location),
+            project_id=project_id,
+            location=location,
+        )
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, str],
+        model: str,
+        litellm_params: GenericLiteLLMParams | None,
+    ) -> dict:  # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
+        access_token, _ = self._mint(litellm_params or GenericLiteLLMParams())
+        return {  # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {access_token}",
+            **headers,
+        }
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        model: str | None,
+        agent: str | None = None,
+        litellm_params: Mapping[str, object] | None = None,
+        stream: bool | None = None,
+    ) -> str:
+        params: Final = (
+            GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams()
+        )
+        collection_url: Final = self._target(api_base, params).collection_url
+        return f"{collection_url}?alt=sse" if stream else collection_url
+
+    def _interaction_by_id_request(
+        self,
+        interaction_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        url_suffix: str = "",
+    ) -> tuple[str, dict]:  # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+        target: Final = self._target(api_base or None, litellm_params)
+        return f"{target.interaction_url(interaction_id)}{url_suffix}", {}  # mutable-ok: same base contract
+
+    def transform_get_interaction_request(
+        self,
+        interaction_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: Mapping[str, str],
+    ) -> tuple[str, dict]:  # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+        return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
+
+    def transform_delete_interaction_request(
+        self,
+        interaction_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: Mapping[str, str],
+    ) -> tuple[str, dict]:  # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+        return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
+
+    def transform_cancel_interaction_request(
+        self,
+        interaction_id: str,
+        api_base: str,
+        litellm_params: GenericLiteLLMParams,
+        headers: Mapping[str, str],
+    ) -> tuple[str, dict]:  # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+        return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel")
--- a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
@ -0,0 +1,231 @@
+import pytest
+
+import litellm
+from litellm.interactions.utils import get_provider_interactions_api_config
+from litellm.llms.gemini.interactions.transformation import (
+    GoogleAIStudioInteractionsConfig,
+)
+from litellm.llms.vertex_ai.interactions.transformation import (
+    VertexAIInteractionsConfig,
+)
+from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions"
+
+
+class MinterRecorder:
+    def __init__(self, resolved_project: str = "creds-proj") -> None:
+        self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = []
+        self.resolved_project = resolved_project
+
+    def __call__(
+        self,
+        credentials: VERTEX_CREDENTIALS_TYPES | None,
+        project_id: str | None,
+    ) -> tuple[str, str]:
+        self.calls.append((credentials, project_id))
+        return "test-token", project_id or self.resolved_project
+
+
+@pytest.fixture
+def minter():
+    return MinterRecorder()
+
+
+@pytest.fixture
+def config(minter):
+    return VertexAIInteractionsConfig(mint_access_token=minter)
+
+
+@pytest.fixture
+def litellm_params():
+    return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json")
+
+
+class TestRegistration:
+    def test_vertex_ai_returns_vertex_config(self):
+        assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig)
+
+    def test_vertex_ai_beta_returns_vertex_config(self):
+        assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig)
+
+    def test_gemini_still_returns_google_ai_studio_config(self):
+        gemini_config = get_provider_interactions_api_config("gemini")
+        assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig)
+        assert not isinstance(gemini_config, VertexAIInteractionsConfig)
+
+    def test_lazy_import_resolves(self):
+        assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig
+
+    def test_custom_llm_provider_is_vertex_ai(self, config):
+        assert config.custom_llm_provider == LlmProviders.VERTEX_AI
+
+
+class TestValidateEnvironment:
+    def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params):
+        headers = config.validate_environment(
+            headers={},
+            model="gemini-omni-flash-preview",
+            litellm_params=litellm_params,
+        )
+
+        assert headers["Authorization"] == "Bearer test-token"
+        assert headers["Content-Type"] == "application/json"
+        assert "x-goog-api-key" not in headers
+        assert "Api-Revision" not in headers
+        assert minter.calls == [("creds.json", "test-proj")]
+
+    def test_caller_authorization_wins(self, config, litellm_params):
+        headers = config.validate_environment(
+            headers={"Authorization": "Bearer caller-token"},
+            model="gemini-omni-flash-preview",
+            litellm_params=litellm_params,
+        )
+
+        assert headers["Authorization"] == "Bearer caller-token"
+
+
+class TestGetCompleteUrl:
+    def test_defaults_to_global_v1beta1(self, config, litellm_params):
+        url = config.get_complete_url(
+            api_base=None,
+            model="gemini-omni-flash-preview",
+            litellm_params=dict(litellm_params),
+        )
+
+        assert url == GLOBAL_BASE
+
+    def test_stream_appends_alt_sse(self, config, litellm_params):
+        url = config.get_complete_url(
+            api_base=None,
+            model="gemini-omni-flash-preview",
+            litellm_params=dict(litellm_params),
+            stream=True,
+        )
+
+        assert url == f"{GLOBAL_BASE}?alt=sse"
+
+    def test_multi_region_location_uses_rep_host(self, config):
+        url = config.get_complete_url(
+            api_base=None,
+            model="gemini-omni-flash-preview",
+            litellm_params={"vertex_project": "test-proj", "vertex_location": "us"},
+        )
+
+        assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions"
+
+    def test_regional_location_uses_regional_host(self, config):
+        url = config.get_complete_url(
+            api_base=None,
+            model="gemini-omni-flash-preview",
+            litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"},
+        )
+
+        assert url == (
+            "https://us-central1-aiplatform.googleapis.com"
+            "/v1beta1/projects/test-proj/locations/us-central1/interactions"
+        )
+
+    def test_location_env_fallback_is_ignored(self, config, monkeypatch):
+        monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5")
+
+        url = config.get_complete_url(
+            api_base=None,
+            model="gemini-omni-flash-preview",
+            litellm_params={"vertex_project": "test-proj"},
+        )
+
+        assert url == GLOBAL_BASE
+
+    def test_api_base_override(self, config, litellm_params):
+        url = config.get_complete_url(
+            api_base="https://proxy.example.test",
+            model="gemini-omni-flash-preview",
+            litellm_params=dict(litellm_params),
+        )
+
+        assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions"
+
+    def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch):
+        monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
+        monkeypatch.setattr(litellm, "vertex_project", None)
+
+        url = config.get_complete_url(
+            api_base=None,
+            model="gemini-omni-flash-preview",
+            litellm_params={"vertex_credentials": "creds.json"},
+        )
+
+        assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions"
+
+    def test_invalid_location_rejected(self, config):
+        with pytest.raises(ValueError, match="Invalid vertex_location"):
+            config.get_complete_url(
+                api_base=None,
+                model="gemini-omni-flash-preview",
+                litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"},
+            )
+
+    def test_missing_project_rejected(self, monkeypatch):
+        monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
+        monkeypatch.setattr(litellm, "vertex_project", None)
+
+        def unresolved_minter(
+            credentials: VERTEX_CREDENTIALS_TYPES | None,
+            project_id: str | None,
+        ) -> tuple[str, str]:
+            return "test-token", ""
+
+        with pytest.raises(ValueError, match="Vertex AI project is required"):
+            VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url(
+                api_base=None,
+                model="gemini-omni-flash-preview",
+                litellm_params={},
+            )
+
+
+class TestInteractionByIdRequests:
+    def test_get_url(self, config, litellm_params):
+        url, request_body = config.transform_get_interaction_request(
+            interaction_id="abc123",
+            api_base="",
+            litellm_params=litellm_params,
+            headers={},
+        )
+
+        assert url == f"{GLOBAL_BASE}/abc123"
+        assert request_body == {}
+
+    def test_get_url_encodes_interaction_id(self, config, litellm_params):
+        url, _ = config.transform_get_interaction_request(
+            interaction_id="id/with space",
+            api_base="",
+            litellm_params=litellm_params,
+            headers={},
+        )
+
+        assert url == f"{GLOBAL_BASE}/id%2Fwith%20space"
+
+    def test_delete_url(self, config, litellm_params):
+        url, request_body = config.transform_delete_interaction_request(
+            interaction_id="abc123",
+            api_base="",
+            litellm_params=litellm_params,
+            headers={},
+        )
+
+        assert url == f"{GLOBAL_BASE}/abc123"
+        assert request_body == {}
+
+    def test_cancel_url(self, config, litellm_params):
+        url, request_body = config.transform_cancel_interaction_request(
+            interaction_id="abc123",
+            api_base="",
+            litellm_params=litellm_params,
+            headers={},
+        )
+
+        assert url == f"{GLOBAL_BASE}/abc123:cancel"
+        assert request_body == {}