From 09a9c241618b98c6e1d7c9c3a8f7cdb56b4c19d8 Mon Sep 17 00:00:00 2001 From: Logan Solonche Date: Tue, 2 Jun 2026 16:33:36 -0400 Subject: [PATCH] Added VertexAI support, including GCP ADC Resolution --- .env.example | 7 +- README.md | 8 ++ model_registry.yaml | 21 ++++ pyproject.toml | 1 + src/skillspector/providers/__init__.py | 6 +- .../providers/nv_build/model_registry.yaml | 21 ++++ .../providers/openai/model_registry.yaml | 21 ++++ .../providers/vertexai/__init__.py | 20 ++++ .../providers/vertexai/model_registry.yaml | 30 +++++ .../providers/vertexai/provider.py | 107 ++++++++++++++++++ 10 files changed, 240 insertions(+), 2 deletions(-) create mode 100644 src/skillspector/providers/vertexai/__init__.py create mode 100644 src/skillspector/providers/vertexai/model_registry.yaml create mode 100644 src/skillspector/providers/vertexai/provider.py diff --git a/.env.example b/.env.example index 24ad4f5..01b60fb 100644 --- a/.env.example +++ b/.env.example @@ -2,7 +2,7 @@ ENV=dev # options: dev|s # Active LLM provider. Selects which provider answers credentials, # metadata, and default-model lookups. Leave unset to default to nv_build. -# Options: openai | anthropic | nv_build +# Options: openai | anthropic | vertexai | nv_build SKILLSPECTOR_PROVIDER= # Provider credentials — set the one matching SKILLSPECTOR_PROVIDER (or @@ -21,6 +21,11 @@ OPENAI_BASE_URL= # For SKILLSPECTOR_PROVIDER=anthropic. ANTHROPIC_API_KEY= +# For SKILLSPECTOR_PROVIDER=vertexai +GOOGLE_APPLICATION_CREDENTIALS= +GOOGLE_CLOUD_PROJECT= +GOOGLE_CLOUD_LOCATION= + # SkillSpector config SKILLSPECTOR_MODEL= # leave empty to use the active provider's bundled default (see README); set to override (e.g. gpt-5.2) # SKILLSPECTOR_MODEL_REGISTRY=./model_registry.yaml # optional override; defaults to each provider's bundled YAML in src/skillspector/providers/ diff --git a/README.md b/README.md index ab84623..c9038fd 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,14 @@ export SKILLSPECTOR_PROVIDER=nv_build export NVIDIA_INFERENCE_KEY=nvapi-... skillspector scan ./my-skill/ +# VertexAI (Google Cloud) +export SKILLSPECTOR_PROVIDER=vertexai +export GOOGLE_APPLICATION_CREDENTIALS=/path/to/service-account.json +export GOOGLE_CLOUD_PROJECT=your-project-id +export GOOGLE_CLOUD_LOCATION=us-central1 +export SKILLSPECTOR_MODEL=gemini-2.5-pro +skillspector scan ./my-skill/ + # Local Ollama or any OpenAI-compatible endpoint export SKILLSPECTOR_PROVIDER=openai export OPENAI_API_KEY=ollama diff --git a/model_registry.yaml b/model_registry.yaml index e1c2b8c..f528891 100644 --- a/model_registry.yaml +++ b/model_registry.yaml @@ -40,3 +40,24 @@ models: "openai/openai/gpt-5.3-chat": context_length: 128000 max_output_tokens: 16384 + + # Google Gemini models (via VertexAI or AI Studio OpenAI-compatible endpoints) + "gemini-2.0-flash": + context_length: 1048576 + max_output_tokens: 8192 + + "gemini-2.5-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-2.5-flash": + context_length: 1048576 + max_output_tokens: 65535 + + "gemini-3.1-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-3.5-flash": + context_length: 1048576 + max_output_tokens: 8192 diff --git a/pyproject.toml b/pyproject.toml index 9cd0e53..0c469f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "langchain-openai>=1.1.10", "langsmith>=0.7.30", "yara-python>=4.5.0", + "google-auth>=2.53.0", ] [project.optional-dependencies] diff --git a/src/skillspector/providers/__init__.py b/src/skillspector/providers/__init__.py index 78bdd17..0c26fa0 100644 --- a/src/skillspector/providers/__init__.py +++ b/src/skillspector/providers/__init__.py @@ -49,6 +49,10 @@ def _select_active_provider() -> ModelMetadataProvider: from .anthropic import AnthropicProvider return AnthropicProvider() + if name == "vertexai": + from .vertexai import VertexAIProvider + + return VertexAIProvider() if name == "nv_build": return NvBuildProvider() if name in ("nv_inference", ""): @@ -63,7 +67,7 @@ def _select_active_provider() -> ModelMetadataProvider: raise ValueError( f"Unknown SKILLSPECTOR_PROVIDER: {name!r}. " - "Expected one of: openai, anthropic, nv_build (or unset)." + "Expected one of: openai, anthropic, vertexai, nv_build (or unset)." ) diff --git a/src/skillspector/providers/nv_build/model_registry.yaml b/src/skillspector/providers/nv_build/model_registry.yaml index aeba04e..2485b6f 100644 --- a/src/skillspector/providers/nv_build/model_registry.yaml +++ b/src/skillspector/providers/nv_build/model_registry.yaml @@ -26,3 +26,24 @@ models: "openai/gpt-oss-120b": context_length: 128000 max_output_tokens: 16384 + + # Google Gemini models (via VertexAI or AI Studio OpenAI-compatible endpoints) + "gemini-2.0-flash": + context_length: 1048576 + max_output_tokens: 8192 + + "gemini-2.5-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-2.5-flash": + context_length: 1048576 + max_output_tokens: 65535 + + "gemini-3.1-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-3.5-flash": + context_length: 1048576 + max_output_tokens: 8192 diff --git a/src/skillspector/providers/openai/model_registry.yaml b/src/skillspector/providers/openai/model_registry.yaml index a4d2606..3145e61 100644 --- a/src/skillspector/providers/openai/model_registry.yaml +++ b/src/skillspector/providers/openai/model_registry.yaml @@ -12,3 +12,24 @@ models: "gpt-5.4": context_length: 1000000 max_output_tokens: 128000 + + # Google Gemini models (via VertexAI or AI Studio OpenAI-compatible endpoints) + "gemini-2.0-flash": + context_length: 1048576 + max_output_tokens: 8192 + + "gemini-2.5-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-2.5-flash": + context_length: 1048576 + max_output_tokens: 65535 + + "gemini-3.1-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-3.5-flash": + context_length: 1048576 + max_output_tokens: 8192 diff --git a/src/skillspector/providers/vertexai/__init__.py b/src/skillspector/providers/vertexai/__init__.py new file mode 100644 index 0000000..0099188 --- /dev/null +++ b/src/skillspector/providers/vertexai/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VertexAI provider package (VertexAI OpenAI-compatibility endpoint).""" + +from .provider import REGISTRY_PATH, VertexAIProvider + +__all__ = ["REGISTRY_PATH", "VertexAIProvider"] diff --git a/src/skillspector/providers/vertexai/model_registry.yaml b/src/skillspector/providers/vertexai/model_registry.yaml new file mode 100644 index 0000000..3fafa8c --- /dev/null +++ b/src/skillspector/providers/vertexai/model_registry.yaml @@ -0,0 +1,30 @@ +# Token-budget metadata for the VertexAIProvider. Bundled with the +# package; consulted whenever the active provider is VertexAIProvider. +# +# Format: +# models: +# "": +# context_length: # total context window in tokens (required) +# max_output_tokens: # model's max output cap (optional) + +models: + # Google Gemini models (via VertexAI) + "gemini-2.0-flash": + context_length: 1048576 + max_output_tokens: 8192 + + "gemini-2.5-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-2.5-flash": + context_length: 1048576 + max_output_tokens: 65535 + + "gemini-3.1-pro": + context_length: 1000000 + max_output_tokens: 65536 + + "gemini-3.5-flash": + context_length: 1048576 + max_output_tokens: 8192 diff --git a/src/skillspector/providers/vertexai/provider.py b/src/skillspector/providers/vertexai/provider.py new file mode 100644 index 0000000..1982614 --- /dev/null +++ b/src/skillspector/providers/vertexai/provider.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VertexAI provider — Gemini models via VertexAI OpenAI-compatible endpoint. + +Reads ``GOOGLE_APPLICATION_CREDENTIALS``, ``GOOGLE_CLOUD_PROJECT``, and +``GOOGLE_CLOUD_LOCATION`` for credentials and constructs the VertexAI +OpenAI-compatible endpoint URL. Uses Google Cloud Application Default +Credentials (ADC) to generate access tokens. Defaults to Gemini 2.5 Flash. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import google.auth +import google.auth.transport.requests + +from skillspector.providers import registry + +REGISTRY_PATH = str(Path(__file__).with_name("model_registry.yaml")) + + +class VertexAIProvider: + """Stock VertexAI credentials + bundled-YAML metadata provider.""" + + DEFAULT_MODEL = "gemini-2.5-flash" + SLOT_DEFAULTS: dict[str, str] = {} + + + def resolve_credentials(self) -> tuple[str, str | None] | None: + """Return ``(access_token, base_url)`` from Google Cloud credentials. + + Uses Application Default Credentials (ADC) via ``google.auth.default()``. + The access token is refreshed from the credentials object and returned + as the API key for the OpenAI-compatible client. + + Returns ``None`` when required environment variables are not set. + + Raises: + google.auth.exceptions.DefaultCredentialsError: When credentials + are configured but invalid or malformed. + ValueError: When project cannot be determined or token refresh fails. + """ + + project_id = os.environ.get("GOOGLE_CLOUD_PROJECT", "").strip() + location = os.environ.get("GOOGLE_CLOUD_LOCATION", "").strip() + + if not project_id or not location: + return None + + # If we get here, the user explicitly configured VertexAI, + # so let authentication errors propagate for debugging + + + credentials, default_project = google.auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + + project = project_id or default_project + if not project: + raise ValueError( + "Could not determine GCP project. Ensure GOOGLE_CLOUD_PROJECT " + "is set or the credentials file contains a project ID." + ) + + credentials.refresh(google.auth.transport.requests.Request()) + + access_token = credentials.token + if not access_token: + raise ValueError( + "Failed to obtain access token from Google Cloud credentials. " + "Ensure GOOGLE_APPLICATION_CREDENTIALS points to a valid " + "service account key file." + ) + + # Construct the VertexAI OpenAI-compatible base URL + base_url = ( + f"https://{location}-aiplatform.googleapis.com/v1beta1/" + f"projects/{project}/locations/{location}/endpoints/openapi" + ) + + return access_token, base_url + + def get_context_length(self, model: str) -> int | None: + return registry.lookup_context_length(REGISTRY_PATH, model) + + def get_max_output_tokens(self, model: str) -> int | None: + return registry.lookup_max_output_tokens(REGISTRY_PATH, model) + + def resolve_model(self, slot: str = "default") -> str: + """Resolve model: ``SKILLSPECTOR_MODEL`` env > slot default > ``DEFAULT_MODEL``.""" + user_input = os.environ.get("SKILLSPECTOR_MODEL", "").strip() + return user_input or self.SLOT_DEFAULTS.get(slot, "") or self.DEFAULT_MODEL \ No newline at end of file