From cc63a1711b179d1af6df744e33eb1d62f744e347 Mon Sep 17 00:00:00 2001 From: Daniel Graham Date: Tue, 5 May 2026 10:11:12 -0400 Subject: [PATCH] Make Gemini extraction model configurable The initial Gemini backend defaulted to 2.5 Flash, but large semantic extraction chunks can benefit from newer models and more output headroom. Move the default to Gemini 3 Flash Preview, add CLI and environment model overrides, and increase the Gemini completion budget while keeping low reasoning effort for cost control. Constraint: Google exposes Gemini through an OpenAI-compatible chat-completions endpoint Rejected: Hardcode Gemini 3.1 Pro as the default | higher cost for routine repository indexing Confidence: medium Scope-risk: narrow Directive: Keep --model and GRAPHIFY_GEMINI_MODEL working before changing Gemini defaults again Tested: uv run --directory vendor/graphify pytest tests/test_llm_backends.py tests/test_chunking.py -q Not-tested: Live Gemini 3 extraction on the full cloud-edge repo before this commit --- README.md | 1 + graphify/__main__.py | 7 +++++++ graphify/llm.py | 26 +++++++++++++++++++++----- graphify/skill.md | 2 +- tests/test_llm_backends.py | 19 +++++++++++++++++-- 5 files changed, 47 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index a59c88b..3e6e7e6 100644 --- a/README.md +++ b/README.md @@ -275,6 +275,7 @@ graphify antigravity install / uninstall graphify extract ./docs # headless LLM extraction for CI (no IDE needed) graphify extract ./docs --backend gemini # explicit backend: gemini, kimi, claude, or openai +graphify extract ./docs --backend gemini --model gemini-3.1-pro-preview graphify extract ./docs --no-cluster # raw extraction only, skip clustering graphify extract ./docs --dedup-llm # LLM tiebreaker for ambiguous entity pairs (uses same API key) diff --git a/graphify/__main__.py b/graphify/__main__.py index 25811bd..9337314 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1096,6 +1096,7 @@ def main() -> None: print(" --label NAME project label in header") print(" extract headless full extraction (AST + semantic LLM) for CI/scripts") print(" --backend B gemini|kimi|claude|openai (default: whichever API key is set)") + print(" --model M override backend default model") print(" --out DIR output dir (default: ); writes /graphify-out/") print(" --no-cluster skip clustering, write raw extraction only") print(" benchmark [graph.json] measure token reduction vs naive full-corpus approach") @@ -1913,6 +1914,7 @@ def main() -> None: sys.exit(1) backend: str | None = None + model: str | None = None out_dir: Path | None = None no_cluster = False dedup_llm = False @@ -1924,6 +1926,10 @@ def main() -> None: backend = args[i + 1]; i += 2 elif a.startswith("--backend="): backend = a.split("=", 1)[1]; i += 1 + elif a == "--model" and i + 1 < len(args): + model = args[i + 1]; i += 2 + elif a.startswith("--model="): + model = a.split("=", 1)[1]; i += 1 elif a == "--out" and i + 1 < len(args): out_dir = Path(args[i + 1]); i += 2 elif a.startswith("--out="): @@ -2067,6 +2073,7 @@ def main() -> None: fresh = _extract_corpus_parallel( [Path(p) for p in uncached_paths], backend=backend, + model=model, root=target, ) except ImportError as exc: diff --git a/graphify/llm.py b/graphify/llm.py index 09ac286..2d9b548 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -61,16 +61,19 @@ BACKENDS: dict[str, dict] = { }, "gemini": { "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/", - "default_model": "gemini-2.5-flash", + "default_model": "gemini-3-flash-preview", "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"], - "pricing": {"input": 0.30, "output": 2.50}, # USD per 1M tokens + "model_env_key": "GRAPHIFY_GEMINI_MODEL", + "pricing": {"input": 0.50, "output": 3.00}, # USD per 1M tokens "temperature": 0, - "reasoning_effort": "none", + "reasoning_effort": "low", + "max_completion_tokens": 16384, }, "openai": { "base_url": "https://api.openai.com/v1", "default_model": "gpt-4.1-mini", "env_key": "OPENAI_API_KEY", + "model_env_key": "GRAPHIFY_OPENAI_MODEL", "pricing": {"input": 0.40, "output": 1.60}, # USD per 1M tokens "temperature": 0, }, @@ -146,6 +149,17 @@ def _format_backend_env_keys(backend: str) -> str: return " or ".join(_backend_env_keys(backend)) +def _default_model_for_backend(backend: str) -> str: + """Return configured model override or backend default model.""" + cfg = BACKENDS[backend] + model_env_key = cfg.get("model_env_key") + if model_env_key: + model = os.environ.get(model_env_key) + if model: + return model + return cfg["default_model"] + + def _call_openai_compat( base_url: str, api_key: str, @@ -153,6 +167,7 @@ def _call_openai_compat( user_message: str, temperature: float | None = 0, reasoning_effort: str | None = None, + max_completion_tokens: int = 8192, ) -> dict: """Call any OpenAI-compatible API (Kimi, OpenAI, etc.) and return parsed JSON.""" try: @@ -170,7 +185,7 @@ def _call_openai_compat( {"role": "system", "content": _EXTRACTION_SYSTEM}, {"role": "user", "content": user_message}, ], - "max_completion_tokens": 8192, + "max_completion_tokens": max_completion_tokens, } if temperature is not None: kwargs["temperature"] = temperature @@ -241,7 +256,7 @@ def extract_files_direct( f"No API key for backend '{backend}'. " f"Set {_format_backend_env_keys(backend)} or pass api_key=." ) - mdl = model or cfg["default_model"] + mdl = model or _default_model_for_backend(backend) user_msg = _read_files(files, root) if backend == "claude": @@ -254,6 +269,7 @@ def extract_files_direct( user_msg, temperature=cfg.get("temperature", 0), reasoning_effort=cfg.get("reasoning_effort"), + max_completion_tokens=cfg.get("max_completion_tokens", 8192), ) diff --git a/graphify/skill.md b/graphify/skill.md index f9f8587..9d238e6 100644 --- a/graphify/skill.md +++ b/graphify/skill.md @@ -194,7 +194,7 @@ This step has two parts: **structural extraction** (deterministic, free) and **s **Before dispatching subagents:** check whether `GEMINI_API_KEY` or `GOOGLE_API_KEY` is set. If neither is set, print this one-liner to the user: > Tip: set `GEMINI_API_KEY` or `GOOGLE_API_KEY` to use Gemini for semantic extraction (`pip install 'graphifyy[gemini]'`). -Print it once, then continue. If `GEMINI_API_KEY` or `GOOGLE_API_KEY` IS set, use `graphify.llm.extract_corpus_parallel(files, backend="gemini")` for semantic extraction instead of dispatching Claude subagents. +Print it once, then continue. If `GEMINI_API_KEY` or `GOOGLE_API_KEY` IS set, use `graphify.llm.extract_corpus_parallel(files, backend="gemini")` for semantic extraction instead of dispatching Claude subagents. The default Gemini model is `gemini-3-flash-preview`; set `GRAPHIFY_GEMINI_MODEL` or pass `--model` in headless CLI flows to override it. **Run Part A (AST) and Part B (semantic) in parallel. Dispatch all semantic subagents AND start AST extraction in the same message. Both can run simultaneously since they operate on different file types. Merge results in Part C as before.** diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py index 34766ca..78121ac 100644 --- a/tests/test_llm_backends.py +++ b/tests/test_llm_backends.py @@ -66,11 +66,26 @@ def test_extract_files_direct_routes_gemini_through_openai_compat(tmp_path, monk assert call.call_args.args[:4] == ( "https://generativelanguage.googleapis.com/v1beta/openai/", "google-key", - "gemini-2.5-flash", + "gemini-3-flash-preview", "=== note.md ===\n# Architecture\n\nThe runner emits a snapshot.\n", ) assert call.call_args.kwargs["temperature"] == 0 - assert call.call_args.kwargs["reasoning_effort"] == "none" + assert call.call_args.kwargs["reasoning_effort"] == "low" + assert call.call_args.kwargs["max_completion_tokens"] == 16384 + + +def test_gemini_model_can_be_overridden_by_env(tmp_path, monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("GOOGLE_API_KEY", "google-key") + monkeypatch.setenv("GRAPHIFY_GEMINI_MODEL", "gemini-3.1-pro-preview") + source = tmp_path / "note.md" + source.write_text("# Architecture\n") + result = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 1, "output_tokens": 1} + + with patch("graphify.llm._call_openai_compat", return_value=result) as call: + llm.extract_files_direct([source], backend="gemini", root=tmp_path) + + assert call.call_args.args[2] == "gemini-3.1-pro-preview" def test_missing_gemini_key_names_both_supported_env_vars(monkeypatch):