LCORE-4705: Implemented vertexai provider

src/providers/inference/registry.py

@@ -10,10 +10,12 @@
 from providers.inference.errors import ProviderNotFoundError
 from providers.inference.openai.provider import OpenAIInferenceProvider
 from providers.inference.types import ModelType, ProviderId, ProviderType
+from providers.inference.vertexai.provider import VertexAIInferenceProvider
 from utils.types import Singleton
 
 PROVIDER_CLASSES: Final[dict[ProviderType, type[InferenceProvider[Any, Any]]]] = {
     "openai": OpenAIInferenceProvider,
+    "vertexai": VertexAIInferenceProvider,
 }
 
 

src/providers/inference/vertexai/README.md

@@ -4,3 +4,6 @@
 
 Vertex AI inference provider package.
 
+## [provider.py](provider.py)
+
+Vertex AI runtime inference provider backed by Pydantic AI.

src/providers/inference/vertexai/__init__.py

@@ -1 +1,5 @@
 """Vertex AI inference provider package."""
+
+from providers.inference.vertexai.provider import VertexAIInferenceProvider
+
+__all__ = ["VertexAIInferenceProvider"]

src/providers/inference/vertexai/provider.py

@@ -0,0 +1,135 @@
+"""Vertex AI runtime inference provider backed by Pydantic AI."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic_ai.embeddings import EmbeddingModel
+from pydantic_ai.embeddings.google import GoogleEmbeddingModel
+from pydantic_ai.models import Model
+from pydantic_ai.models.google import GoogleModel
+from pydantic_ai.providers.google_cloud import GoogleCloudProvider
+from sentence_transformers import CrossEncoder
+
+from models.config import VertexAIInferenceProviderEntry
+from providers.inference.base import InferenceProvider
+from providers.inference.catalog import CatalogModel, DiscoveredModel
+from providers.inference.errors import ModelTypeNotSupportedError
+from providers.inference.model_type_inference import (
+    infer_model_type_from_pai_known_names,
+)
+from providers.inference.types import ModelType
+
+
+class VertexAIInferenceProvider(
+    InferenceProvider[VertexAIInferenceProviderEntry, GoogleCloudProvider]
+):
+    """Inference provider for Google Vertex AI (Google Cloud) APIs."""
+
+    def _build_provider(self) -> GoogleCloudProvider:
+        """Build the Pydantic AI Google Cloud provider from configuration.
+
+        Uses Application Default Credentials. Set GOOGLE_APPLICATION_CREDENTIALS
+        when a credentials file is required.
+
+        Returns:
+            A Pydantic AI Google Cloud provider using the configured project and
+            location.
+
+        """
+        config = self._entry.config
+        return GoogleCloudProvider(project=config.project, location=config.location)
+
+    async def _discover_models(self) -> list[DiscoveredModel]:
+        """List models from the Vertex AI model catalog.
+
+        Keeps each catalog resource name as the provider model id so Model Garden
+        publisher paths remain usable for inference.
+
+        Returns:
+            Discovered models with inferred types, before allow-list filtering.
+
+        """
+        discovered: list[DiscoveredModel] = []
+        pager = await self._get_provider.client.aio.models.list()
+        async for model in pager:
+            if not model.name:
+                continue
+            discovered.append(
+                DiscoveredModel(
+                    provider_model_id=model.name,
+                    model_type=self._infer_model_type(
+                        model.name,
+                        supported_actions=model.supported_actions,
+                    ),
+                )
+            )
+        return discovered
+
+    def _create_llm(self, model: CatalogModel) -> Model:
+        """Create a Google chat model for a catalog entry.
+
+        Args:
+            model: Catalog entry describing the LLM to instantiate.
+
+        Returns:
+            A Pydantic AI Google model bound to this provider.
+
+        """
+        return GoogleModel(
+            model.provider_model_id,
+            provider=self._get_provider,
+        )
+
+    def _create_embedding(self, model: CatalogModel) -> EmbeddingModel:
+        """Create a Google embedding model for a catalog entry.
+
+        Args:
+            model: Catalog entry describing the embedding model to instantiate.
+
+        Returns:
+            A Pydantic AI Google embedding model bound to this provider.
+
+        """
+        return GoogleEmbeddingModel(
+            model.provider_model_id,
+            provider=self._get_provider,
+        )
+
+    def _create_rerank(self, model: CatalogModel) -> CrossEncoder:
+        """Reject rerank requests; this provider has no rerank API.
+
+        Args:
+            model: Catalog entry that requested rerank support.
+
+        Raises:
+            ModelTypeNotSupportedError: Always, because rerank is unsupported.
+
+        """
+        _ = model
+        raise ModelTypeNotSupportedError("rerank", self.provider_id)
+
+    def _infer_model_type(
+        self,
+        provider_model_id: str,
+        **kwargs: Any,
+    ) -> ModelType:
+        """Infer model type from known PAI names, then Vertex supported actions.
+
+        Args:
+            provider_model_id: Model id from Vertex or from an allow-list entry.
+            **kwargs: Optional supported_actions from catalog discovery.
+
+        Returns:
+            embedding when the model id is a known embedding model or
+            supported_actions includes embedContent; otherwise llm.
+
+        """
+        short_id = provider_model_id.rsplit("/", 1)[-1]
+        if known := infer_model_type_from_pai_known_names(short_id):
+            return known
+
+        supported_actions = kwargs.get("supported_actions") or []
+        if "embedContent" in supported_actions:
+            return "embedding"
+        return "llm"

tests/unit/providers/inference/vertexai/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for the Vertex AI inference provider."""

