From a9cb6929613e81adc72a7a84ac4404662dd760e2 Mon Sep 17 00:00:00 2001 From: Daniel Graham Date: Tue, 5 May 2026 08:59:37 -0400 Subject: [PATCH 1/2] Prefer accessible semantic extraction backends Gemini is often the cheaper available quota for low-stakes semantic graph extraction, while OpenAI is a useful fallback. Extend the direct extraction backend registry, CLI validation, docs, and tests so headless extraction can use GEMINI_API_KEY, GOOGLE_API_KEY, or OPENAI_API_KEY without changing the existing Claude and Kimi paths. Constraint: Gemini supports OpenAI-compatible chat completions at the Google generative-language endpoint Rejected: Native google-genai integration | higher dependency and response-shape churn for the same chat-completions path Confidence: medium Scope-risk: moderate Directive: Keep backend detection explicit and test every accepted API-key environment variable before adding new providers Tested: uv run --directory vendor/graphify pytest tests/test_llm_backends.py tests/test_chunking.py -q Not-tested: Live Gemini/OpenAI API calls; no GEMINI_API_KEY or OPENAI_API_KEY present in this environment --- README.md | 4 +- graphify/__main__.py | 25 +++++++----- graphify/build.py | 2 +- graphify/dedup.py | 12 +++--- graphify/llm.py | 70 +++++++++++++++++++++++++++----- graphify/skill.md | 6 +-- pyproject.toml | 2 + tests/test_llm_backends.py | 82 ++++++++++++++++++++++++++++++++++++++ 8 files changed, 172 insertions(+), 31 deletions(-) create mode 100644 tests/test_llm_backends.py diff --git a/README.md b/README.md index f5d9ef74..a59c88b9 100644 --- a/README.md +++ b/README.md @@ -221,7 +221,7 @@ The MCP server gives your assistant structured access: `query_graph`, `get_node` - **Code files** — processed locally via tree-sitter. Nothing leaves your machine. - **Video / audio** — transcribed locally with faster-whisper. Nothing leaves your machine. -- **Docs, PDFs, images** — sent to your AI assistant for semantic extraction (via the `/graphify` skill, using whatever model your IDE session runs). Headless `graphify extract` requires `ANTHROPIC_API_KEY` (Claude) or `MOONSHOT_API_KEY` (Kimi). The `--dedup-llm` flag uses the same key. +- **Docs, PDFs, images** — sent to your AI assistant for semantic extraction (via the `/graphify` skill, using whatever model your IDE session runs). Headless `graphify extract` requires `GEMINI_API_KEY` / `GOOGLE_API_KEY` (Gemini), `MOONSHOT_API_KEY` (Kimi), `ANTHROPIC_API_KEY` (Claude), or `OPENAI_API_KEY` (OpenAI). The `--dedup-llm` flag uses the same key. - No telemetry, no usage tracking, no analytics. --- @@ -274,7 +274,7 @@ graphify kiro install / uninstall graphify antigravity install / uninstall graphify extract ./docs # headless LLM extraction for CI (no IDE needed) -graphify extract ./docs --backend claude # explicit backend: claude (ANTHROPIC_API_KEY) or kimi (MOONSHOT_API_KEY) +graphify extract ./docs --backend gemini # explicit backend: gemini, kimi, claude, or openai 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 7db76524..25811bd6 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1095,7 +1095,7 @@ def main() -> None: print(" --top-k-edges N per-symbol outbound edges in inspector (default 12)") print(" --label NAME project label in header") print(" extract headless full extraction (AST + semantic LLM) for CI/scripts") - print(" --backend B kimi|claude (default: whichever API key is set)") + print(" --backend B gemini|kimi|claude|openai (default: whichever API key is set)") 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") @@ -1573,8 +1573,13 @@ def main() -> None: ok = _rebuild_code(watch_path, force=force) if ok: print("Code graph updated. For doc/paper/image changes run /graphify --update in your AI assistant.") - if not os.environ.get("MOONSHOT_API_KEY") and not os.environ.get("GRAPHIFY_NO_TIPS"): - print("Tip: set MOONSHOT_API_KEY to use Kimi K2.6 for semantic extraction — 3x cheaper, richer graphs. pip install 'graphifyy[kimi]'") + if not ( + os.environ.get("GEMINI_API_KEY") + or os.environ.get("GOOGLE_API_KEY") + or os.environ.get("MOONSHOT_API_KEY") + or os.environ.get("GRAPHIFY_NO_TIPS") + ): + print("Tip: set GEMINI_API_KEY or GOOGLE_API_KEY to use Gemini for semantic extraction.") else: print("Nothing to update or rebuild failed — check output above.", file=sys.stderr) sys.exit(1) @@ -1896,7 +1901,7 @@ def main() -> None: # has an API key set. if len(sys.argv) < 3: print( - "Usage: graphify extract [--backend kimi|claude] " + "Usage: graphify extract [--backend gemini|kimi|claude|openai] " "[--out DIR] [--no-cluster]", file=sys.stderr, ) @@ -1939,13 +1944,16 @@ def main() -> None: detect_backend as _detect_backend, estimate_cost as _estimate_cost, extract_corpus_parallel as _extract_corpus_parallel, + _format_backend_env_keys, + _get_backend_api_key, ) if backend is None: backend = _detect_backend() if backend is None: print( - "error: no LLM API key found. Set MOONSHOT_API_KEY (kimi) " - "or ANTHROPIC_API_KEY (claude), or pass --backend.", + "error: no LLM API key found. Set GEMINI_API_KEY or GOOGLE_API_KEY " + "(gemini), MOONSHOT_API_KEY (kimi), ANTHROPIC_API_KEY (claude), " + "or OPENAI_API_KEY (openai), or pass --backend.", file=sys.stderr, ) sys.exit(1) @@ -1956,10 +1964,9 @@ def main() -> None: file=sys.stderr, ) sys.exit(1) - env_key = _BACKENDS[backend]["env_key"] - if not os.environ.get(env_key): + if not _get_backend_api_key(backend): print( - f"error: backend '{backend}' requires {env_key} to be set.", + f"error: backend '{backend}' requires {_format_backend_env_keys(backend)} to be set.", file=sys.stderr, ) sys.exit(1) diff --git a/graphify/build.py b/graphify/build.py index 76a83578..d0ad04d2 100644 --- a/graphify/build.py +++ b/graphify/build.py @@ -128,7 +128,7 @@ def build( directed=True produces a DiGraph that preserves edge direction (source→target). directed=False (default) produces an undirected Graph for backward compatibility. dedup=True (default) runs entity deduplication before building the graph. - dedup_llm_backend: if set (e.g. "claude" or "kimi"), uses LLM to resolve + dedup_llm_backend: if set (e.g. "gemini", "claude", or "kimi"), uses LLM to resolve ambiguous pairs in the 75–92 Jaro-Winkler score zone. Extractions are merged in order. For nodes with the same ID, the last diff --git a/graphify/dedup.py b/graphify/dedup.py index af5cef81..6efe4e2f 100644 --- a/graphify/dedup.py +++ b/graphify/dedup.py @@ -255,11 +255,13 @@ def _llm_tiebreak( ) -> None: """Batch-resolve ambiguous pairs (score in [low, high)) via LLM.""" try: - from graphify.llm import BACKENDS - import os - env_key = BACKENDS.get(backend, {}).get("env_key", "") - if not os.environ.get(env_key): - print(f"[graphify] --dedup-llm: {env_key} not set, skipping LLM tiebreaker.", flush=True) + from graphify.llm import BACKENDS, _format_backend_env_keys, _get_backend_api_key + if backend not in BACKENDS: + print(f"[graphify] --dedup-llm: unknown backend {backend!r}, skipping LLM tiebreaker.", flush=True) + return + if not _get_backend_api_key(backend): + env_keys = _format_backend_env_keys(backend) + print(f"[graphify] --dedup-llm: {env_keys} not set, skipping LLM tiebreaker.", flush=True) return except ImportError: return diff --git a/graphify/llm.py b/graphify/llm.py index 9cba0ff8..09ac2866 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1,5 +1,6 @@ -# Direct LLM backend for semantic extraction — supports Claude and Kimi K2.6. -# Used by `graphify . --backend kimi` and the benchmark scripts. +# Direct LLM backend for semantic extraction — supports Claude, Kimi K2.6, +# Gemini, and OpenAI. +# Used by `graphify extract . --backend gemini` and the benchmark scripts. # The default graphify pipeline uses Claude Code subagents via skill.md; # this module provides a direct API path for non-Claude-Code environments. from __future__ import annotations @@ -58,6 +59,21 @@ BACKENDS: dict[str, dict] = { "pricing": {"input": 0.74, "output": 4.66}, # USD per 1M tokens "temperature": None, # kimi-k2.6 enforces its own fixed temperature; sending any value raises 400 }, + "gemini": { + "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/", + "default_model": "gemini-2.5-flash", + "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"], + "pricing": {"input": 0.30, "output": 2.50}, # USD per 1M tokens + "temperature": 0, + "reasoning_effort": "none", + }, + "openai": { + "base_url": "https://api.openai.com/v1", + "default_model": "gpt-4.1-mini", + "env_key": "OPENAI_API_KEY", + "pricing": {"input": 0.40, "output": 1.60}, # USD per 1M tokens + "temperature": 0, + }, } _EXTRACTION_SYSTEM = """\ @@ -107,19 +123,43 @@ def _parse_llm_json(raw: str) -> dict: return {"nodes": [], "edges": [], "hyperedges": []} +def _backend_env_keys(backend: str) -> list[str]: + """Return accepted API-key environment variables for a backend.""" + cfg = BACKENDS[backend] + keys = cfg.get("env_keys") + if keys: + return list(keys) + return [cfg["env_key"]] + + +def _get_backend_api_key(backend: str) -> str: + """Return the first configured API key for backend, or an empty string.""" + for env_key in _backend_env_keys(backend): + value = os.environ.get(env_key) + if value: + return value + return "" + + +def _format_backend_env_keys(backend: str) -> str: + """Return user-facing accepted API-key variable names.""" + return " or ".join(_backend_env_keys(backend)) + + def _call_openai_compat( base_url: str, api_key: str, model: str, user_message: str, temperature: float | None = 0, + reasoning_effort: str | None = None, ) -> dict: """Call any OpenAI-compatible API (Kimi, OpenAI, etc.) and return parsed JSON.""" try: from openai import OpenAI except ImportError as exc: raise ImportError( - "Kimi/OpenAI-compatible extraction requires the openai package. " + "Gemini/Kimi/OpenAI-compatible extraction requires the openai package. " "Run: pip install openai" ) from exc @@ -134,6 +174,8 @@ def _call_openai_compat( } if temperature is not None: kwargs["temperature"] = temperature + if reasoning_effort is not None: + kwargs["reasoning_effort"] = reasoning_effort # Kimi-k2.6 is a reasoning model — disable thinking so content isn't empty if "moonshot" in base_url: kwargs["extra_body"] = {"thinking": {"type": "disabled"}} @@ -193,11 +235,11 @@ def extract_files_direct( raise ValueError(f"Unknown backend {backend!r}. Available: {sorted(BACKENDS)}") cfg = BACKENDS[backend] - key = api_key or os.environ.get(cfg["env_key"], "") + key = api_key or _get_backend_api_key(backend) if not key: raise ValueError( f"No API key for backend '{backend}'. " - f"Set {cfg['env_key']} or pass api_key=." + f"Set {_format_backend_env_keys(backend)} or pass api_key=." ) mdl = model or cfg["default_model"] user_msg = _read_files(files, root) @@ -205,7 +247,14 @@ def extract_files_direct( if backend == "claude": return _call_claude(key, mdl, user_msg) else: - return _call_openai_compat(cfg["base_url"], key, mdl, user_msg, temperature=cfg.get("temperature", 0)) + return _call_openai_compat( + cfg["base_url"], + key, + mdl, + user_msg, + temperature=cfg.get("temperature", 0), + reasoning_effort=cfg.get("reasoning_effort"), + ) def _estimate_file_tokens(path: Path) -> int: @@ -468,11 +517,10 @@ def estimate_cost(backend: str, input_tokens: int, output_tokens: int) -> float: def detect_backend() -> str | None: """Return the name of whichever backend has an API key set, or None. - Kimi is checked first (opt-in). Falls back to Claude if ANTHROPIC_API_KEY is set. + Gemini is checked first, then Kimi, Claude, and OpenAI. Claude is the default for the skill.md subagent pipeline and is never forced here. """ - if os.environ.get("MOONSHOT_API_KEY"): - return "kimi" - if os.environ.get("ANTHROPIC_API_KEY"): - return "claude" + for backend in ("gemini", "kimi", "claude", "openai"): + if _get_backend_api_key(backend): + return backend return None diff --git a/graphify/skill.md b/graphify/skill.md index 24d40d8a..f9f85874 100644 --- a/graphify/skill.md +++ b/graphify/skill.md @@ -191,10 +191,10 @@ After transcription: This step has two parts: **structural extraction** (deterministic, free) and **semantic extraction** (LLM, costs tokens). -**Before dispatching subagents:** check whether `MOONSHOT_API_KEY` is set. If it is NOT set, print this one-liner to the user: -> Tip: set `MOONSHOT_API_KEY` to use Kimi K2.6 for semantic extraction — 3x cheaper, richer graphs (`pip install 'graphifyy[kimi]'`). +**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 `MOONSHOT_API_KEY` IS set, use `graphify.llm.extract_corpus_parallel(files, backend="kimi")` 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. **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/pyproject.toml b/pyproject.toml index 80a9166d..5e74eb37 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,8 @@ leiden = ["graspologic; python_version < '3.13'"] office = ["python-docx", "openpyxl"] video = ["faster-whisper", "yt-dlp"] kimi = ["openai", "tiktoken"] +gemini = ["openai", "tiktoken"] +openai = ["openai", "tiktoken"] sql = ["tree-sitter-sql"] all = ["mcp", "neo4j", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken", "tree-sitter-sql"] diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py new file mode 100644 index 00000000..34766cac --- /dev/null +++ b/tests/test_llm_backends.py @@ -0,0 +1,82 @@ +"""Tests for direct semantic-extraction backend selection.""" + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from graphify import llm + + +def _clear_backend_env(monkeypatch): + for env_key in ( + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + "MOONSHOT_API_KEY", + "ANTHROPIC_API_KEY", + "OPENAI_API_KEY", + ): + monkeypatch.delenv(env_key, raising=False) + + +def test_gemini_accepts_gemini_api_key(monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("GEMINI_API_KEY", "gemini-key") + + assert llm.detect_backend() == "gemini" + assert llm._get_backend_api_key("gemini") == "gemini-key" + + +def test_gemini_accepts_google_api_key(monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("GOOGLE_API_KEY", "google-key") + + assert llm.detect_backend() == "gemini" + assert llm._get_backend_api_key("gemini") == "google-key" + + +def test_backend_detection_prefers_gemini(monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "openai-key") + monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key") + monkeypatch.setenv("MOONSHOT_API_KEY", "moonshot-key") + monkeypatch.setenv("GEMINI_API_KEY", "gemini-key") + + assert llm.detect_backend() == "gemini" + + +def test_openai_backend_detected(monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "openai-key") + + assert llm.detect_backend() == "openai" + assert llm._get_backend_api_key("openai") == "openai-key" + + +def test_extract_files_direct_routes_gemini_through_openai_compat(tmp_path, monkeypatch): + _clear_backend_env(monkeypatch) + monkeypatch.setenv("GOOGLE_API_KEY", "google-key") + source = tmp_path / "note.md" + source.write_text("# Architecture\n\nThe runner emits a snapshot.\n") + result = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 1, "output_tokens": 1} + + with patch("graphify.llm._call_openai_compat", return_value=result) as call: + assert llm.extract_files_direct([source], backend="gemini", root=tmp_path) is result + + assert call.call_args.args[:4] == ( + "https://generativelanguage.googleapis.com/v1beta/openai/", + "google-key", + "gemini-2.5-flash", + "=== 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" + + +def test_missing_gemini_key_names_both_supported_env_vars(monkeypatch): + _clear_backend_env(monkeypatch) + + with pytest.raises(ValueError) as exc: + llm.extract_files_direct([Path("missing.md")], backend="gemini") + + assert "GEMINI_API_KEY or GOOGLE_API_KEY" in str(exc.value) From cc63a1711b179d1af6df744e33eb1d62f744e347 Mon Sep 17 00:00:00 2001 From: Daniel Graham Date: Tue, 5 May 2026 10:11:12 -0400 Subject: [PATCH 2/2] 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 a59c88b9..3e6e7e60 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 25811bd6..93373148 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 09ac2866..2d9b5489 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 f9f85874..9d238e6a 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 34766cac..78121aca 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):