From 050cc894e7a85d49d24ea86f867ce2f305b3e258 Mon Sep 17 00:00:00 2001 From: Rohit C Prasad Date: Sat, 25 Jul 2026 16:19:00 -0700 Subject: [PATCH] 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. --- coworker/providers/__init__.py | 2 + coworker/providers/capabilities.py | 4 +- coworker/providers/matrix.py | 20 +++ coworker/providers/registry.py | 108 +++++++++++- coworker/providers/vertex_provider.py | 198 +++++++++++++++++++++ tests/test_vertex_provider.py | 237 ++++++++++++++++++++++++++ 6 files changed, 565 insertions(+), 4 deletions(-) create mode 100644 coworker/providers/vertex_provider.py create mode 100644 tests/test_vertex_provider.py diff --git a/coworker/providers/__init__.py b/coworker/providers/__init__.py index 0c526ab1..6c35ef95 100644 --- a/coworker/providers/__init__.py +++ b/coworker/providers/__init__.py @@ -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", diff --git a/coworker/providers/capabilities.py b/coworker/providers/capabilities.py index aea3921f..f096ff3a 100644 --- a/coworker/providers/capabilities.py +++ b/coworker/providers/capabilities.py @@ -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 ) diff --git a/coworker/providers/matrix.py b/coworker/providers/matrix.py index 60cff5c7..7a8daba6 100644 --- a/coworker/providers/matrix.py +++ b/coworker/providers/matrix.py @@ -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" + ), } diff --git a/coworker/providers/registry.py b/coworker/providers/registry.py index 165c027a..1318b86e 100644 --- a/coworker/providers/registry.py +++ b/coworker/providers/registry.py @@ -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( diff --git a/coworker/providers/vertex_provider.py b/coworker/providers/vertex_provider.py new file mode 100644 index 00000000..3110dd64 --- /dev/null +++ b/coworker/providers/vertex_provider.py @@ -0,0 +1,198 @@ +"""Google Vertex AI provider — one entry in Settings, three wire paths by model family. + +Routed ids look like `vertex:/`; 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 `/` 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) diff --git a/tests/test_vertex_provider.py b/tests/test_vertex_provider.py new file mode 100644 index 00000000..b1f52b95 --- /dev/null +++ b/tests/test_vertex_provider.py @@ -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"]