From 33d3efd3b24ce021de7eed0eaf373c320d888ea0 Mon Sep 17 00:00:00 2001 From: Rohit C Prasad Date: Sun, 26 Jul 2026 22:26:54 -0700 Subject: [PATCH] Vertex: countTokens verify, global-location host, honest region help Model listing 403/404s under plain ADC; countTokens is free and proves project+location+API in one call. Verified live: Gemini (global), Qwen MaaS (us-south1). --- coworker/providers/registry.py | 34 +++++++++++++++++++-------- coworker/providers/vertex_provider.py | 10 +++++++- tests/test_vertex_provider.py | 23 ++++++++++++++---- 3 files changed, 51 insertions(+), 16 deletions(-) diff --git a/coworker/providers/registry.py b/coworker/providers/registry.py index 49df1995..a07ec316 100644 --- a/coworker/providers/registry.py +++ b/coworker/providers/registry.py @@ -400,9 +400,10 @@ DESCRIPTORS: list[ProviderDescriptor] = [ "location", "Location", secret=False, - placeholder="us-east5", - help="The region your Vertex AI models are enabled in " - "(Claude models: us-east5 or europe-west1).", + placeholder="global", + help="Use `global` for the newest Gemini and Claude models. Some models " + "are regional — Model Garden lists each (Claude also: us-east5 / " + "europe-west1; Qwen3 Coder: us-south1).", ), ProviderField( "auth_method", @@ -695,9 +696,16 @@ def _verify_bedrock(fields: dict[str, Any], timeout: float) -> dict[str, Any]: return {"ok": True} +# Verify probe: countTokens on a stable Gemini model — free (no generation), works with +# plain ADC (the model list/GET endpoints 403/404 under user credentials — checked live +# 2026-07-26), and exercises project + location + API enablement in one call. +_VERTEX_PROBE_MODEL = "gemini-2.5-flash" +_VERTEX_PROBE_BODY = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} + + def _verify_vertex(fields: dict[str, Any], timeout: float) -> dict[str, Any]: - """One cheap read-only call (list Google's publisher models) through the SELECTED - auth method: ADC / service-account bearer, or the express API key header.""" + """One cheap call (countTokens) through the SELECTED auth method: ADC / + service-account bearer, or the express API key header.""" import httpx from .vertex_provider import load_credentials @@ -713,9 +721,11 @@ def _verify_vertex(fields: dict[str, Any], timeout: float) -> dict[str, Any]: return {"ok": False, "error": "Enter a Vertex API key to test."} try: # Express mode is global — no region host, no project in the path. - resp = httpx.get( - "https://aiplatform.googleapis.com/v1/publishers/google/models", + resp = httpx.post( + "https://aiplatform.googleapis.com/v1/publishers/google/models/" + f"{_VERTEX_PROBE_MODEL}:countTokens", headers={"x-goog-api-key": key}, + json=_VERTEX_PROBE_BODY, timeout=timeout, ) except Exception as exc: @@ -754,11 +764,15 @@ def _verify_vertex(fields: dict[str, Any], timeout: float) -> dict[str, Any]: 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})."} + from .vertex_provider import _regional_host + try: - resp = httpx.get( - f"https://{location}-aiplatform.googleapis.com/v1/projects/{project}" - f"/locations/{location}/publishers/google/models", + resp = httpx.post( + f"https://{_regional_host(location)}/v1/projects/{project}" + f"/locations/{location}/publishers/google/models/" + f"{_VERTEX_PROBE_MODEL}:countTokens", headers={"Authorization": f"Bearer {creds.token}"}, + json=_VERTEX_PROBE_BODY, timeout=timeout, ) except Exception as exc: diff --git a/coworker/providers/vertex_provider.py b/coworker/providers/vertex_provider.py index 0891339a..66054fe1 100644 --- a/coworker/providers/vertex_provider.py +++ b/coworker/providers/vertex_provider.py @@ -46,6 +46,14 @@ _SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] _FAMILIES = ("gemini", "claude", "openweight") +def _regional_host(location: Optional[str]) -> str: + """Vertex REST host for a location — `global` (newer Gemini models) has no region + prefix (checked live 2026-07-26).""" + if not location or location == "global": + return "aiplatform.googleapis.com" + return f"{location}-aiplatform.googleapis.com" + + 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).""" @@ -191,7 +199,7 @@ class VertexProvider(ProviderClient): client = self._clients.get("openweight") if client is None: base = ( - f"https://{self._location}-aiplatform.googleapis.com/v1/projects/" + f"https://{_regional_host(self._location)}/v1/projects/" f"{self._project}/locations/{self._location}/endpoints/openapi" ) client = OpenAIProvider(api_key=creds.token, base_url=base) diff --git a/tests/test_vertex_provider.py b/tests/test_vertex_provider.py index 48055b03..89947f27 100644 --- a/tests/test_vertex_provider.py +++ b/tests/test_vertex_provider.py @@ -105,6 +105,16 @@ def test_openweight_builds_maas_endpoint_and_refreshes_bearer(): assert rebuilt._api_key == "tok-3" +def test_openweight_global_location_has_no_region_host(): + creds = _FakeCreds() + p = VertexProvider(project="proj", location="global", credentials=creds) + client = p._openweight_client() + assert client._base_url == ( + "https://aiplatform.googleapis.com/v1/projects/proj" + "/locations/global/endpoints/openapi" + ) + + # -- credentials ------------------------------------------------------------------------ @@ -240,7 +250,7 @@ def test_verify_vertex_api_key_method(monkeypatch): captured: dict = {} - def fake_get(url, headers=None, timeout=None, **kw): + def fake_post(url, headers=None, json=None, timeout=None, **kw): captured["url"] = url captured["headers"] = headers @@ -249,7 +259,7 @@ def test_verify_vertex_api_key_method(monkeypatch): return _Resp() - monkeypatch.setattr(httpx, "get", fake_get) + monkeypatch.setattr(httpx, "post", fake_post) out = verify_provider_key( "vertex", fields={ @@ -262,7 +272,10 @@ def test_verify_vertex_api_key_method(monkeypatch): assert out == {"ok": True} assert captured["headers"]["x-goog-api-key"] == "AQ.k" # Express mode is global — no region host, no project in the path. - assert captured["url"] == "https://aiplatform.googleapis.com/v1/publishers/google/models" + assert captured["url"] == ( + "https://aiplatform.googleapis.com/v1/publishers/google/models/" + "gemini-2.5-flash:countTokens" + ) def test_verify_vertex_service_account_requires_json(): @@ -306,7 +319,7 @@ def _patch_verify(monkeypatch, creds: Any, status_code: Optional[int]): monkeypatch.setattr(vp, "load_credentials", lambda raw: creds) captured: dict = {} - def fake_get(url, headers=None, timeout=None, **kw): + def fake_post(url, headers=None, json=None, timeout=None, **kw): captured["url"] = url captured["headers"] = headers @@ -317,7 +330,7 @@ def _patch_verify(monkeypatch, creds: Any, status_code: Optional[int]): resp.status_code = status_code return resp - monkeypatch.setattr(httpx, "get", fake_get) + monkeypatch.setattr(httpx, "post", fake_post) return captured