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):