beyhan/graphify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make Gemini extraction model configurable

Browse Source 
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
This commit is contained in:
Daniel Graham
2026-05-05 10:11:12 -04:00
parent a9cb692961
commit cc63a1711b
5 changed files with 47 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -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)
graphify/__main__.py
+7
View File
@@ -1096,6 +1096,7 @@ def main() -> None:
print(" --label NAME project label in header")
print(" extract <path> 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: <path>); writes <DIR>/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:
graphify/llm.py
+21 -5
View File
@@ -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),
)
graphify/skill.md
+1 -1
View File
@@ -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.**
tests/test_llm_backends.py
+17 -2
View File
@@ -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):