From b8f8f4fc245978fb9ff21946fbfc5923178e0742 Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Mon, 19 Oct 2026 14:59:44 +0200 Subject: [PATCH] Implemented vertexai provider --- src/providers/inference/registry.py | 2 + src/providers/inference/vertexai/README.md | 3 + src/providers/inference/vertexai/__init__.py | 4 + src/providers/inference/vertexai/provider.py | 135 ++++++++++++ .../providers/inference/vertexai/__init__.py | 1 + .../inference/vertexai/test_provider.py | 193 ++++++++++++++++++ 6 files changed, 338 insertions(+) create mode 100644 src/providers/inference/vertexai/provider.py create mode 100644 tests/unit/providers/inference/vertexai/__init__.py create mode 100644 tests/unit/providers/inference/vertexai/test_provider.py diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index 163022c86..8e2a45466 100644 --- a/src/providers/inference/registry.py +++ b/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, } diff --git a/src/providers/inference/vertexai/README.md b/src/providers/inference/vertexai/README.md index 32ace8618..9fe63a493 100644 --- a/src/providers/inference/vertexai/README.md +++ b/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. diff --git a/src/providers/inference/vertexai/__init__.py b/src/providers/inference/vertexai/__init__.py index 4039d7de8..25119e73a 100644 --- a/src/providers/inference/vertexai/__init__.py +++ b/src/providers/inference/vertexai/__init__.py @@ -1 +1,5 @@ """Vertex AI inference provider package.""" + +from providers.inference.vertexai.provider import VertexAIInferenceProvider + +__all__ = ["VertexAIInferenceProvider"] diff --git a/src/providers/inference/vertexai/provider.py b/src/providers/inference/vertexai/provider.py new file mode 100644 index 000000000..80b797ea3 --- /dev/null +++ b/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" diff --git a/tests/unit/providers/inference/vertexai/__init__.py b/tests/unit/providers/inference/vertexai/__init__.py new file mode 100644 index 000000000..3744e78f4 --- /dev/null +++ b/tests/unit/providers/inference/vertexai/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the Vertex AI inference provider.""" diff --git a/tests/unit/providers/inference/vertexai/test_provider.py b/tests/unit/providers/inference/vertexai/test_provider.py new file mode 100644 index 000000000..808f1202b --- /dev/null +++ b/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 + )