mirror of
https://github.com/andrewyng/openworker.git
synced 2026-09-03 04:49:26 +00:00
Add Google Vertex AI provider with per-family dispatch
gemini/ and claude/ ids reuse the native providers; openweight/ goes through the MaaS OpenAI-compat endpoint with an auto-refreshed google-auth bearer. Credentials: service-account JSON or Application Default Credentials.
This commit is contained in:
@@ -22,6 +22,7 @@ from .registry import (
|
||||
verify_provider_key,
|
||||
)
|
||||
from .router import ProviderRouter
|
||||
from .vertex_provider import VertexProvider
|
||||
|
||||
__all__ = [
|
||||
"AssistantTurn",
|
||||
@@ -33,6 +34,7 @@ __all__ = [
|
||||
"BedrockProvider",
|
||||
"GeminiProvider",
|
||||
"OpenAIProvider",
|
||||
"VertexProvider",
|
||||
"resolve_api_key",
|
||||
"capabilities_for",
|
||||
"ProviderRouter",
|
||||
|
||||
@@ -33,8 +33,8 @@ def capabilities_for(model: str) -> ModelCapabilities:
|
||||
# The family segment decides: Claude keeps its native capabilities; everything else
|
||||
# stays conservative until probed (Converse tool calling works across families, but
|
||||
# parallel calls and vision vary per model).
|
||||
if provider == "bedrock":
|
||||
if name.startswith("claude/"):
|
||||
if provider in ("bedrock", "vertex"):
|
||||
if name.startswith(("claude/", "gemini/")):
|
||||
return ModelCapabilities(
|
||||
tools=True, vision=True, pdf=True, parallel_tool_calls=True, streaming=True
|
||||
)
|
||||
|
||||
@@ -137,6 +137,26 @@ MATRIX: dict[str, ModelEntry] = {
|
||||
"bedrock:other/mistral.mistral-large-3-v1:0": ModelEntry(
|
||||
"Mistral Large 3 · AWS Bedrock"
|
||||
),
|
||||
# Vertex ids carry a family segment too (gemini/ and claude/ → native paths,
|
||||
# openweight/ → the MaaS OpenAI-compat endpoint, keeping the publisher segment).
|
||||
"vertex:gemini/gemini-3.1-pro-preview": ModelEntry(
|
||||
"Gemini 3.1 Pro · Vertex AI", _AGENTIC_VISION
|
||||
),
|
||||
"vertex:gemini/gemini-3.6-flash": ModelEntry(
|
||||
"Gemini 3.6 Flash · Vertex AI", _AGENTIC_VISION
|
||||
),
|
||||
"vertex:claude/claude-sonnet-4-6": ModelEntry(
|
||||
"Claude Sonnet 4.6 · Vertex AI", _AGENTIC_VISION
|
||||
),
|
||||
"vertex:claude/claude-haiku-4-5": ModelEntry(
|
||||
"Claude Haiku 4.5 · Vertex AI", _AGENTIC_VISION
|
||||
),
|
||||
"vertex:openweight/meta/llama-4-maverick-17b-128e-instruct-maas": ModelEntry(
|
||||
"Llama 4 Maverick · Vertex AI"
|
||||
),
|
||||
"vertex:openweight/qwen/qwen3-coder-480b-a35b-instruct-maas": ModelEntry(
|
||||
"Qwen3 Coder · Vertex AI"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -10,7 +10,8 @@ Today: `openai` (the default, with an optional custom endpoint that covers Azure
|
||||
`/openai/v1` and any OpenAI-compliant gateway), `anthropic` (native Messages API via
|
||||
`AnthropicProvider`), `gemini` (native Google GenAI API via `GeminiProvider`), `bedrock`
|
||||
(models in the user's own AWS account — Claude natively, everything else via Converse),
|
||||
and `ollama` (local, OpenAI-compatible `/v1`).
|
||||
`vertex` (the user's own GCP project — Gemini and Claude natively, open-weight via the
|
||||
MaaS endpoint), and `ollama` (local, OpenAI-compatible `/v1`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -24,6 +25,7 @@ from .base import ProviderClient
|
||||
from .bedrock_provider import BedrockProvider
|
||||
from .gemini_provider import GeminiProvider
|
||||
from .openai_provider import OpenAIProvider
|
||||
from .vertex_provider import VertexProvider
|
||||
|
||||
DEFAULT_OLLAMA_URL = "http://localhost:11434"
|
||||
|
||||
@@ -145,6 +147,20 @@ def _build_bedrock(profile: dict[str, Any], secrets: Any) -> ProviderClient:
|
||||
)
|
||||
|
||||
|
||||
def _build_vertex(profile: dict[str, Any], secrets: Any) -> ProviderClient:
|
||||
# Blank service_account_json → Application Default Credentials, resolved at call time.
|
||||
p = profile or {}
|
||||
|
||||
def get(key: str) -> Optional[str]:
|
||||
return (p.get(key) or "").strip() or None
|
||||
|
||||
return VertexProvider(
|
||||
project=get("project"),
|
||||
location=get("location"),
|
||||
service_account_json=get("service_account_json"),
|
||||
)
|
||||
|
||||
|
||||
def _build_ollama(profile: dict[str, Any], secrets: Any) -> ProviderClient:
|
||||
# Ollama's OpenAI-compatible endpoint ignores the key but the SDK requires a non-empty
|
||||
# string, so we pass a placeholder. `base_url` comes from the stored profile (or the default).
|
||||
@@ -319,6 +335,40 @@ DESCRIPTORS: list[ProviderDescriptor] = [
|
||||
blurb="Runs models inside your own AWS account. Claude uses Anthropic's native "
|
||||
"Bedrock path; every other model goes through the Converse API.",
|
||||
),
|
||||
ProviderDescriptor(
|
||||
name="vertex",
|
||||
title="Vertex AI (Google Cloud)",
|
||||
needs_key=True,
|
||||
fields=[
|
||||
ProviderField(
|
||||
"project",
|
||||
"GCP project ID",
|
||||
secret=False,
|
||||
placeholder="my-project-123",
|
||||
),
|
||||
ProviderField(
|
||||
"location",
|
||||
"Location",
|
||||
secret=False,
|
||||
placeholder="us-east5",
|
||||
help="The region your Vertex AI models are enabled in "
|
||||
"(Claude models: us-east5 or europe-west1).",
|
||||
),
|
||||
ProviderField(
|
||||
"service_account_json",
|
||||
"Service-account JSON (optional)",
|
||||
secret=True,
|
||||
required=False,
|
||||
help="Paste the JSON key or a path to it. Leave blank to use "
|
||||
"Application Default Credentials "
|
||||
"(`gcloud auth application-default login`).",
|
||||
),
|
||||
],
|
||||
build=_build_vertex,
|
||||
recommended_model="gemini/gemini-3.6-flash",
|
||||
blurb="Runs models inside your own Google Cloud project. Gemini and Claude use "
|
||||
"their native APIs; open-weight models go through the Vertex MaaS endpoint.",
|
||||
),
|
||||
# OpenAI-compatible vendors, listed as first-class providers so users don't need to know the
|
||||
# "point the OpenAI slot at a different endpoint" trick (owner call, 2026-07-04). Each keeps
|
||||
# its own key profile; the endpoint is prefilled and editable (regional variants in `help`).
|
||||
@@ -539,6 +589,58 @@ def _verify_bedrock(fields: dict[str, Any], timeout: float) -> dict[str, Any]:
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
def _verify_vertex(fields: dict[str, Any], timeout: float) -> dict[str, Any]:
|
||||
"""Resolve credentials (service account or ADC), mint a bearer, and list Google's
|
||||
publisher models in the given project/location — one cheap read-only call."""
|
||||
import httpx
|
||||
|
||||
from .vertex_provider import load_credentials
|
||||
|
||||
project = (fields.get("project") or "").strip()
|
||||
location = (fields.get("location") or "").strip()
|
||||
try:
|
||||
creds = load_credentials(fields.get("service_account_json"))
|
||||
if creds is None:
|
||||
import google.auth
|
||||
|
||||
creds, _ = google.auth.default(
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"]
|
||||
)
|
||||
from google.auth.transport.requests import Request
|
||||
|
||||
creds.refresh(Request())
|
||||
except Exception as exc:
|
||||
kind = exc.__class__.__name__
|
||||
if kind == "DefaultCredentialsError":
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "No Google Cloud credentials found — paste a service-account "
|
||||
"JSON, or run `gcloud auth application-default login` first.",
|
||||
}
|
||||
if kind in ("RefreshError", "MalformedError", "JSONDecodeError", "ValueError"):
|
||||
return {"ok": False, "error": "Google rejected the credentials."}
|
||||
return {"ok": False, "error": f"Couldn't load Google credentials ({kind})."}
|
||||
try:
|
||||
resp = httpx.get(
|
||||
f"https://{location}-aiplatform.googleapis.com/v1/projects/{project}"
|
||||
f"/locations/{location}/publishers/google/models",
|
||||
headers={"Authorization": f"Bearer {creds.token}"},
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as exc:
|
||||
return {"ok": False, "error": f"Couldn't reach Vertex AI ({exc.__class__.__name__})."}
|
||||
if resp.status_code < 300:
|
||||
return {"ok": True}
|
||||
if resp.status_code in (401, 403):
|
||||
return {
|
||||
"ok": False,
|
||||
"error": "Credentials work but lack Vertex AI access in this project.",
|
||||
}
|
||||
if resp.status_code == 404:
|
||||
return {"ok": False, "error": "Project or location not found on Vertex AI."}
|
||||
return {"ok": False, "error": f"Vertex AI returned HTTP {resp.status_code}."}
|
||||
|
||||
|
||||
def verify_provider_key(
|
||||
name: str,
|
||||
*,
|
||||
@@ -550,7 +652,7 @@ def verify_provider_key(
|
||||
"""Validate a provider's credentials with one cheap, read-only call (list models) — the same
|
||||
pattern connectors use to validate tokens. Transient: callers pass the key directly so a user
|
||||
can Test before saving. Never raises; returns {ok, error?}. Multi-field cloud providers
|
||||
(Bedrock) take their whole form via `fields`; everyone else uses api_key/base_url.
|
||||
(Bedrock, Vertex) take their whole form via `fields`; everyone else uses api_key/base_url.
|
||||
"""
|
||||
import httpx
|
||||
|
||||
@@ -558,6 +660,8 @@ def verify_provider_key(
|
||||
key = (api_key or "").strip()
|
||||
if name == "bedrock":
|
||||
return _verify_bedrock(fields or {}, timeout)
|
||||
if name == "vertex":
|
||||
return _verify_vertex(fields or {}, timeout)
|
||||
try:
|
||||
if name == "anthropic":
|
||||
resp = httpx.get(
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Google Vertex AI provider — one entry in Settings, three wire paths by model family.
|
||||
|
||||
Routed ids look like `vertex:<family>/<model id>`; the router strips `vertex:` and this
|
||||
provider splits the family segment, reusing an existing provider class per family:
|
||||
|
||||
- `gemini/…` → the native `GeminiProvider` over `genai.Client(vertexai=True)`.
|
||||
- `claude/…` → the native `AnthropicProvider` over the SDK's `AnthropicVertex` client.
|
||||
- `openweight/…` → `OpenAIProvider` against Vertex's OpenAI-compatible MaaS endpoint
|
||||
(Llama, Qwen, DeepSeek, …; ids keep their publisher segment, e.g. `openweight/meta/…`).
|
||||
|
||||
An id with no recognized family segment is best-effort routed by name (gemini* → Gemini,
|
||||
claude* → Claude, anything else → MaaS) so a raw id pasted without the add-model dropdown
|
||||
still works.
|
||||
|
||||
Credentials: an explicit service-account JSON (pasted content or a file path) when the
|
||||
profile has one, else Application Default Credentials (`gcloud auth application-default
|
||||
login`). The MaaS path authenticates with a google-auth bearer token that expires ~hourly —
|
||||
this wrapper refreshes it and rebuilds the OpenAI sub-client as needed; the two native SDK
|
||||
clients take the credentials object and refresh internally.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Optional
|
||||
|
||||
from .anthropic_provider import AnthropicProvider
|
||||
from .base import AssistantTurn, ModelCapabilities, ProviderClient
|
||||
from .capabilities import capabilities_for
|
||||
from .gemini_provider import GeminiProvider
|
||||
from .openai_provider import OpenAIProvider
|
||||
|
||||
_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||
|
||||
_FAMILIES = ("gemini", "claude", "openweight")
|
||||
|
||||
|
||||
def load_credentials(service_account_json: Optional[str]) -> Any:
|
||||
"""Explicit service-account JSON (content or path) → Credentials; blank → None (the
|
||||
SDKs and the token path then fall back to Application Default Credentials)."""
|
||||
raw = (service_account_json or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
from google.oauth2 import service_account
|
||||
|
||||
if raw.startswith("{"):
|
||||
info = json.loads(raw)
|
||||
return service_account.Credentials.from_service_account_info(
|
||||
info, scopes=_SCOPES
|
||||
)
|
||||
return service_account.Credentials.from_service_account_file(raw, scopes=_SCOPES)
|
||||
|
||||
|
||||
class VertexProvider(ProviderClient):
|
||||
"""Family dispatcher: splits `<family>/<model id>` and delegates to the sub-client."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
project: Optional[str] = None,
|
||||
location: Optional[str] = None,
|
||||
service_account_json: Optional[str] = None,
|
||||
credentials: Any = None,
|
||||
gemini_client: Optional[ProviderClient] = None,
|
||||
claude_client: Optional[ProviderClient] = None,
|
||||
openweight_client: Optional[ProviderClient] = None,
|
||||
):
|
||||
self._project = project
|
||||
self._location = location
|
||||
self._service_account_json = service_account_json
|
||||
self._credentials = credentials # test seam; normally resolved lazily
|
||||
# Test seams: pre-built sub-providers skip the SDK construction below.
|
||||
self._clients: dict[str, ProviderClient] = {}
|
||||
if gemini_client is not None:
|
||||
self._clients["gemini"] = gemini_client
|
||||
if claude_client is not None:
|
||||
self._clients["claude"] = claude_client
|
||||
if openweight_client is not None:
|
||||
self._clients["openweight"] = openweight_client
|
||||
self._openweight_injected = openweight_client is not None
|
||||
|
||||
@staticmethod
|
||||
def _split(model: str) -> tuple[str, str]:
|
||||
if "/" in model:
|
||||
family, rest = model.split("/", 1)
|
||||
if family in _FAMILIES:
|
||||
return family, rest
|
||||
# Raw id without a family segment: route by name, best effort.
|
||||
if model.startswith("gemini"):
|
||||
return "gemini", model
|
||||
if model.startswith("claude"):
|
||||
return "claude", model
|
||||
return "openweight", model
|
||||
|
||||
# -- credentials -------------------------------------------------------------
|
||||
def _explicit_credentials(self) -> Any:
|
||||
"""The service-account credentials, or None to let each SDK use ADC."""
|
||||
if self._credentials is None:
|
||||
self._credentials = load_credentials(self._service_account_json)
|
||||
return self._credentials
|
||||
|
||||
def _bearer_credentials(self) -> Any:
|
||||
"""Credentials for the MaaS bearer token: explicit service account, else ADC."""
|
||||
creds = self._explicit_credentials()
|
||||
if creds is None:
|
||||
import google.auth
|
||||
|
||||
try:
|
||||
creds, _ = google.auth.default(scopes=_SCOPES)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"No Google Cloud credentials found — paste a service-account JSON "
|
||||
"in Settings ▸ Models, or run `gcloud auth application-default login`."
|
||||
) from exc
|
||||
self._credentials = creds
|
||||
return creds
|
||||
|
||||
# -- family sub-clients --------------------------------------------------------
|
||||
def _family_client(self, family: str) -> ProviderClient:
|
||||
if family == "openweight":
|
||||
return self._openweight_client()
|
||||
client = self._clients.get(family)
|
||||
if client is None:
|
||||
if family == "gemini":
|
||||
from google import genai
|
||||
|
||||
client = GeminiProvider(
|
||||
client=genai.Client(
|
||||
vertexai=True,
|
||||
project=self._project,
|
||||
location=self._location,
|
||||
credentials=self._explicit_credentials(),
|
||||
)
|
||||
)
|
||||
else:
|
||||
from anthropic import AnthropicVertex
|
||||
|
||||
client = AnthropicProvider(
|
||||
client=AnthropicVertex(
|
||||
project_id=self._project,
|
||||
region=self._location,
|
||||
credentials=self._explicit_credentials(),
|
||||
)
|
||||
)
|
||||
self._clients[family] = client
|
||||
return client
|
||||
|
||||
def _openweight_client(self) -> ProviderClient:
|
||||
"""OpenAIProvider over the Vertex MaaS endpoint, rebuilt whenever the bearer
|
||||
token has to be refreshed (google-auth tokens expire ~hourly)."""
|
||||
if self._openweight_injected:
|
||||
return self._clients["openweight"]
|
||||
creds = self._bearer_credentials()
|
||||
if not getattr(creds, "valid", False):
|
||||
from google.auth.transport.requests import Request
|
||||
|
||||
creds.refresh(Request())
|
||||
self._clients.pop("openweight", None) # stale token — rebuild below
|
||||
client = self._clients.get("openweight")
|
||||
if client is None:
|
||||
base = (
|
||||
f"https://{self._location}-aiplatform.googleapis.com/v1/projects/"
|
||||
f"{self._project}/locations/{self._location}/endpoints/openapi"
|
||||
)
|
||||
client = OpenAIProvider(api_key=creds.token, base_url=base)
|
||||
self._clients["openweight"] = client
|
||||
return client
|
||||
|
||||
# -- ProviderClient -------------------------------------------------------------
|
||||
def complete(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: Optional[list[dict[str, Any]]] = None,
|
||||
**settings: Any,
|
||||
) -> AssistantTurn:
|
||||
family, rest = self._split(model)
|
||||
return self._family_client(family).complete(
|
||||
model=rest, messages=messages, tools=tools, **settings
|
||||
)
|
||||
|
||||
def stream(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: Optional[list[dict[str, Any]]] = None,
|
||||
**settings: Any,
|
||||
):
|
||||
family, rest = self._split(model)
|
||||
return self._family_client(family).stream(
|
||||
model=rest, messages=messages, tools=tools, **settings
|
||||
)
|
||||
|
||||
def capabilities(self, model: str) -> ModelCapabilities:
|
||||
qualified = model if model.startswith("vertex:") else f"vertex:{model}"
|
||||
return capabilities_for(qualified)
|
||||
@@ -0,0 +1,237 @@
|
||||
"""Google Vertex AI provider — 3-way family dispatch, bearer refresh, registry glue."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from coworker.providers import capabilities_for
|
||||
from coworker.providers.base import AssistantTurn, ProviderClient, StreamChunk
|
||||
from coworker.providers.vertex_provider import VertexProvider, load_credentials
|
||||
|
||||
# -- family dispatch ----------------------------------------------------------------
|
||||
|
||||
|
||||
class _Recorder(ProviderClient):
|
||||
def __init__(self):
|
||||
self.seen: list[str] = []
|
||||
|
||||
def complete(self, *, model, messages, tools=None, **settings):
|
||||
self.seen.append(model)
|
||||
return AssistantTurn(text="ok")
|
||||
|
||||
def stream(self, *, model, messages, tools=None, **settings):
|
||||
self.seen.append(model)
|
||||
yield StreamChunk(turn=AssistantTurn(text="ok"))
|
||||
|
||||
def capabilities(self, model):
|
||||
return capabilities_for(model)
|
||||
|
||||
|
||||
def _provider(**kw) -> tuple[VertexProvider, _Recorder, _Recorder, _Recorder]:
|
||||
gemini, claude, openweight = _Recorder(), _Recorder(), _Recorder()
|
||||
p = VertexProvider(
|
||||
project="proj",
|
||||
location="us-east5",
|
||||
gemini_client=gemini,
|
||||
claude_client=claude,
|
||||
openweight_client=openweight,
|
||||
**kw,
|
||||
)
|
||||
return p, gemini, claude, openweight
|
||||
|
||||
|
||||
def test_family_dispatch():
|
||||
p, gemini, claude, openweight = _provider()
|
||||
msgs = [{"role": "user", "content": "x"}]
|
||||
p.complete(model="gemini/gemini-3.6-flash", messages=msgs)
|
||||
p.complete(model="claude/claude-sonnet-4-6", messages=msgs)
|
||||
# Openweight ids keep their publisher segment — only the FIRST slash splits.
|
||||
p.complete(
|
||||
model="openweight/meta/llama-4-maverick-17b-128e-instruct-maas", messages=msgs
|
||||
)
|
||||
assert gemini.seen == ["gemini-3.6-flash"]
|
||||
assert claude.seen == ["claude-sonnet-4-6"]
|
||||
assert openweight.seen == ["meta/llama-4-maverick-17b-128e-instruct-maas"]
|
||||
|
||||
|
||||
def test_raw_ids_route_by_name():
|
||||
p, gemini, claude, openweight = _provider()
|
||||
msgs = [{"role": "user", "content": "x"}]
|
||||
p.complete(model="gemini-2.5-pro", messages=msgs)
|
||||
p.complete(model="claude-haiku-4-5", messages=msgs)
|
||||
p.complete(model="deepseek-ai/deepseek-v4-maas", messages=msgs)
|
||||
assert gemini.seen == ["gemini-2.5-pro"]
|
||||
assert claude.seen == ["claude-haiku-4-5"]
|
||||
assert openweight.seen == ["deepseek-ai/deepseek-v4-maas"]
|
||||
|
||||
|
||||
# -- openweight bearer refresh ---------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeCreds:
|
||||
"""google-auth-shaped credentials: `valid` flips true after refresh()."""
|
||||
|
||||
def __init__(self, token: str = "tok-1"):
|
||||
self.token = token
|
||||
self.valid = False
|
||||
self.refreshes = 0
|
||||
|
||||
def refresh(self, request):
|
||||
self.refreshes += 1
|
||||
self.token = f"tok-{self.refreshes + 1}"
|
||||
self.valid = True
|
||||
|
||||
|
||||
def test_openweight_builds_maas_endpoint_and_refreshes_bearer():
|
||||
creds = _FakeCreds()
|
||||
p = VertexProvider(project="proj", location="us-east5", credentials=creds)
|
||||
client = p._openweight_client()
|
||||
assert creds.refreshes == 1
|
||||
assert client._api_key == "tok-2"
|
||||
assert client._base_url == (
|
||||
"https://us-east5-aiplatform.googleapis.com/v1/projects/proj"
|
||||
"/locations/us-east5/endpoints/openapi"
|
||||
)
|
||||
# Token still valid → the same sub-client is reused, no extra refresh.
|
||||
assert p._openweight_client() is client
|
||||
assert creds.refreshes == 1
|
||||
# Token expired → refresh and rebuild with the new bearer.
|
||||
creds.valid = False
|
||||
rebuilt = p._openweight_client()
|
||||
assert rebuilt is not client
|
||||
assert creds.refreshes == 2
|
||||
assert rebuilt._api_key == "tok-3"
|
||||
|
||||
|
||||
# -- credentials ------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_load_credentials_blank_means_adc():
|
||||
assert load_credentials(None) is None
|
||||
assert load_credentials(" ") is None
|
||||
|
||||
|
||||
def test_load_credentials_bad_json_raises():
|
||||
with pytest.raises(Exception):
|
||||
load_credentials('{"type": "service_account"') # malformed JSON
|
||||
|
||||
|
||||
# -- capabilities / matrix ----------------------------------------------------------------
|
||||
|
||||
|
||||
def test_vertex_capabilities_from_matrix_and_fallback():
|
||||
assert capabilities_for("vertex:gemini/gemini-3.6-flash").vision
|
||||
assert capabilities_for("vertex:claude/claude-sonnet-4-6").pdf
|
||||
curated_ow = capabilities_for(
|
||||
"vertex:openweight/meta/llama-4-maverick-17b-128e-instruct-maas"
|
||||
)
|
||||
assert curated_ow.tools
|
||||
# Custom ids fall back on the family segment.
|
||||
assert capabilities_for("vertex:gemini/gemini-4.0-preview").vision
|
||||
custom_ow = capabilities_for("vertex:openweight/some-org/new-model-maas")
|
||||
assert custom_ow.tools and not custom_ow.parallel_tool_calls
|
||||
|
||||
|
||||
# -- registry / manager glue ----------------------------------------------------------------
|
||||
|
||||
|
||||
def test_vertex_descriptor_and_builder():
|
||||
from coworker.providers.registry import build_provider_client, get_descriptor
|
||||
|
||||
d = get_descriptor("vertex")
|
||||
assert d is not None and d.needs_key
|
||||
assert [f.key for f in d.fields] == ["project", "location", "service_account_json"]
|
||||
assert [f.key for f in d.fields if f.required] == ["project", "location"]
|
||||
sa = next(f for f in d.fields if f.key == "service_account_json")
|
||||
assert sa.secret and "Application Default Credentials" in sa.help
|
||||
|
||||
from coworker.providers.matrix import models_for_provider
|
||||
|
||||
assert d.recommended_model in models_for_provider("vertex")
|
||||
|
||||
p = build_provider_client(
|
||||
"vertex", {"project": "proj", "location": "europe-west1"}, None
|
||||
)
|
||||
assert isinstance(p, VertexProvider)
|
||||
assert p._project == "proj" and p._location == "europe-west1"
|
||||
|
||||
|
||||
def test_vertex_configured_needs_project_and_location():
|
||||
from coworker.providers.registry import descriptor_configured, get_descriptor
|
||||
|
||||
d = get_descriptor("vertex")
|
||||
assert not descriptor_configured(d, {})
|
||||
assert not descriptor_configured(d, {"project": "proj"})
|
||||
assert descriptor_configured(d, {"project": "proj", "location": "us-east5"})
|
||||
|
||||
|
||||
def test_router_routes_vertex_ids():
|
||||
from coworker.providers.router import ProviderRouter
|
||||
|
||||
router = ProviderRouter.__new__(ProviderRouter)
|
||||
model = "vertex:openweight/meta/llama-4-maverick-17b-128e-instruct-maas"
|
||||
assert router._provider_name(model) == "vertex"
|
||||
assert ProviderRouter._bare(model) == (
|
||||
"openweight/meta/llama-4-maverick-17b-128e-instruct-maas"
|
||||
)
|
||||
|
||||
|
||||
# -- verify ---------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _patch_verify(monkeypatch, creds: Any, status_code: Optional[int]):
|
||||
import httpx
|
||||
|
||||
import coworker.providers.vertex_provider as vp
|
||||
|
||||
monkeypatch.setattr(vp, "load_credentials", lambda raw: creds)
|
||||
captured: dict = {}
|
||||
|
||||
def fake_get(url, headers=None, timeout=None, **kw):
|
||||
captured["url"] = url
|
||||
captured["headers"] = headers
|
||||
|
||||
class _Resp:
|
||||
pass
|
||||
|
||||
resp = _Resp()
|
||||
resp.status_code = status_code
|
||||
return resp
|
||||
|
||||
monkeypatch.setattr(httpx, "get", fake_get)
|
||||
return captured
|
||||
|
||||
|
||||
def test_verify_vertex_ok(monkeypatch):
|
||||
from coworker.providers.registry import verify_provider_key
|
||||
|
||||
creds = _FakeCreds()
|
||||
captured = _patch_verify(monkeypatch, creds, 200)
|
||||
out = verify_provider_key(
|
||||
"vertex",
|
||||
fields={"project": "proj", "location": "us-east5", "service_account_json": "x"},
|
||||
)
|
||||
assert out == {"ok": True}
|
||||
assert creds.refreshes == 1
|
||||
assert "proj/locations/us-east5/publishers/google/models" in captured["url"]
|
||||
assert captured["headers"]["Authorization"] == "Bearer tok-2"
|
||||
|
||||
|
||||
def test_verify_vertex_maps_permission_errors(monkeypatch):
|
||||
from coworker.providers.registry import verify_provider_key
|
||||
|
||||
_patch_verify(monkeypatch, _FakeCreds(), 403)
|
||||
out = verify_provider_key(
|
||||
"vertex",
|
||||
fields={"project": "proj", "location": "us-east5", "service_account_json": "x"},
|
||||
)
|
||||
assert not out["ok"] and "Vertex AI access" in out["error"]
|
||||
|
||||
_patch_verify(monkeypatch, _FakeCreds(), 404)
|
||||
out = verify_provider_key(
|
||||
"vertex",
|
||||
fields={"project": "nope", "location": "us-east5", "service_account_json": "x"},
|
||||
)
|
||||
assert not out["ok"] and "not found" in out["error"]
|
||||
Reference in New Issue
Block a user