tests/unit/providers/inference/vertexai/test_provider.py

@@ -0,0 +1,193 @@
+"""Unit tests for the Vertex AI inference provider."""
+
+# pylint: disable=redefined-outer-name,protected-access
+
+from typing import Any
+
+import pytest
+from pydantic_ai.embeddings.google import GoogleEmbeddingModel
+from pydantic_ai.models.google import GoogleModel
+from pydantic_ai.providers.google_cloud import GoogleCloudProvider
+
+from models.config import VertexAIInferenceProviderEntry
+from providers.inference.catalog import CatalogModel
+from providers.inference.errors import ModelTypeNotSupportedError
+from providers.inference.vertexai.provider import VertexAIInferenceProvider
+
+
+class _AsyncIter:
+    """Minimal async iterator for mocked Vertex model list results."""
+
+    def __init__(self, items: list[Any]) -> None:
+        """Store items to yield."""
+        self._items = iter(items)
+
+    def __aiter__(self) -> "_AsyncIter":
+        """Return self as the async iterator."""
+        return self
+
+    async def __anext__(self) -> Any:
+        """Yield the next item or stop."""
+        try:
+            return next(self._items)
+        except StopIteration as exc:
+            raise StopAsyncIteration from exc
+
+
+@pytest.fixture
+def provider() -> VertexAIInferenceProvider:
+    """Create a Vertex AI provider."""
+    entry = VertexAIInferenceProviderEntry.model_validate(
+        {
+            "type": "vertexai",
+            "id": "vertexai",
+            "config": {
+                "project": "test-project",
+                "location": "us-central1",
+            },
+        }
+    )
+    return VertexAIInferenceProvider(entry)
+
+
+def test_build_provider_uses_config(provider: VertexAIInferenceProvider) -> None:
+    """Builds a Pydantic AI Google Cloud provider with project and location."""
+    pai_provider = provider._build_provider()
+
+    assert isinstance(pai_provider, GoogleCloudProvider)
+    assert pai_provider.name == "google-cloud"
+    assert pai_provider.client.vertexai is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("catalog", "expected_type"),
+    [
+        (
+            ("publishers/google/models/gemini-3.5-flash", ["generateContent"]),
+            "llm",
+        ),
+        (
+            ("publishers/google/models/text-embedding-005", ["embedContent"]),
+            "embedding",
+        ),
+        (
+            ("publishers/meta/models/llama-3.3-70b-instruct-maas", ["generateContent"]),
+            "llm",
+        ),
+        (
+            ("custom-model", None),
+            "llm",
+        ),
+    ],
+)
+async def test_discover_models_keeps_catalog_name(
+    provider: VertexAIInferenceProvider,
+    mocker: Any,
+    catalog: tuple[str, list[str] | None],
+    expected_type: str,
+) -> None:
+    """Keeps catalog resource names and types from known names or actions."""
+    name, supported_actions = catalog
+    catalog_model = mocker.MagicMock()
+    catalog_model.name = name
+    catalog_model.supported_actions = supported_actions
+
+    pai_provider = mocker.MagicMock(spec=GoogleCloudProvider)
+    pai_provider.client.aio.models.list = mocker.AsyncMock(
+        return_value=_AsyncIter([catalog_model])
+    )
+    provider._provider = pai_provider
+
+    discovered = await provider._discover_models()
+
+    assert len(discovered) == 1
+    assert discovered[0].provider_model_id == name
+    assert discovered[0].model_type == expected_type
+    pai_provider.client.aio.models.list.assert_awaited_once_with()
+
+
+def test_create_llm_uses_google_model(
+    provider: VertexAIInferenceProvider,
+) -> None:
+    """Creates a Pydantic AI Google model using the provider model ID."""
+    pai_provider = provider._build_provider()
+    provider._provider = pai_provider
+
+    model = provider._create_llm(
+        CatalogModel(
+            name="chat",
+            provider_model_id="gemini-3.5-flash",
+            model_type="llm",
+        )
+    )
+
+    assert isinstance(model, GoogleModel)
+    assert model.model_name == "gemini-3.5-flash"
+    assert model.provider is pai_provider
+
+
+def test_create_embedding_uses_google_embedding_model(
+    provider: VertexAIInferenceProvider,
+) -> None:
+    """Creates a Pydantic AI Google embedding model using the provider model ID."""
+    pai_provider = provider._build_provider()
+    provider._provider = pai_provider
+
+    model = provider._create_embedding(
+        CatalogModel(
+            name="embed",
+            provider_model_id="text-embedding-005",
+            model_type="embedding",
+        )
+    )
+
+    assert isinstance(model, GoogleEmbeddingModel)
+    assert model.model_name == "text-embedding-005"
+    assert model._provider is pai_provider
+
+
+def test_create_rerank_is_unsupported(
+    provider: VertexAIInferenceProvider,
+) -> None:
+    """Rejects reranking because the Vertex AI provider does not support it."""
+    with pytest.raises(ModelTypeNotSupportedError) as exc_info:
+        provider._create_rerank(
+            CatalogModel(
+                name="rerank",
+                provider_model_id="rerank-model",
+                model_type="rerank",
+            )
+        )
+
+    assert exc_info.value.model_type == "rerank"
+    assert exc_info.value.provider_id == "vertexai"
+
+
+@pytest.mark.parametrize(
+    ("model_id", "supported_actions", "expected_type"),
+    [
+        ("text-embedding-005", None, "embedding"),
+        ("publishers/google/models/text-embedding-005", None, "embedding"),
+        ("gemini-3.5-flash", ["embedContent"], "llm"),
+        ("publishers/google/models/gemini-3.5-flash", ["embedContent"], "llm"),
+        ("unknown-model", ["generateContent"], "llm"),
+        ("unknown-model", ["embedContent"], "embedding"),
+        ("unknown-model", None, "llm"),
+        ("unknown-model", [], "llm"),
+    ],
+)
+def test_infer_model_type_prefers_known_names_then_actions(
+    provider: VertexAIInferenceProvider,
+    model_id: str,
+    supported_actions: list[str] | None,
+    expected_type: str,
+) -> None:
+    """Prefers PAI known names (including path suffixes), then actions."""
+    assert (
+        provider._infer_model_type(
+            model_id,
+            supported_actions=supported_actions,
+        )
+        == expected_type
+    )