Merge PR #735: Add Gemini and OpenAI semantic extraction backends (preserve Ollama priority)

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), `MOONSHOT_API_KEY` (Kimi), or a running Ollama instance (`OLLAMA_BASE_URL`). 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), `OPENAI_API_KEY` (OpenAI), or a running Ollama instance (`OLLAMA_BASE_URL`). The `--dedup-llm` flag uses the same key.
 - No telemetry, no usage tracking, no analytics.
 
 ---
@@ -274,7 +274,8 @@ 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, openai, or ollama
+graphify extract ./docs --backend gemini --model gemini-3.1-pro-preview
 graphify extract ./docs --backend ollama       # local Ollama (set OLLAMA_BASE_URL / OLLAMA_MODEL)
 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

@@ -1096,8 +1096,8 @@ 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 <path>          headless full extraction (AST + semantic LLM) for CI/scripts")
-        print("    --backend B             kimi|claude (default: whichever API key is set)")
-        print("    --model <name>          override the backend's default model")
+        print("    --backend B             gemini|kimi|claude|openai|ollama (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("    --global                also merge the resulting graph into the global graph")
@@ -1609,8 +1609,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)
@@ -2000,7 +2005,7 @@ def main() -> None:
         # has an API key set.
         if len(sys.argv) < 3:
             print(
-                "Usage: graphify extract <path> [--backend kimi|claude] "
+                "Usage: graphify extract <path> [--backend gemini|kimi|claude|openai] "
                 "[--out DIR] [--no-cluster]",
                 file=sys.stderr,
             )
@@ -2012,10 +2017,10 @@ 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
-        model_override: str | None = None
         global_merge = False
         global_repo_tag: str | None = None
         args = sys.argv[3:]
@@ -2027,9 +2032,9 @@ def main() -> None:
             elif a.startswith("--backend="):
                 backend = a.split("=", 1)[1]; i += 1
             elif a == "--model" and i + 1 < len(args):
-                model_override = args[i + 1]; i += 2
+                model = args[i + 1]; i += 2
             elif a.startswith("--model="):
-                model_override = a.split("=", 1)[1]; i += 1
+                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="):
@@ -2054,13 +2059,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)
@@ -2071,10 +2079,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)
@@ -2175,7 +2182,7 @@ def main() -> None:
                     fresh = _extract_corpus_parallel(
                         [Path(p) for p in uncached_paths],
                         backend=backend,
-                        model=model_override,
+                        model=model,
                         root=target,
                     )
                 except ImportError as exc:

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

graphify/dedup.py

@@ -265,11 +265,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

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
@@ -68,6 +69,24 @@ BACKENDS: dict[str, dict] = {
         "temperature": 0,
         "max_tokens": 16384,
     },
+    "gemini": {
+        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
+        "default_model": "gemini-3-flash-preview",
+        "env_keys": ["GEMINI_API_KEY", "GOOGLE_API_KEY"],
+        "model_env_key": "GRAPHIFY_GEMINI_MODEL",
+        "pricing": {"input": 0.50, "output": 3.00},  # USD per 1M tokens
+        "temperature": 0,
+        "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,
+    },
 }
 
 
@@ -130,13 +149,48 @@ 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 _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,
     model: str,
     user_message: str,
     temperature: float | None = 0,
-    max_tokens: int = 8192,
+    reasoning_effort: str | None = None,
+    max_completion_tokens: int = 8192,
     *,
     backend: str = "",
 ) -> dict:
@@ -146,7 +200,7 @@ def _call_openai_compat(
     except ImportError as exc:
         pkg_hint = "graphifyy[kimi]" if backend == "kimi" else "openai"
         raise ImportError(
-            f"{'Ollama' if backend == 'ollama' else 'Kimi'}/OpenAI-compatible extraction requires the openai package. "
+            "Gemini/Kimi/Ollama/OpenAI-compatible extraction requires the openai package. "
             f"Run: pip install {pkg_hint}"
         ) from exc
 
@@ -157,10 +211,12 @@ def _call_openai_compat(
             {"role": "system", "content": _EXTRACTION_SYSTEM},
             {"role": "user", "content": user_message},
         ],
-        "max_completion_tokens": max_tokens,
+        "max_completion_tokens": max_completion_tokens,
     }
     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"}}
@@ -228,22 +284,31 @@ 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 and backend == "ollama":
         key = "ollama"  # Ollama ignores auth but openai client requires non-empty
     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"]
+    mdl = model or _default_model_for_backend(backend)
     user_msg = _read_files(files, root)
     max_out = _resolve_max_tokens(cfg.get("max_tokens", 8192))
 
     if backend == "claude":
         return _call_claude(key, mdl, user_msg, max_tokens=max_out)
     else:
-        return _call_openai_compat(cfg["base_url"], key, mdl, user_msg, temperature=cfg.get("temperature", 0), max_tokens=max_out, backend=backend)
+        return _call_openai_compat(
+            cfg["base_url"],
+            key,
+            mdl,
+            user_msg,
+            temperature=cfg.get("temperature", 0),
+            reasoning_effort=cfg.get("reasoning_effort"),
+            max_completion_tokens=cfg.get("max_completion_tokens", max_out),
+            backend=backend,
+        )
 
 
 def _estimate_file_tokens(path: Path) -> int:
@@ -506,13 +571,16 @@ 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.
 
-    Priority: kimi → ollama (if OLLAMA_BASE_URL set) → claude.
-    Ollama is opt-in via env var — never auto-probed.
+    Priority: gemini → kimi → ollama (opt-in via OLLAMA_BASE_URL) → claude → openai.
+    Ollama is opt-in via env var — never auto-probed without OLLAMA_BASE_URL set.
+    Claude is the default for the skill.md subagent pipeline and is never forced here.
     """
-    if os.environ.get("MOONSHOT_API_KEY"):
-        return "kimi"
+    for backend in ("gemini", "kimi"):
+        if _get_backend_api_key(backend):
+            return backend
     if os.environ.get("OLLAMA_BASE_URL"):
         return "ollama"
-    if os.environ.get("ANTHROPIC_API_KEY"):
-        return "claude"
+    for backend in ("claude", "openai"):
+        if _get_backend_api_key(backend):
+            return backend
     return None

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. 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.**
 

pyproject.toml

@@ -56,6 +56,8 @@ office = ["python-docx", "openpyxl"]
 video = ["faster-whisper", "yt-dlp"]
 kimi = ["openai", "tiktoken"]
 ollama = ["openai"]
+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"]
 

tests/test_llm_backends.py

@@ -0,0 +1,97 @@
+"""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-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"] == "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):
+    _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)