diff --git a/README.md b/README.md index f5d9ef7..a59c88b 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 7db7652..25811bd 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 76a8357..d0ad04d 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 af5cef8..6efe4e2 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 9cba0ff..09ac286 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 24d40d8..f9f8587 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 80a9166..5e74eb3 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 0000000..34766ca --- /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)