From 29ace33c27dbdec84f26d12d73181880871afa17 Mon Sep 17 00:00:00 2001 From: Elior Erez Date: Tue, 25 Aug 2026 09:58:16 -0400 Subject: [PATCH 1/4] Add Vertex AI backend for orgs that disallow raw API keys The existing "gemini" backend only talks to the API-key-only Generative Language API, which some GCP orgs disallow outright via policy. Adds a "vertex" backend using google-genai's Vertex AI mode, authenticated via Application Default Credentials (service account, Workload Identity Federation, or `gcloud auth application-default login`) instead of a static key -- no env_key/env_keys, following the same keyless pattern already used for the "bedrock" backend's AWS credential chain. Configured via GOOGLE_CLOUD_PROJECT (required) and GOOGLE_CLOUD_LOCATION (defaults to us-central1), with GRAPHIFY_VERTEX_MODEL to override the default gemini-2.5-flash. Thinking is disabled by default (thinking_budget=0) -- confirmed live that a small max_output_tokens can otherwise be silently consumed entirely by (billed, never-returned) thinking tokens before any extraction JSON is emitted, and thinking traces carry no value for a fixed-schema extraction task anyway. Verified end-to-end against a real GCP project: auto-detection via GOOGLE_CLOUD_PROJECT, the plain-text call path, and the full JSON extraction path (including token accounting) all confirmed working. --- graphify/llm.py | 146 ++++++++++++++++++++++++++++++++++++++++++++++-- pyproject.toml | 6 ++ 2 files changed, 146 insertions(+), 6 deletions(-) diff --git a/graphify/llm.py b/graphify/llm.py index ae2119773..7d6ce38ac 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -146,6 +146,31 @@ def _resolve_ollama_base_url(default: str) -> str: "max_completion_tokens": 16384, "vision": True, }, + "vertex": { + # Vertex AI via Application Default Credentials (service account, + # Workload Identity Federation, or `gcloud auth application-default + # login`) -- a different Google product from "gemini" above, which + # only talks to the API-key-only Generative Language API and cannot + # authenticate via ADC at all. Needed for orgs whose policy disallows + # raw API keys outright (confirmed live: Red Hat's GCP org policy). + # No env_key/env_keys -- like "bedrock", auth comes from the + # environment's credential chain, not a key graphify reads itself. + "default_model": "gemini-2.5-flash", + "model_env_key": "GRAPHIFY_VERTEX_MODEL", + "pricing": {"input": 0.50, "output": 3.00}, # USD per 1M tokens (gemini-2.5-flash) + "temperature": 0, + # Thinking tokens are billed but never appear in the extraction JSON, + # and can silently consume the entire max_output_tokens budget before + # any real output is emitted (confirmed live: a small max_output_tokens + # returned an empty response with the whole budget spent on thinking). + # Zeroed out by default for the same reason "gemini" above defaults to + # reasoning_effort="low" -- same model family, same rationale, just a + # harder cutoff since native Vertex offers an exact budget instead of + # a qualitative effort level. + "thinking_budget": 0, + "max_tokens": 8192, + "vision": True, + }, "openai": { # OPENAI_BASE_URL points the backend at any OpenAI-compatible server # (llama.cpp, vLLM, LM Studio, ...); OPENAI_MODEL overrides the default @@ -971,6 +996,19 @@ def _bedrock_content(user_message: str, refs: list[_ImageRef]) -> list[dict]: return content +def _vertex_content(user_message: str, refs: list[_ImageRef]): + """Build the google-genai `contents` value (list of Parts) for Vertex AI.""" + from google.genai import types + + parts = [ + types.Part.from_bytes(data=r.raw, mime_type=r.media_type) + for r in refs + if r.raw + ] + parts.append(types.Part.from_text(text=_with_image_notes(user_message, refs))) + return parts + + _LLM_JSON_MAX_BYTES = 10 * 1024 * 1024 # 10 MB hard cap before json.loads (F-016) @@ -1313,7 +1351,14 @@ def _get_backend_api_key(backend: str) -> str: def _format_backend_env_keys(backend: str) -> str: """Return user-facing accepted API-key variable names.""" keys = _backend_env_keys(backend) - return " or ".join(keys) if keys else "AWS_PROFILE or AWS_REGION" + if keys: + return " or ".join(keys) + # Keyless backends authenticate via an ambient credential chain rather + # than a graphify-read env var, so there's no key list to print — name + # the actual credential signal instead. + if backend == "vertex": + return "GOOGLE_CLOUD_PROJECT (with Application Default Credentials configured)" + return "AWS_PROFILE or AWS_REGION" def _default_model_for_backend(backend: str) -> str: @@ -1882,6 +1927,68 @@ def _call_bedrock(model: str, user_message: str, max_tokens: int = 8192, *, deep return result +def _vertex_client(): + """Construct a Vertex-AI-mode google-genai client using ADC. + + Project/location come from GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION — + the same names google-genai's own Vertex mode and `gcloud` already use, so + a working `gcloud` environment (service account, Workload Identity + Federation, or a user's `gcloud auth application-default login`) needs no + graphify-specific config beyond selecting `--backend vertex`. + """ + try: + from google import genai + except ImportError as exc: + raise ImportError(_backend_pkg_hint("google-genai", "vertex")) from exc + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "").strip() + if not project: + raise ValueError( + "Vertex AI backend requires GOOGLE_CLOUD_PROJECT to be set " + "(the GCP project ID Vertex AI calls are billed/scoped to)." + ) + location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1").strip() + return genai.Client(vertexai=True, project=project, location=location) + + +def _call_vertex(model: str, user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False, images: list[_ImageRef] | None = None) -> dict: + """Call Vertex AI's Gemini models via google-genai, authenticated by ADC.""" + from google.genai import errors as genai_errors, types + + client = _vertex_client() + config = types.GenerateContentConfig( + system_instruction=_extraction_system(deep=deep_mode), + temperature=BACKENDS["vertex"].get("temperature", 0), + max_output_tokens=max_tokens, + thinking_config=types.ThinkingConfig(thinking_budget=BACKENDS["vertex"].get("thinking_budget", 0)), + ) + try: + resp = client.models.generate_content( + model=model, + contents=_vertex_content(user_message, images or []), + config=config, + ) + except genai_errors.APIError as exc: + raise RuntimeError(f"Vertex AI API error ({exc.code}): {exc.message}") from exc + + raw_content = resp.text + result = _parse_llm_json(raw_content or "{}") + usage = resp.usage_metadata + result["input_tokens"] = (usage.prompt_token_count if usage else 0) or 0 + result["output_tokens"] = (usage.candidates_token_count if usage else 0) or 0 + result["model"] = model + finish_reason = resp.candidates[0].finish_reason if resp.candidates else None + result["finish_reason"] = "length" if finish_reason and finish_reason.name == "MAX_TOKENS" else "stop" + if _response_is_hollow(raw_content, result) and result["finish_reason"] != "length": + print( + "[graphify] vertex returned a hollow response; treating as " + "truncation so adaptive retry can bisect the chunk.", + file=sys.stderr, + ) + result["finish_reason"] = "length" + return result + + def extract_files_direct( files: list[Path], backend: str | None = None, @@ -1911,7 +2018,8 @@ def extract_files_direct( "No LLM backend configured. Set one of: GEMINI_API_KEY, ANTHROPIC_API_KEY, " "OPENAI_API_KEY, DEEPSEEK_API_KEY, MOONSHOT_API_KEY, " "AZURE_OPENAI_API_KEY+AZURE_OPENAI_ENDPOINT, OLLAMA_BASE_URL, " - "or AWS credentials. Pass backend= explicitly to select a provider." + "AWS credentials, or GOOGLE_CLOUD_PROJECT+Application Default " + "Credentials (Vertex AI). Pass backend= explicitly to select a provider." ) if backend not in BACKENDS: raise ValueError(f"Unknown backend {backend!r}. Available: {sorted(BACKENDS)}") @@ -1931,7 +2039,7 @@ def extract_files_direct( file=sys.stderr, ) key = "ollama" - if not key and backend not in ("bedrock", "claude-cli"): + if not key and backend not in ("bedrock", "vertex", "claude-cli"): raise ValueError( f"No API key for backend '{backend}'. " f"Set {_format_backend_env_keys(backend)} or pass api_key=." @@ -1957,6 +2065,8 @@ def extract_files_direct( result = _call_claude_cli(user_msg, max_tokens=max_out, deep_mode=deep_mode, images=image_refs) elif backend == "bedrock": result = _call_bedrock(mdl, user_msg, max_tokens=max_out, deep_mode=deep_mode, images=image_refs) + elif backend == "vertex": + result = _call_vertex(mdl, user_msg, max_tokens=max_out, deep_mode=deep_mode, images=image_refs) elif backend == "azure": endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT", "").strip() if not endpoint: @@ -2855,7 +2965,7 @@ def _call_llm( ollama_url = _resolve_ollama_base_url(cfg.get("base_url", "")) _validate_ollama_base_url(ollama_url) key = "ollama" - if not key and backend not in ("bedrock", "claude-cli"): + if not key and backend not in ("bedrock", "vertex", "claude-cli"): raise ValueError( f"No API key for backend '{backend}'. Set {_format_backend_env_keys(backend)}." ) @@ -2957,6 +3067,23 @@ def _rec(inp, out) -> None: _rec(bu.get("inputTokens", 0), bu.get("outputTokens", 0)) return _bedrock_response_text(resp, default="") + if backend == "vertex": + from google.genai import types + + client = _vertex_client() + resp = client.models.generate_content( + model=mdl, + contents=prompt, + config=types.GenerateContentConfig( + max_output_tokens=max_tokens, + thinking_config=types.ThinkingConfig(thinking_budget=BACKENDS["vertex"].get("thinking_budget", 0)), + ), + ) + vu = resp.usage_metadata + if vu: + _rec(vu.prompt_token_count, vu.candidates_token_count) + return resp.text or "" + if backend == "azure": endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT", "").strip() if not endpoint: @@ -3103,7 +3230,8 @@ def _validate_ollama_base_url(url: str, *, warn: bool = True) -> None: def detect_backend() -> str | None: """Return the name of whichever backend has an API key set, or None. - Priority: gemini → kimi → claude → openai → deepseek → azure → bedrock → ollama (last, opt-in). + Priority: gemini → kimi → claude → openai → deepseek → azure → bedrock → + vertex → ollama (last, opt-in). Ollama is intentionally checked LAST so a paid API key (Anthropic/OpenAI/etc.) is never silently shadowed by an incidental OLLAMA_BASE_URL in the environment @@ -3118,6 +3246,12 @@ def detect_backend() -> str | None: return "azure" if os.environ.get("AWS_PROFILE") or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION"): return "bedrock" + # GEMINI_API_KEY/GOOGLE_API_KEY (checked first, above) wins over this when + # both happen to be set -- vertex is the fallback for environments (e.g. + # under an org policy that disallows raw API keys) where a static Gemini + # key was never an option in the first place. + if os.environ.get("GOOGLE_CLOUD_PROJECT"): + return "vertex" # Honor Ollama's own OLLAMA_HOST here too, not just OLLAMA_BASE_URL (#1940) — # otherwise a user who set the standard Ollama var but no --backend still # gets "no LLM API key found". Empty default -> falsy when neither is set, @@ -3127,7 +3261,7 @@ def detect_backend() -> str | None: _validate_ollama_base_url(ollama_url) return "ollama" for name in BACKENDS: - if name not in ("gemini", "kimi", "claude", "openai", "deepseek", "azure", "bedrock", "ollama", "claude-cli"): + if name not in ("gemini", "kimi", "claude", "openai", "deepseek", "azure", "bedrock", "vertex", "ollama", "claude-cli"): if _get_backend_api_key(name): return name return None diff --git a/pyproject.toml b/pyproject.toml index 237ed4341..9a01708a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,12 @@ ollama = ["openai"] bedrock = ["boto3"] anthropic = ["anthropic"] gemini = ["openai", "tiktoken"] +# Vertex AI (ADC/service-account/Workload-Identity-Federation auth) -- a +# separate Google product from the "gemini" extra above, which only talks to +# the API-key-only Generative Language API. google-genai's Vertex mode is the +# only backend here with no API key at all, for orgs whose policy disallows +# raw API keys outright. +vertex = ["google-genai"] openai = ["openai", "tiktoken"] chinese = ["jieba"] sql = ["tree-sitter-sql"] From 129c9fba424804644fd1441f58623f149eb45361 Mon Sep 17 00:00:00 2001 From: Elior Erez Date: Tue, 25 Aug 2026 11:02:32 -0400 Subject: [PATCH 2/4] fix(llm): Check Ollama before Vertex in detect_backend priority GOOGLE_CLOUD_PROJECT is often set globally for other GCP tools (gcloud, terraform, etc.), not specifically for graphify. An explicit Ollama configuration (OLLAMA_BASE_URL/OLLAMA_HOST) should not be shadowed by an ambient GCP environment variable. Fixes the finding: 'Vertex autodetection shadows explicit Ollama configuration' --- graphify/llm.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/graphify/llm.py b/graphify/llm.py index 7d6ce38ac..ac4c912f2 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -3231,13 +3231,14 @@ def detect_backend() -> str | None: """Return the name of whichever backend has an API key set, or None. Priority: gemini → kimi → claude → openai → deepseek → azure → bedrock → - vertex → ollama (last, opt-in). - - Ollama is intentionally checked LAST so a paid API key (Anthropic/OpenAI/etc.) - is never silently shadowed by an incidental OLLAMA_BASE_URL in the environment - — see security finding F-002/F-029. Setting OLLAMA_BASE_URL alongside a paid - key now keeps you on the paid backend; remove the paid key (or pass - --backend ollama explicitly) to route to the local model. + ollama → vertex (last). + + Ollama is checked before Vertex because OLLAMA_BASE_URL/OLLAMA_HOST is an + explicit graphify-specific configuration, while GOOGLE_CLOUD_PROJECT is often + set globally for other GCP tools (gcloud, terraform, etc.). An explicit Ollama + configuration should not be shadowed by an ambient GCP environment variable. + Both are checked after paid API keys to avoid silently shadowing a paid key + with a free/local backend (security finding F-002/F-029). """ for backend in ("gemini", "kimi", "claude", "openai", "deepseek"): if _get_backend_api_key(backend): @@ -3246,12 +3247,6 @@ def detect_backend() -> str | None: return "azure" if os.environ.get("AWS_PROFILE") or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION"): return "bedrock" - # GEMINI_API_KEY/GOOGLE_API_KEY (checked first, above) wins over this when - # both happen to be set -- vertex is the fallback for environments (e.g. - # under an org policy that disallows raw API keys) where a static Gemini - # key was never an option in the first place. - if os.environ.get("GOOGLE_CLOUD_PROJECT"): - return "vertex" # Honor Ollama's own OLLAMA_HOST here too, not just OLLAMA_BASE_URL (#1940) — # otherwise a user who set the standard Ollama var but no --backend still # gets "no LLM API key found". Empty default -> falsy when neither is set, @@ -3260,6 +3255,14 @@ def detect_backend() -> str | None: if ollama_url: _validate_ollama_base_url(ollama_url) return "ollama" + # GEMINI_API_KEY/GOOGLE_API_KEY (checked first, above) wins over this when + # both happen to be set -- vertex is the fallback for environments (e.g. + # under an org policy that disallows raw API keys) where a static Gemini + # key was never an option in the first place. Checked after Ollama because + # GOOGLE_CLOUD_PROJECT is often ambient (set for gcloud/terraform/etc.), + # while OLLAMA_BASE_URL is graphify-specific. + if os.environ.get("GOOGLE_CLOUD_PROJECT"): + return "vertex" for name in BACKENDS: if name not in ("gemini", "kimi", "claude", "openai", "deepseek", "azure", "bedrock", "vertex", "ollama", "claude-cli"): if _get_backend_api_key(name): From 1a122e1b09f2b615a79b0e85482666ca005af076 Mon Sep 17 00:00:00 2001 From: Elior Erez Date: Tue, 25 Aug 2026 14:22:39 -0400 Subject: [PATCH 3/4] NO-ISSUE: fix vertex backend GOOGLE_CLOUD_PROJECT check in cli.py Vertex backend failed at runtime with 'backend vertex requires GOOGLE_CLOUD_PROJECT to be set' even when GOOGLE_CLOUD_PROJECT was genuinely set in the environment (confirmed in osac run 32878010795). Root cause: cli.py lines 3497-3535 pre-flight check for backends without API keys has special-case exemptions for ollama (localhost), bedrock (AWS env vars), and claude-cli (binary check), but no exemption for vertex. Since _get_backend_api_key('vertex') always returns empty (vertex has no env_key in BACKENDS, uses ADC instead), allow_no_key stayed False and triggered the error regardless of GOOGLE_CLOUD_PROJECT being set. Fix: add vertex branch mirroring bedrock pattern, checking for GOOGLE_CLOUD_PROJECT env var (same check _vertex_client() does). --- graphify/cli.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/graphify/cli.py b/graphify/cli.py index 5b7339726..c36ac734d 100644 --- a/graphify/cli.py +++ b/graphify/cli.py @@ -3517,6 +3517,8 @@ def _parse_float(name: str, raw: str) -> float: or os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("AWS_ACCESS_KEY_ID") ) + elif backend == "vertex": + allow_no_key = bool(os.environ.get("GOOGLE_CLOUD_PROJECT")) elif backend == "claude-cli": import shutil as _shutil allow_no_key = _shutil.which("claude") is not None From 653aa4eddc1e5ffe524a8d928b9e6d7c3c31cdf3 Mon Sep 17 00:00:00 2001 From: Elior Erez Date: Tue, 25 Aug 2026 16:09:32 -0400 Subject: [PATCH 4/4] Configure retry/backoff for the vertex backend's Vertex AI client Every other backend in this module wires _resolve_max_retries() into its SDK client so a burst of 429s during a large parallel run gets absorbed instead of dropping the chunk -- google-genai needs this set explicitly via http_options.retry_options. Confirmed live that an unconfigured client does NOT retry 429s at all (every field of an empty HttpRetryOptions() defaults to None), which is exactly why a large first-run corpus (~1100 docs across the osac mono-repo) surfaced "Resource exhausted" as hard chunk failures instead of transient delays (9 of 16 semantic chunks failed outright, producing a smaller graph that graphify's own shrink guard correctly refused to publish). --- graphify/llm.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/graphify/llm.py b/graphify/llm.py index ac4c912f2..0cfe5731d 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1938,6 +1938,7 @@ def _vertex_client(): """ try: from google import genai + from google.genai import types except ImportError as exc: raise ImportError(_backend_pkg_hint("google-genai", "vertex")) from exc @@ -1948,7 +1949,26 @@ def _vertex_client(): "(the GCP project ID Vertex AI calls are billed/scoped to)." ) location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1").strip() - return genai.Client(vertexai=True, project=project, location=location) + # Every other backend in this module wires _resolve_max_retries() into its + # SDK client (see _azure_client, _call_bedrock) specifically so a burst of + # 429s during a large parallel run gets absorbed instead of dropping the + # chunk (see _resolve_max_retries' own docstring, #1523) -- google-genai + # needs this set explicitly via http_options.retry_options; confirmed live + # that an unconfigured client does NOT retry 429s at all (every field of an + # empty HttpRetryOptions() defaults to None, i.e. no retry behavior), + # which is exactly why a large first-run corpus (~1100 docs) surfaced + # "Resource exhausted" as hard chunk failures instead of transient delays. + return genai.Client( + vertexai=True, + project=project, + location=location, + http_options=types.HttpOptions( + retry_options=types.HttpRetryOptions( + attempts=_resolve_max_retries() + 1, + http_status_codes=[429, 500, 502, 503, 504], + ) + ), + ) def _call_vertex(model: str, user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False, images: list[_ImageRef] | None = None) -> dict: