"""Vertex AI client -- Gemini via Google Cloud with service-account auth.
Lightweight pure-httpx path: gets an access token via ``gcloud auth
print-access-token`` (works when ``GOOGLE_APPLICATION_CREDENTIALS``
points at a valid service-account JSON), then POSTs to the Vertex AI
REST API.
No dependency on ``google-cloud-aiplatform`` or ``google-genai`` --
those would add ~50 MB of transitive deps. Users who want the full
SDK can install it separately; this module gives the small-dep path.
Compatible with Vertex AI Gemini 2.5 Flash / Pro models on
``us-central1``. Configure with:
- ``GOOGLE_APPLICATION_CREDENTIALS`` -- path to service-account JSON
- ``GOOGLE_CLOUD_PROJECT`` -- project id
- ``MORIE_VERTEX_REGION`` (optional) -- default ``us-central1``
"""
from __future__ import annotations
import json
import os
import time
from dataclasses import dataclass
from typing import Any
try:
import httpx as _httpx # type: ignore
except ImportError: # pragma: no cover
_httpx = None
[docs]
@dataclass
class VertexConfig:
"""Resolved Vertex configuration from env vars."""
project: str
location: str = "us-central1"
model: str = "gemini-2.5-flash"
token_ttl_s: int = 3300 # access tokens last ~1h; refresh at 55m
gcloud_path: str = "gcloud"
[docs]
def resolve_config() -> VertexConfig:
"""Read Vertex config from environment, with sensible defaults."""
project = os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get("MORIE_EE_PROJECT")
if not project:
raise RuntimeError(
"Vertex requires GOOGLE_CLOUD_PROJECT (or MORIE_EE_PROJECT). "
"Set via scripts/envset GOOGLE_CLOUD_PROJECT=your-project-id."
)
gcloud = os.environ.get("GCLOUD_PATH") or "gcloud"
# Try known install paths if bare 'gcloud' isn't on PATH.
for candidate in (gcloud, "/mnt/nvme/google-cloud-sdk/bin/gcloud", "/opt/google-cloud-sdk/bin/gcloud"):
if os.path.exists(candidate) or _which(candidate):
gcloud = candidate
break
return VertexConfig(
project=project,
location=os.environ.get("VERTEX_LOCATION", "us-central1"),
model=os.environ.get("VERTEX_MODEL", "gemini-2.5-flash"),
gcloud_path=gcloud,
)
def _which(cmd: str) -> str | None:
from shutil import which
return which(cmd)
_TOKEN_CACHE: dict[str, Any] = {"token": None, "expires_at": 0.0}
def _access_token(cfg: VertexConfig) -> str:
"""Get a valid GCP access token, caching until close to TTL expiry."""
now = time.time()
if _TOKEN_CACHE["token"] and now < _TOKEN_CACHE["expires_at"]:
return _TOKEN_CACHE["token"]
from ._interactive import launcher
sp = launcher("Getting a Google Cloud access token (gcloud)")
try:
out = sp.run(
[cfg.gcloud_path, "auth", "print-access-token"],
check=True,
capture_output=True,
text=True,
encoding="utf-8",
errors="replace",
timeout=30,
).stdout.strip()
except sp.CalledProcessError as exc:
raise RuntimeError(f"gcloud print-access-token failed: {exc.stderr.strip()}") from exc
except FileNotFoundError as exc:
raise RuntimeError(f"gcloud not found at {cfg.gcloud_path}. Set GCLOUD_PATH env var.") from exc
if not out:
raise RuntimeError("gcloud returned empty access token.")
_TOKEN_CACHE["token"] = out
_TOKEN_CACHE["expires_at"] = now + cfg.token_ttl_s
return out
[docs]
def ask_gemini(
prompt: str,
*,
model: str | None = None,
instructions: str | None = None,
temperature: float = 0.1,
max_output_tokens: int = 2048,
timeout_s: float = 120.0,
cfg: VertexConfig | None = None,
) -> str:
"""Send a single-turn prompt to Gemini via Vertex AI; return the text.
Parameters
----------
prompt : str
The user prompt.
model : str, optional
Model override. Defaults to VERTEX_MODEL env or
``gemini-2.5-flash``.
instructions : str, optional
Instructions for the model (persona, constraints), sent as the
``systemInstruction`` field of the Gemini REST API.
temperature : float, default 0.1
Deterministic for scientific work. Match morie.llm default.
max_output_tokens : int, default 2048
Gemini 2.5 Flash supports up to 8192; 2048 is a sane default.
timeout_s : float, default 120
HTTP client timeout.
cfg : VertexConfig, optional
Pre-resolved config. If None, ``resolve_config()`` is called.
Returns
-------
str
The generated text, trimmed.
Raises
------
ImportError
If ``httpx`` isn't installed.
RuntimeError
If gcloud auth, project env, or Vertex call fails.
"""
if _httpx is None:
raise ImportError("morie.vertex requires httpx. Install with: pip install httpx")
cfg = cfg or resolve_config()
model = model or cfg.model
token = _access_token(cfg)
endpoint = (
f"https://{cfg.location}-aiplatform.googleapis.com/v1/"
f"projects/{cfg.project}/locations/{cfg.location}/"
f"publishers/google/models/{model}:generateContent"
)
payload: dict[str, Any] = {
"contents": [
{"role": "user", "parts": [{"text": prompt}]},
],
"generationConfig": {
"temperature": float(temperature),
"maxOutputTokens": int(max_output_tokens),
},
}
if instructions:
payload["systemInstruction"] = {
"parts": [{"text": instructions}],
}
with _httpx.Client(timeout=timeout_s) as client:
r = client.post(
endpoint,
headers={
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
},
json=payload,
)
if r.status_code != 200:
raise RuntimeError(f"Vertex API returned {r.status_code}: {r.text[:400]}")
data = r.json()
try:
parts = data["candidates"][0]["content"]["parts"]
return "".join(p.get("text", "") for p in parts).strip()
except (KeyError, IndexError) as exc:
raise RuntimeError(f"unexpected Vertex response shape: {json.dumps(data)[:400]}") from exc
[docs]
def health_check() -> dict[str, Any]:
"""Run a tiny smoke test -- returns a dict suitable for JSON logging."""
out: dict[str, Any] = {"ok": False, "error": None, "model": None}
try:
cfg = resolve_config()
out["project"] = cfg.project
out["location"] = cfg.location
out["model"] = cfg.model
reply = ask_gemini(
"reply with just OK and nothing else",
cfg=cfg,
temperature=0.0,
max_output_tokens=8,
)
out["reply"] = reply
out["ok"] = bool(reply)
except Exception as exc:
out["error"] = f"{type(exc).__name__}: {exc}"
return out