diff --git a/README.md b/README.md index b92ace14..5cb9c5ee 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), `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) diff --git a/graphify/__main__.py b/graphify/__main__.py index fb28cf42..4eae7315 100644 --- a/graphify/__main__.py +++ b/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 headless full extraction (AST + semantic LLM) for CI/scripts") - print(" --backend B kimi|claude (default: whichever API key is set)") - print(" --model 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: ); writes /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 [--backend kimi|claude] " + "Usage: graphify extract [--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: diff --git a/graphify/build.py b/graphify/build.py index 6154c16d..00d3eedf 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 8a4188e6..0b862459 100644 --- a/graphify/dedup.py +++ b/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 diff --git a/graphify/llm.py b/graphify/llm.py index 77fcbb78..1ad663d9 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 @@ -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 diff --git a/graphify/skill.md b/graphify/skill.md index 24d40d8a..9d238e6a 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. 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/pyproject.toml b/pyproject.toml index 5833ceef..00a50088 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py new file mode 100644 index 00000000..78121aca --- /dev/null +++ b/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)