Files
graphify/graphify/llm.py
T
Safi 250834a0c4 add anthropic extra and fix backend-missing hints for uv tool installs
uv tool install graphifyy runs graphify in an isolated venv, so when a backend
needs anthropic/openai/boto3 the old 'Run: pip install anthropic' advice never
reached it. Worse, claude was the only backend with no [extra] at all, so a user
with ANTHROPIC_API_KEY set could not satisfy it without --with anthropic.

Add an anthropic optional extra (and include it in [all]), and replace the
package-missing ImportError messages with a shared _backend_pkg_hint that points
at 'uv tool install graphifyy[<extra>] --force' first, then the pip/venv path.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-03 00:08:03 +01:00

1500 lines
63 KiB
Python

# 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
import json
import os
import re
import sys
import time
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
# `_read_files` truncates each file at this many characters before joining into
# the user message. Token estimates use the same cap so packing matches reality.
_FILE_CHAR_CAP = 20_000
# `_read_files` also wraps each file in a `=== {rel} ===\n...\n\n` separator;
# this is roughly the per-file overhead in characters that the prompt adds.
_PER_FILE_OVERHEAD_CHARS = 80
# Coarse fallback used only when `tiktoken` is not installed. 1 token ≈ 4 chars
# is the standard heuristic for English/code on BPE tokenizers.
_CHARS_PER_TOKEN = 4
def _get_tokenizer():
"""Return a tiktoken encoder for accurate token counts, or None if tiktoken
is not installed. We use `cl100k_base` (GPT-4 / GPT-3.5-turbo) as a proxy:
Kimi-K2 ships a tiktoken-based tokenizer with very similar BPE behaviour,
and Claude's tokenizer has a comparable token-to-char ratio for prose/code.
Estimates only need to be within ~5%, not exact.
"""
try:
import tiktoken
except ImportError:
return None
try:
return tiktoken.get_encoding("cl100k_base")
except Exception: # network failure on first-use download, etc.
return None
# Cached at import time. None if tiktoken is unavailable; consumers must handle.
_TOKENIZER = _get_tokenizer()
BACKENDS: dict[str, dict] = {
"claude": {
"base_url": "https://api.anthropic.com",
"default_model": "claude-sonnet-4-6",
"env_key": "ANTHROPIC_API_KEY",
"pricing": {"input": 3.0, "output": 15.0}, # USD per 1M tokens
"temperature": 0,
"max_tokens": 16384,
},
"kimi": {
"base_url": "https://api.moonshot.ai/v1",
"default_model": "kimi-k2.6",
"env_key": "MOONSHOT_API_KEY",
"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
"max_tokens": 16384,
},
"ollama": {
"base_url": os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
"default_model": os.environ.get("OLLAMA_MODEL", "qwen2.5-coder:7b"),
"env_key": "OLLAMA_API_KEY",
"pricing": {"input": 0.0, "output": 0.0},
"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,
},
"deepseek": {
"base_url": "https://api.deepseek.com",
"default_model": "deepseek-v4-flash",
"env_key": "DEEPSEEK_API_KEY",
"model_env_key": "GRAPHIFY_DEEPSEEK_MODEL",
"pricing": {"input": 0.14, "output": 0.28}, # USD per 1M tokens (v4-flash)
# deepseek-reasoner / thinking-mode models silently ignore temperature;
# deepseek-chat / v4-flash (non-thinking) accept 0-2. Safe to send 0.
"temperature": 0,
"max_tokens": 16384,
},
"bedrock": {
"default_model": "anthropic.claude-3-5-sonnet-20241022-v2:0",
"model_env_key": "GRAPHIFY_BEDROCK_MODEL",
"pricing": {"input": 3.0, "output": 15.0}, # USD per 1M tokens
"temperature": 0,
"max_tokens": 16384,
},
"claude-cli": {
# Routes through the locally-installed `claude` CLI (Claude Code) using
# `-p --output-format json`. Authenticates via the user's existing
# Pro/Max subscription instead of a separate ANTHROPIC_API_KEY — costs
# are billed to the plan, not pay-as-you-go API credit.
"default_model": "claude-code-plan",
"pricing": {"input": 0.0, "output": 0.0},
"temperature": 0,
"max_tokens": 16384,
},
}
def _custom_providers_path(global_: bool = True) -> Path:
if global_:
return Path.home() / ".graphify" / "providers.json"
return Path(".graphify") / "providers.json"
def provider_base_url_ok(base_url: str, name: str, *, warn: bool = True) -> bool:
"""Structural safety check for a custom-provider base_url.
A custom provider receives the full corpus plus the user's API key, so its
base_url is an exfiltration channel. We deliberately do NOT run the ingest
SSRF guard here: that blocks private/internal IPs, which would wrongly reject
legitimate on-prem corporate LLM gateways. Instead we reject non-http(s)
schemes outright and warn loudly when the corpus would leave over plaintext
http to a non-loopback host. The primary control against trusting injected
config is the GRAPHIFY_ALLOW_LOCAL_PROVIDERS gate on project-local files.
"""
from urllib.parse import urlparse
try:
parsed = urlparse(base_url)
except Exception:
if warn:
print(f"[graphify] WARNING: provider {name!r} has an unparseable base_url; ignoring.", file=sys.stderr)
return False
if parsed.scheme not in ("http", "https"):
if warn:
print(
f"[graphify] WARNING: provider {name!r} base_url scheme {parsed.scheme!r} is not "
"http/https; ignoring.",
file=sys.stderr,
)
return False
host = (parsed.hostname or "").lower()
is_loopback = host in ("localhost", "127.0.0.1", "::1") or host.startswith("127.")
if warn and parsed.scheme == "http" and not is_loopback:
print(
f"[graphify] WARNING: provider {name!r} sends your corpus to {host!r} over plaintext "
"http. Use https unless this is a trusted local endpoint.",
file=sys.stderr,
)
return True
def _load_custom_providers() -> dict[str, dict]:
# A project-local ./.graphify/providers.json travels with a cloned or shared
# repo and defines where the corpus + API key are sent, so loading it
# silently is a corpus/key exfiltration vector. Require an explicit opt-in;
# the user's own global ~/.graphify/providers.json stays trusted.
local_path = _custom_providers_path(global_=False)
global_path = _custom_providers_path(global_=True)
allow_local = os.environ.get("GRAPHIFY_ALLOW_LOCAL_PROVIDERS", "").strip().lower() in ("1", "true", "yes")
if local_path.is_file() and not allow_local:
print(
f"[graphify] WARNING: ignoring project-local {local_path} (custom providers control "
"where your corpus and API key are sent). Set GRAPHIFY_ALLOW_LOCAL_PROVIDERS=1 to load it.",
file=sys.stderr,
)
providers: dict[str, dict] = {}
paths = [local_path, global_path] if allow_local else [global_path]
for path in paths:
if path.is_file():
try:
data = json.loads(path.read_text(encoding="utf-8"))
if isinstance(data, dict):
for name, cfg in data.items():
if not (isinstance(name, str) and isinstance(cfg, dict)):
continue
if name in BACKENDS or name in providers:
continue
if not provider_base_url_ok(str(cfg.get("base_url", "")), name):
continue
if "pricing" not in cfg:
cfg = dict(cfg, pricing={"input": 0.0, "output": 0.0})
providers[name] = cfg
except Exception:
pass
return providers
BACKENDS.update(_load_custom_providers())
def _resolve_max_tokens(default: int) -> int:
"""Honour GRAPHIFY_MAX_OUTPUT_TOKENS env var override, else use backend default."""
raw = os.environ.get("GRAPHIFY_MAX_OUTPUT_TOKENS", "").strip()
if raw:
try:
v = int(raw)
if v > 0:
return v
except ValueError:
pass
return default
_EXTRACTION_SYSTEM = """\
You are a graphify semantic extraction agent. Extract a knowledge graph fragment from the files provided.
Output ONLY valid JSON — no explanation, no markdown fences, no preamble.
Rules:
- EXTRACTED: relationship explicit in source (import, call, citation, reference)
- INFERRED: reasonable inference (shared data structure, implied dependency)
- AMBIGUOUS: uncertain — flag for review, do not omit
Node ID format: lowercase, only [a-z0-9_], no dots or slashes.
Format: {stem}_{entity} where stem = filename without extension, entity = symbol name (both normalised).
Output exactly this schema:
{"nodes":[{"id":"stem_entity","label":"Human Readable Name","file_type":"code|document|paper|image|rationale|concept","source_file":"relative/path","source_location":null,"source_url":null,"captured_at":null,"author":null,"contributor":null}],"edges":[{"source":"node_id","target":"node_id","relation":"calls|implements|references|cites|conceptually_related_to|shares_data_with|semantically_similar_to","confidence":"EXTRACTED|INFERRED|AMBIGUOUS","confidence_score":1.0,"source_file":"relative/path","source_location":null,"weight":1.0}],"hyperedges":[],"input_tokens":0,"output_tokens":0}
"""
_DEEP_EXTRACTION_SUFFIX = """\
DEEP_MODE: include additional INFERRED edges only for concrete architectural
signals (shared data contracts, explicit lifecycle coupling, or multi-step flow
dependencies visible in the sources). Avoid broad conceptual similarity edges.
Mark uncertain ones AMBIGUOUS instead of omitting.
"""
def _extraction_system(*, deep: bool = False) -> str:
"""Return the semantic-extraction system prompt, optionally in deep mode."""
if not deep:
return _EXTRACTION_SYSTEM
return _EXTRACTION_SYSTEM + _DEEP_EXTRACTION_SUFFIX
def _read_files(paths: list[Path], root: Path) -> str:
"""Return file contents formatted for the extraction prompt."""
parts: list[str] = []
for p in paths:
try:
rel = p.relative_to(root)
except ValueError:
rel = p
try:
content = p.read_text(encoding="utf-8", errors="replace")
except OSError:
continue
parts.append(f"=== {rel} ===\n{content[:20000]}")
return "\n\n".join(parts)
_LLM_JSON_MAX_BYTES = 10 * 1024 * 1024 # 10 MB hard cap before json.loads (F-016)
def _parse_llm_json(raw: str) -> dict:
"""Strip optional markdown fences and parse JSON. Returns empty fragment on failure.
Caps the input at `_LLM_JSON_MAX_BYTES` so a hostile or runaway model
response cannot exhaust memory inside `json.loads` (F-016).
"""
if len(raw) > _LLM_JSON_MAX_BYTES:
print(
f"[graphify] LLM response exceeds {_LLM_JSON_MAX_BYTES} bytes "
f"({len(raw)} bytes); refusing to parse and dropping chunk.",
file=sys.stderr,
)
return {"nodes": [], "edges": [], "hyperedges": []}
# Strategy 1: strip whitespace, then handle markdown fences anywhere in the
# text (not only at offset 0 — the original code only stripped fences when
# `raw.startswith("```")`, missing the common case where Claude prepends a
# preamble like "Here's the extracted entities:\n\n```json\n{...}\n```").
stripped = raw.strip()
fence_start = stripped.find("```")
if fence_start != -1:
after_fence = stripped[fence_start + 3 :]
# Optional language tag (json, JSON, javascript, etc.) up to newline.
nl = after_fence.find("\n")
if nl != -1 and after_fence[:nl].strip().lower() in {"json", "javascript", "js", ""}:
after_fence = after_fence[nl + 1 :]
fence_end = after_fence.rfind("```")
if fence_end != -1:
stripped = after_fence[:fence_end].strip()
else:
stripped = after_fence.strip()
try:
return json.loads(stripped)
except json.JSONDecodeError:
pass
# Strategy 2: extract the first balanced JSON object found anywhere in
# the text. Handles the case where Claude wraps the JSON in prose without
# any markdown fence ("The extracted graph is { ... }. Hope this helps!").
start = stripped.find("{")
if start != -1:
depth = 0
in_string = False
escape = False
for i in range(start, len(stripped)):
ch = stripped[i]
if escape:
escape = False
continue
if ch == "\\":
escape = True
continue
if ch == '"':
in_string = not in_string
continue
if in_string:
continue
if ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
try:
return json.loads(stripped[start : i + 1])
except json.JSONDecodeError:
break
print(
f"[graphify] LLM returned invalid JSON, skipping chunk "
f"(first 200 chars: {raw[:200]!r})",
file=sys.stderr,
)
return {"nodes": [], "edges": [], "hyperedges": []}
def _response_is_hollow(raw_content: str | None, parsed: dict) -> bool:
"""Detect a successful HTTP response that yielded no usable extraction.
A local model under load (most often Ollama) can return HTTP 200 with an
empty / null `message.content`, with whitespace, or with a half-generated
JSON prefix that fails to parse. All of these collapse to a "successful"
call producing zero nodes and zero edges. Without this check the chunk
is silently dropped from the corpus because no exception is raised and
`finish_reason` is `"stop"` rather than `"length"`. By flagging the
result as hollow, callers can re-route it through the same bisection
path used for context-window overflow and `finish_reason="length"`.
"""
if raw_content is None or not raw_content.strip():
return True
nodes = parsed.get("nodes")
edges = parsed.get("edges")
hyperedges = parsed.get("hyperedges")
return not nodes and not edges and not 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)
env_key = cfg.get("env_key")
if env_key:
return [env_key]
return []
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."""
keys = _backend_env_keys(backend)
return " or ".join(keys) if keys else "AWS_PROFILE or AWS_REGION"
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 _backend_pkg_hint(pkg: str, extra: str) -> str:
"""Package-missing message that works for the recommended `uv tool` install.
`uv tool install graphifyy` puts graphify in an isolated venv, so a plain
`pip install <pkg>` never reaches it - the friction a user hits when a
backend needs anthropic/openai/boto3 and the only advice was "pip install".
Point at the extra and the uv path first, then the pip/venv fallback.
"""
return (
f"the '{pkg}' package is required for this backend but is not installed. "
f"Install it with: uv tool install \"graphifyy[{extra}]\" --force "
f"(uv tool), or pip install {pkg} (pip/venv install)."
)
def _call_openai_compat(
base_url: str,
api_key: str,
model: str,
user_message: str,
temperature: float | None = 0,
reasoning_effort: str | None = None,
max_completion_tokens: int = 8192,
*,
backend: str = "",
deep_mode: bool = False,
) -> dict:
"""Call any OpenAI-compatible API (Kimi, OpenAI, etc.) and return parsed JSON."""
try:
from openai import OpenAI
except ImportError as exc:
extra = backend if backend in ("kimi", "gemini", "openai", "ollama") else "openai"
raise ImportError(_backend_pkg_hint("openai", extra)) from exc
# Local backends (ollama, llama.cpp, vLLM) routinely take >60s for a
# single chunk on a large model — far longer than the openai SDK's
# default. Honour GRAPHIFY_API_TIMEOUT (seconds) for explicit override;
# default to 600s, which is long enough for a 31B model on a 16k chunk
# but still bounds runaway connections (issue #792 addendum).
timeout_raw = os.environ.get("GRAPHIFY_API_TIMEOUT", "").strip()
timeout_s: float = 600.0
if timeout_raw:
try:
v = float(timeout_raw)
if v > 0:
timeout_s = v
except ValueError:
pass
client = OpenAI(api_key=api_key, base_url=base_url, timeout=timeout_s)
kwargs: dict = {
"model": model,
"messages": [
{"role": "system", "content": _extraction_system(deep=deep_mode)},
{"role": "user", "content": user_message},
],
"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"}}
# Ollama defaults num_ctx to 2048 and silently truncates prompts larger
# than that — the symptom is hollow 200 OK responses after the first few
# chunks (#798). We derive num_ctx from the actual prompt size so we don't
# over-allocate KV-cache VRAM. Over-allocation (e.g. 128k slots for an 8k
# prompt on a 31B model) exhausts VRAM by chunk 4 and produces the same
# hollow-200 symptom — just from a different direction (#798 follow-up).
# Formula: actual input tokens + output cap + system prompt headroom.
# Capped at 131072 (enough for the default 60k token_budget); env var wins.
if backend == "ollama":
num_ctx_raw = os.environ.get("GRAPHIFY_OLLAMA_NUM_CTX", "").strip()
# Auto-derive num_ctx from actual chunk size regardless — used as the
# fallback and for the mismatch check below.
estimated_input = len(user_message) // _CHARS_PER_TOKEN + 400
auto_num_ctx = min(estimated_input + max_completion_tokens + 2000, 131072)
auto_num_ctx = max(auto_num_ctx, 8192)
if num_ctx_raw:
try:
num_ctx = int(num_ctx_raw)
except ValueError:
# Bad env var: fall through to auto-derivation (not 131072 —
# hardcoding the cap is what causes OOM on constrained VRAM).
print(
f"[graphify] GRAPHIFY_OLLAMA_NUM_CTX={num_ctx_raw!r} is not a valid integer; "
f"using auto-derived value ({auto_num_ctx}).",
file=sys.stderr,
)
num_ctx = auto_num_ctx
else:
# Warn when the pinned value is smaller than the estimated input —
# Ollama silently truncates the prompt and returns empty responses.
if num_ctx < estimated_input:
print(
f"[graphify] warning: GRAPHIFY_OLLAMA_NUM_CTX={num_ctx} is smaller than "
f"the estimated chunk input (~{estimated_input} tokens). Ollama will "
f"silently truncate the prompt and return empty responses. "
f"Try --token-budget {max(1024, num_ctx // 3)} or increase NUM_CTX.",
file=sys.stderr,
)
else:
# Estimate input tokens: user_message chars / 4 (standard BPE
# heuristic) + 400 for the system prompt, then add output headroom.
num_ctx = auto_num_ctx
keep_alive = os.environ.get("GRAPHIFY_OLLAMA_KEEP_ALIVE", "30m")
kwargs["extra_body"] = {"options": {"num_ctx": num_ctx}, "keep_alive": keep_alive}
resp = client.chat.completions.create(**kwargs)
if not resp.choices or resp.choices[0].message is None:
raise ValueError("LLM returned empty or filtered response")
raw_content = resp.choices[0].message.content
result = _parse_llm_json(raw_content or "{}")
result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0
result["output_tokens"] = resp.usage.completion_tokens if resp.usage else 0
result["model"] = model
# `finish_reason == "length"` means the model hit max_completion_tokens
# mid-generation. The JSON we got back is truncated; callers should
# treat this as a signal to retry with smaller input.
result["finish_reason"] = resp.choices[0].finish_reason
# An overwhelmed local model (typically Ollama) can return HTTP 200 with
# empty / null content or unparseable half-generated JSON. The call looks
# successful, `finish_reason` is `"stop"`, and the chunk would be silently
# dropped from the corpus. Re-label as `"length"` so the adaptive retry
# layer bisects the chunk — same recovery as a true truncation.
if _response_is_hollow(raw_content, result) and result["finish_reason"] != "length":
print(
f"[graphify] {backend or 'backend'} returned a hollow response "
f"(content={'empty' if not (raw_content or '').strip() else 'no nodes/edges'}, "
f"output_tokens={result['output_tokens']}); "
"treating as truncation so adaptive retry can bisect the chunk.",
file=sys.stderr,
)
result["finish_reason"] = "length"
output_tokens = result["output_tokens"]
if output_tokens < 50 and backend == "ollama":
print(
"[graphify] warning: ollama returned very few tokens — likely causes: "
"(1) VRAM pressure: check `nvidia-smi` and reduce chunk size with "
"--token-budget (e.g. --token-budget 4096) or set "
"GRAPHIFY_OLLAMA_NUM_CTX to a smaller value; "
"(2) model too small for JSON instruction following — "
"try a larger model with --model (e.g. --model qwen2.5-coder:14b).",
file=sys.stderr,
)
return result
def _call_claude(api_key: str, model: str, user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False) -> dict:
"""Call Anthropic Claude directly (not via OpenAI compat layer)."""
try:
import anthropic
except ImportError as exc:
raise ImportError(_backend_pkg_hint("anthropic", "anthropic")) from exc
client = anthropic.Anthropic(api_key=api_key)
resp = client.messages.create(
model=model,
max_tokens=max_tokens,
system=_extraction_system(deep=deep_mode),
messages=[{"role": "user", "content": user_message}],
)
raw_content = resp.content[0].text if resp.content else None
result = _parse_llm_json(raw_content or "{}")
result["input_tokens"] = resp.usage.input_tokens if resp.usage else 0
result["output_tokens"] = resp.usage.output_tokens if resp.usage else 0
result["model"] = model
# Normalise Anthropic's `stop_reason` to the OpenAI-compat `finish_reason`
# vocabulary so the adaptive-retry layer doesn't have to know which
# backend produced the result.
result["finish_reason"] = "length" if resp.stop_reason == "max_tokens" else "stop"
if _response_is_hollow(raw_content, result) and result["finish_reason"] != "length":
print(
"[graphify] claude returned a hollow response; treating as "
"truncation so adaptive retry can bisect the chunk.",
file=sys.stderr,
)
result["finish_reason"] = "length"
return result
def _call_claude_cli(user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False) -> dict:
"""Call Claude via the locally-installed Claude Code CLI (`claude -p`).
Routes through the user's Claude Code subscription auth instead of a separate
ANTHROPIC_API_KEY. Useful for Pro/Max subscribers who don't want to provision
a pay-as-you-go API key just to run graphify's semantic pass.
"""
import platform
import shutil
import subprocess
# On Windows, npm installs `claude` as both `claude.ps1` and `claude.cmd`
# alongside each other. When PATHEXT lists `.PS1` before `.CMD`,
# `shutil.which("claude")` returns `claude.ps1`, which `CreateProcess`
# cannot execute directly — it raises `[WinError 2] The system cannot
# find the file specified`. `claude.cmd` IS executable by CreateProcess,
# so prefer it explicitly on Windows. See issue #1072.
claude_cmd = "claude"
if platform.system() == "Windows":
cmd_path = shutil.which("claude.cmd")
if cmd_path:
claude_cmd = cmd_path
elif shutil.which("claude") is None:
raise RuntimeError(
"Claude Code CLI not found on $PATH. Install from "
"https://claude.ai/code and run `claude` once to authenticate."
)
elif shutil.which("claude") is None:
raise RuntimeError(
"Claude Code CLI not found on $PATH. Install from "
"https://claude.ai/code and run `claude` once to authenticate."
)
# Use --system-prompt (replaces) instead of --append-system-prompt (adds
# to Claude Code's default coding-agent prompt). The default prompt
# pushes the model towards markdown + prose explanations, which conflict
# with the "raw JSON only" extraction instruction and cause ~30-50% of
# responses to come back wrapped in ```json fences or prefixed with a
# preamble — both of which fail the strict json.loads in _parse_llm_json.
# Replacing the default prompt eliminates the conflict at the source.
# Side benefit: cache-creation tokens per call drop ~19% in practice.
cli_args = [
claude_cmd, "-p",
"--output-format", "json",
"--no-session-persistence",
"--system-prompt", _extraction_system(deep=deep_mode),
]
# claude-cli defaults to Opus, which is overkill for the structured-JSON
# extraction graphify performs. GRAPHIFY_CLAUDE_CLI_MODEL=haiku (or
# sonnet, or a full model ID like claude-haiku-4-5-20251001) lets users
# opt into a cheaper / faster model. Default behaviour unchanged when
# the env var is unset.
cli_model = os.environ.get("GRAPHIFY_CLAUDE_CLI_MODEL", "").strip()
if cli_model:
cli_args.extend(["--model", cli_model])
proc = subprocess.run(
cli_args,
input=user_message,
capture_output=True,
text=True,
encoding="utf-8", # Force UTF-8 — prevents UnicodeEncodeError on Windows cp1252
timeout=600,
check=False,
)
if proc.returncode != 0:
raise RuntimeError(
f"claude -p exited {proc.returncode}: {proc.stderr.strip()[:500]}"
)
try:
envelope = json.loads(proc.stdout)
except json.JSONDecodeError as exc:
raise RuntimeError(
f"claude -p produced unparseable JSON envelope: {exc}; "
f"first 500 chars of stdout: {proc.stdout[:500]!r}"
) from exc
raw_content = envelope.get("result", "")
result = _parse_llm_json(raw_content or "{}")
usage = envelope.get("usage") or {}
result["input_tokens"] = (
int(usage.get("input_tokens", 0) or 0)
+ int(usage.get("cache_read_input_tokens", 0) or 0)
+ int(usage.get("cache_creation_input_tokens", 0) or 0)
)
result["output_tokens"] = int(usage.get("output_tokens", 0) or 0)
model_usage = envelope.get("modelUsage") or {}
result["model"] = next(iter(model_usage), "claude-code-plan")
stop_reason = envelope.get("stop_reason", "")
result["finish_reason"] = "length" if stop_reason == "max_tokens" else "stop"
if _response_is_hollow(raw_content, result) and result["finish_reason"] != "length":
print(
"[graphify] claude-cli returned a hollow response; treating as "
"truncation so adaptive retry can bisect the chunk.",
file=sys.stderr,
)
result["finish_reason"] = "length"
return result
def _call_bedrock(model: str, user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False) -> dict:
"""Call AWS Bedrock via boto3 Converse API using the standard AWS credential chain."""
try:
import boto3
import botocore.exceptions
except ImportError as exc:
raise ImportError(
"AWS Bedrock extraction requires boto3. Run: pip install graphifyy[bedrock]"
) from exc
region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1"
profile = os.environ.get("AWS_PROFILE")
session = boto3.Session(profile_name=profile, region_name=region)
client = session.client("bedrock-runtime")
try:
resp = client.converse(
modelId=model,
system=[{"text": _extraction_system(deep=deep_mode)}],
messages=[{"role": "user", "content": [{"text": user_message}]}],
inferenceConfig={"maxTokens": max_tokens, "temperature": 0},
)
except botocore.exceptions.ClientError as exc:
code = exc.response["Error"]["Code"]
msg = exc.response["Error"]["Message"]
raise RuntimeError(f"Bedrock API error ({code}): {msg}") from exc
text = resp.get("output", {}).get("message", {}).get("content", [{}])[0].get("text", "{}")
result = _parse_llm_json(text)
usage = resp.get("usage", {})
result["input_tokens"] = usage.get("inputTokens", 0)
result["output_tokens"] = usage.get("outputTokens", 0)
result["model"] = model
result["finish_reason"] = "length" if resp.get("stopReason") == "max_tokens" else "stop"
if _response_is_hollow(text, result) and result["finish_reason"] != "length":
print(
"[graphify] bedrock returned a hollow response; treating as "
"truncation so adaptive retry can bisect the chunk.",
file=sys.stderr,
)
result["finish_reason"] = "length"
return result
def extract_files_direct(
files: list[Path],
backend: str | None = None,
api_key: str | None = None,
model: str | None = None,
root: Path = Path("."),
*,
deep_mode: bool = False,
) -> dict:
"""Extract semantic nodes/edges from a list of files using the given backend.
Returns dict with nodes, edges, hyperedges, input_tokens, output_tokens.
Raises ValueError for unknown backends or when no API key is configured.
Raises ImportError if SDK missing.
"""
if backend is None:
backend = detect_backend()
if backend is None:
raise ValueError(
"No LLM backend configured. Set one of: GEMINI_API_KEY, ANTHROPIC_API_KEY, "
"OPENAI_API_KEY, DEEPSEEK_API_KEY, MOONSHOT_API_KEY, OLLAMA_BASE_URL, "
"or AWS credentials. Pass backend= explicitly to select a provider."
)
if backend not in BACKENDS:
raise ValueError(f"Unknown backend {backend!r}. Available: {sorted(BACKENDS)}")
cfg = BACKENDS[backend]
key = api_key or _get_backend_api_key(backend)
if not key and backend == "ollama":
# Ollama ignores auth but the OpenAI client library requires a non-empty
# string. Use a placeholder and surface a visible warning so this never
# silently routes traffic without the user realising — see F-029.
ollama_url = os.environ.get("OLLAMA_BASE_URL", cfg.get("base_url", ""))
_validate_ollama_base_url(ollama_url)
print(
"[graphify] WARNING: ollama backend selected with no OLLAMA_API_KEY set; "
f"sending corpus to {ollama_url}. Set OLLAMA_API_KEY (any non-empty value) "
"to suppress this warning.",
file=sys.stderr,
)
key = "ollama"
if not key and backend not in ("bedrock", "claude-cli"):
raise ValueError(
f"No API key for backend '{backend}'. "
f"Set {_format_backend_env_keys(backend)} or pass api_key=."
)
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, deep_mode=deep_mode)
if backend == "claude-cli":
return _call_claude_cli(user_msg, max_tokens=max_out, deep_mode=deep_mode)
if backend == "bedrock":
return _call_bedrock(mdl, user_msg, max_tokens=max_out, deep_mode=deep_mode)
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=_resolve_max_tokens(cfg.get("max_completion_tokens", 8192)),
backend=backend,
deep_mode=deep_mode,
)
def _estimate_file_tokens(path: Path) -> int:
"""Estimate the prompt-token cost of a single file under `_read_files` rules.
Uses tiktoken (`cl100k_base`) when available for accurate counts. Falls back
to the chars/4 heuristic if tiktoken is not installed. Both paths cap at
`_FILE_CHAR_CAP` to match `_read_files`'s truncation, plus a constant for
the `=== rel ===` separator. Returns 0 for unreadable paths so they don't
blow up packing.
"""
if _TOKENIZER is None:
try:
size = path.stat().st_size
except OSError:
return 0
chars = min(size, _FILE_CHAR_CAP) + _PER_FILE_OVERHEAD_CHARS
return chars // _CHARS_PER_TOKEN
try:
content = path.read_text(encoding="utf-8", errors="replace")[:_FILE_CHAR_CAP]
except OSError:
return 0
return len(_TOKENIZER.encode(content)) + (_PER_FILE_OVERHEAD_CHARS // _CHARS_PER_TOKEN)
def _pack_chunks_by_tokens(
files: list[Path],
token_budget: int,
) -> list[list[Path]]:
"""Greedily pack files into chunks that fit a token budget.
Files are first grouped by parent directory so related artifacts share a
chunk (cross-file edges are more likely to be extracted within a chunk
than across chunks). Within each directory, files are added one at a
time; a chunk is closed when adding the next file would exceed the
budget. A single file larger than the budget gets its own chunk and the
caller is expected to handle the API error if it actually overflows the
model's context window — packing can't shrink one big file.
"""
if token_budget <= 0:
raise ValueError(f"token_budget must be positive, got {token_budget}")
by_dir: dict[Path, list[Path]] = {}
for f in files:
by_dir.setdefault(f.parent, []).append(f)
chunks: list[list[Path]] = []
current: list[Path] = []
current_tokens = 0
for directory in sorted(by_dir):
for path in by_dir[directory]:
cost = _estimate_file_tokens(path)
if current and current_tokens + cost > token_budget:
chunks.append(current)
current = []
current_tokens = 0
current.append(path)
current_tokens += cost
if current:
chunks.append(current)
return chunks
_CONTEXT_EXCEEDED_MARKERS = (
"context size",
"context length",
"context_length",
"context window",
"n_keep",
"exceeds the available",
"n_ctx",
"maximum context",
"too many tokens",
"prompt is too long",
"context_length_exceeded",
)
def _looks_like_context_exceeded(exc: BaseException) -> bool:
"""Heuristically classify an exception as a context-window overflow.
Different backends raise different exception types and messages for the
same underlying problem ("the prompt + max_completion_tokens did not fit
in the model's context window"). We match on substrings of the stringified
exception so the retry layer can recover without depending on a specific
SDK class. False positives are cheap (we'll re-extract on halves and
likely recover); false negatives are expensive (chunk fails entirely).
"""
msg = str(exc).lower()
return any(marker in msg for marker in _CONTEXT_EXCEEDED_MARKERS)
def _extract_with_adaptive_retry(
chunk: list[Path],
backend: str,
api_key: str | None,
model: str | None,
root: Path,
max_depth: int,
_depth: int = 0,
*,
deep_mode: bool = False,
) -> dict:
"""Extract a chunk; if the response is truncated (`finish_reason="length"`)
or the API rejects the prompt as too large for the model's context window,
split the chunk in half and recurse.
Three signals drive the retry, all funnelled through the same code:
- `finish_reason == "length"` — the model accepted the input but ran out of
`max_completion_tokens` mid-output. The truncated JSON is unparseable, so
we discard it and re-extract on smaller inputs that produce shorter
outputs.
- context-window-exceeded API errors — the model rejected the input
outright (HTTP 400 from LM Studio, llama.cpp, vLLM, OpenAI, etc.).
Without a retry the whole chunk would fail with no output. Splitting in
half is the same recovery as for the `length` case and works for the
same reason.
- hollow successful responses — the model returned HTTP 200 with empty,
null, or unparseable content (typical of a local Ollama under load).
`_call_openai_compat` re-labels these as `finish_reason="length"` so they
take the same recovery path; without that the chunk would be silently
dropped from the corpus.
Recursion is capped at `max_depth` to bound worst-case cost. A chunk of N
files can split into up to 2**max_depth pieces — at depth=3 that's 8x. If
still failing at the cap, we surface the (likely empty) result with a
warning rather than infinite-loop.
A single-file chunk that overflows is unrecoverable here — we can't make
one file smaller than itself, so we return what we got and warn.
"""
try:
result = extract_files_direct(
chunk, backend=backend, api_key=api_key, model=model, root=root, deep_mode=deep_mode
)
except Exception as exc: # noqa: BLE001 — re-raise unless it's a known context overflow
if not _looks_like_context_exceeded(exc):
raise
if len(chunk) <= 1:
print(
f"[graphify] single-file chunk {chunk[0]} exceeds model context "
f"and cannot be split further: {exc}",
file=sys.stderr,
)
return {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, "model": model, "finish_reason": "stop"}
if _depth >= max_depth:
print(
f"[graphify] chunk of {len(chunk)} still overflows context at "
f"recursion depth {_depth} (max {max_depth}) — dropping",
file=sys.stderr,
)
return {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, "model": model, "finish_reason": "stop"}
print(
f"[graphify] chunk of {len(chunk)} exceeded context at depth "
f"{_depth} ({type(exc).__name__}); splitting in half and retrying",
file=sys.stderr,
)
mid = len(chunk) // 2
left = _extract_with_adaptive_retry(
chunk[:mid], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode
)
right = _extract_with_adaptive_retry(
chunk[mid:], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode
)
return {
"nodes": left.get("nodes", []) + right.get("nodes", []),
"edges": left.get("edges", []) + right.get("edges", []),
"hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []),
"input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0),
"output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0),
"model": model,
"finish_reason": "stop",
}
if result.get("finish_reason") != "length":
return result
if len(chunk) <= 1:
print(
f"[graphify] single-file chunk {chunk[0]} truncated at "
f"max_completion_tokens — partial result kept",
file=sys.stderr,
)
return result
if _depth >= max_depth:
print(
f"[graphify] chunk of {len(chunk)} still truncated at recursion "
f"depth {_depth} (max {max_depth}) — partial result kept",
file=sys.stderr,
)
return result
print(
f"[graphify] chunk of {len(chunk)} truncated at depth {_depth}, "
f"splitting into halves of {len(chunk) // 2} and "
f"{len(chunk) - len(chunk) // 2}",
file=sys.stderr,
)
mid = len(chunk) // 2
left = _extract_with_adaptive_retry(
chunk[:mid], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode
)
right = _extract_with_adaptive_retry(
chunk[mid:], backend, api_key, model, root, max_depth, _depth + 1, deep_mode=deep_mode
)
return {
"nodes": left.get("nodes", []) + right.get("nodes", []),
"edges": left.get("edges", []) + right.get("edges", []),
"hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []),
"input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0),
"output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0),
"model": result.get("model"),
# Both halves either succeeded or have already surfaced their own
# truncation warning; the merged result is no longer truncated as a
# logical unit.
"finish_reason": "stop",
}
def extract_corpus_parallel(
files: list[Path],
backend: str = "kimi",
api_key: str | None = None,
model: str | None = None,
root: Path = Path("."),
chunk_size: int = 20,
on_chunk_done: Callable | None = None,
token_budget: int | None = 60_000,
max_concurrency: int = 4,
max_retry_depth: int = 3,
deep_mode: bool = False,
) -> dict:
"""Extract a corpus in chunks, merging results.
Chunking strategy:
- If `token_budget` is set (default 60_000), files are packed to fit
the budget and grouped by parent directory. This avoids the worst
case where 20 randomly-grouped files exceed a model's context
window in a single request.
- If `token_budget=None`, falls back to the legacy fixed-count
`chunk_size` packing for backwards compatibility.
Concurrency:
- Chunks run in parallel via a thread pool capped at `max_concurrency`
(default 4 — conservative to stay under provider rate limits).
- Set `max_concurrency=1` to force sequential execution.
Adaptive retry on truncation:
- When the LLM returns `finish_reason="length"` (output truncated at
`max_completion_tokens`), the chunk is split in half and each half
re-extracted recursively, up to `max_retry_depth` levels deep
(default 3 → max 8x expansion of one chunk).
- This is signal-driven: chunks too dense to fit in one response
self-heal by splitting until they do, while well-sized chunks pay
no extra cost. Set `max_retry_depth=0` to disable retries.
`on_chunk_done(idx, total, chunk_result)` fires once per chunk as it
completes (in completion order, not submission order). `idx` is the
chunk's submission index so callers can correlate progress. The
callback fires once per top-level chunk; recursive splits are merged
transparently before the callback is invoked.
Returns merged dict with nodes, edges, hyperedges, input_tokens,
output_tokens. Failed chunks are logged to stderr and skipped — one bad
chunk does not abort the run.
"""
if token_budget is not None:
chunks = _pack_chunks_by_tokens(files, token_budget=token_budget)
else:
chunks = [files[i:i + chunk_size] for i in range(0, len(files), chunk_size)]
merged: dict = {
"nodes": [], "edges": [], "hyperedges": [],
"input_tokens": 0, "output_tokens": 0,
"failed_chunks": 0, # count of chunks that raised — loud failure on chunk errors
}
total = len(chunks)
def _run_one(idx: int, chunk: list[Path]) -> tuple[int, dict | None, Exception | None]:
t0 = time.time()
try:
result = _extract_with_adaptive_retry(
chunk,
backend=backend,
api_key=api_key,
model=model,
root=root,
max_depth=max_retry_depth,
deep_mode=deep_mode,
)
result["elapsed_seconds"] = round(time.time() - t0, 2)
return idx, result, None
except Exception as exc: # noqa: BLE001 — caller-facing surface, log + continue
return idx, None, exc
# Ollama serves one request at a time per loaded model on a single GPU.
# Four concurrent 60k-token requests cause VRAM pressure and hollow
# responses after 3-4 chunks (#798). Force serial unless the user opts in.
if backend == "ollama" and os.environ.get("GRAPHIFY_OLLAMA_PARALLEL", "").strip() != "1":
max_concurrency = 1
# claude-cli shells out to a Claude Code session; parallel subprocesses conflict
# over session state. Force serial unless the user explicitly opts in.
if backend == "claude-cli" and os.environ.get("GRAPHIFY_CLAUDE_CLI_PARALLEL", "").strip() != "1":
max_concurrency = 1
workers = max(1, min(max_concurrency, total))
if workers == 1:
# Avoid thread pool overhead for single-worker runs (and keep
# callback ordering identical to the pre-refactor sequential path).
for idx, chunk in enumerate(chunks):
_, result, exc = _run_one(idx, chunk)
if exc is not None:
print(f"[graphify] chunk {idx + 1}/{total} failed: {exc}", file=sys.stderr)
merged["failed_chunks"] += 1
continue
assert result is not None
_merge_into(merged, result)
if callable(on_chunk_done):
on_chunk_done(idx, total, result)
else:
with ThreadPoolExecutor(max_workers=workers) as pool:
futures = [pool.submit(_run_one, idx, chunk) for idx, chunk in enumerate(chunks)]
for future in as_completed(futures):
idx, result, exc = future.result()
if exc is not None:
print(
f"[graphify] chunk {idx + 1}/{total} failed: {exc}",
file=sys.stderr,
)
merged["failed_chunks"] += 1
continue
assert result is not None
_merge_into(merged, result)
if callable(on_chunk_done):
on_chunk_done(idx, total, result)
# Loud failure summary — surface chunk failures at end so they're never
# buried mid-log. Exit 0 preserved for caller compatibility; the
# summary block makes the problem visible.
if merged["failed_chunks"] > 0:
print(
f"[graphify] WARNING: {merged['failed_chunks']}/{total} semantic chunk(s) failed"
" — see errors above. Partial results returned.",
file=sys.stderr,
)
return merged
def _merge_into(merged: dict, result: dict) -> None:
"""Append a chunk result into the running merged accumulator."""
merged["nodes"].extend(result.get("nodes", []))
merged["edges"].extend(result.get("edges", []))
merged["hyperedges"].extend(result.get("hyperedges", []))
merged["input_tokens"] += result.get("input_tokens", 0)
merged["output_tokens"] += result.get("output_tokens", 0)
def _call_llm(prompt: str, *, backend: str, max_tokens: int = 200) -> str:
"""Send a plain-text prompt to `backend` and return the model's text reply.
Used by lightweight callers (e.g. `graphify.dedup` LLM tiebreaker) that
don't need the full extraction prompt or JSON-shaped output. Mirrors the
backend dispatch logic of `extract_files_direct` but skips the
`_EXTRACTION_SYSTEM` prompt and JSON parsing.
Previously `graphify.dedup` imported a `_call_llm` symbol that did not
exist in this module, so the LLM tiebreaker silently no-op'd on
`ImportError` (F-038). Adding the function here re-enables it.
"""
if backend not in BACKENDS:
raise ValueError(f"Unknown backend {backend!r}")
cfg = BACKENDS[backend]
key = _get_backend_api_key(backend)
if not key and backend == "ollama":
ollama_url = os.environ.get("OLLAMA_BASE_URL", cfg.get("base_url", ""))
_validate_ollama_base_url(ollama_url)
key = "ollama"
if not key and backend not in ("bedrock", "claude-cli"):
raise ValueError(
f"No API key for backend '{backend}'. Set {_format_backend_env_keys(backend)}."
)
mdl = _default_model_for_backend(backend)
if backend == "claude":
try:
import anthropic
except ImportError as exc:
raise ImportError(_backend_pkg_hint("anthropic", "anthropic")) from exc
client = anthropic.Anthropic(api_key=key)
resp = client.messages.create(
model=mdl,
max_tokens=max_tokens,
messages=[{"role": "user", "content": prompt}],
)
return resp.content[0].text if resp.content else ""
if backend == "claude-cli":
import shutil, subprocess
if shutil.which("claude") is None:
raise RuntimeError("Claude Code CLI not found on $PATH")
proc = subprocess.run(
["claude", "-p", "--output-format", "json", "--no-session-persistence"],
input=prompt,
capture_output=True,
text=True,
encoding="utf-8", # Force UTF-8 — prevents UnicodeEncodeError on Windows cp1252
timeout=600,
check=False,
)
if proc.returncode != 0:
raise RuntimeError(f"claude -p exited {proc.returncode}: {proc.stderr.strip()[:500]}")
try:
envelope = json.loads(proc.stdout)
except json.JSONDecodeError as exc:
raise RuntimeError(f"claude -p produced unparseable JSON envelope: {exc}") from exc
return envelope.get("result", "")
if backend == "bedrock":
try:
import boto3
except ImportError as exc:
raise ImportError(_backend_pkg_hint("boto3", "bedrock")) from exc
region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1"
profile = os.environ.get("AWS_PROFILE")
session = boto3.Session(profile_name=profile, region_name=region)
client = session.client("bedrock-runtime")
resp = client.converse(
modelId=mdl,
messages=[{"role": "user", "content": [{"text": prompt}]}],
inferenceConfig={"maxTokens": max_tokens, "temperature": 0},
)
return resp.get("output", {}).get("message", {}).get("content", [{}])[0].get("text", "")
# OpenAI-compatible (kimi, openai, gemini, ollama)
try:
from openai import OpenAI
except ImportError as exc:
raise ImportError(_backend_pkg_hint("openai", "openai")) from exc
client = OpenAI(api_key=key, base_url=cfg["base_url"])
kwargs: dict = {
"model": mdl,
"messages": [{"role": "user", "content": prompt}],
"max_completion_tokens": max_tokens,
}
temperature = cfg.get("temperature", 0)
if temperature is not None:
kwargs["temperature"] = temperature
if cfg.get("reasoning_effort"):
kwargs["reasoning_effort"] = cfg["reasoning_effort"]
if "moonshot" in cfg["base_url"]:
kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
resp = client.chat.completions.create(**kwargs)
if not resp.choices or resp.choices[0].message is None:
raise ValueError("LLM returned empty or filtered response")
return resp.choices[0].message.content or ""
def estimate_cost(backend: str, input_tokens: int, output_tokens: int) -> float:
"""Estimate USD cost for a given token count using published pricing."""
if backend not in BACKENDS:
return 0.0
p = BACKENDS[backend]["pricing"]
return (input_tokens * p["input"] + output_tokens * p["output"]) / 1_000_000
def _ollama_host_is_link_local_or_metadata(host: str) -> bool:
"""True if *host* is, or resolves to, a link-local / cloud-metadata address.
Resolves the name so an alias pointing at 169.254.169.254 is caught too, not
just a literal IP. General private/LAN addresses are deliberately NOT treated
as metadata: people do run Ollama on trusted LAN boxes, so those only warn.
"""
import ipaddress
import socket
if host in ("metadata.google.internal", "metadata.google.com", "0.0.0.0", "::", "[::]"): # nosec B104 - blocklist, not a bind
return True
if host.startswith("169.254."): # link-local literal, includes the metadata IP
return True
try:
infos = socket.getaddrinfo(host, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
except (socket.gaierror, UnicodeError, OSError):
return False
for info in infos:
try:
ip = ipaddress.ip_address(info[4][0])
except ValueError:
continue
if ip.is_link_local: # 169.254.0.0/16 and fe80::/10 (includes the metadata IP)
return True
return False
def _validate_ollama_base_url(url: str, *, warn: bool = True) -> None:
"""Warn if OLLAMA_BASE_URL looks unsafe; hard-block link-local/metadata (F3).
Sending an entire corpus to a non-loopback http:// endpoint silently leaks
proprietary code, but some users genuinely run Ollama on a LAN host they
trust, so a general non-loopback target only warns. A link-local or cloud
metadata address (169.254.x, metadata.google.*, or any host that resolves to
one) is never a legitimate Ollama host and is a classic SSRF target, so we
fail closed with a ValueError there regardless of *warn*. Pass warn=False for
an early gate that should hard-block but leave the user-facing warning to the
later in-flow call.
"""
try:
from urllib.parse import urlparse
parsed = urlparse(url)
except Exception:
if warn:
print(
f"[graphify] WARNING: OLLAMA_BASE_URL={url!r} is not a parseable URL.",
file=sys.stderr,
)
return
if parsed.scheme not in ("http", "https"):
if warn:
print(
f"[graphify] WARNING: OLLAMA_BASE_URL has unexpected scheme {parsed.scheme!r}; "
"expected http or https.",
file=sys.stderr,
)
return
host = (parsed.hostname or "").lower()
if _ollama_host_is_link_local_or_metadata(host):
raise ValueError(
f"OLLAMA_BASE_URL points at a link-local/metadata address ({host!r}); refusing to "
"send the corpus there. Set it to a real Ollama host."
)
is_loopback = host in ("localhost", "127.0.0.1", "::1") or host.startswith("127.")
if warn and not is_loopback:
scheme_note = " (UNENCRYPTED)" if parsed.scheme == "http" else ""
print(
f"[graphify] WARNING: OLLAMA_BASE_URL points to non-loopback host {host!r}{scheme_note}. "
"Your full corpus will be sent to that endpoint. "
"Set OLLAMA_BASE_URL=http://localhost:11434/v1 to keep extraction local.",
file=sys.stderr,
)
def detect_backend() -> str | None:
"""Return the name of whichever backend has an API key set, or None.
Priority: gemini → kimi → claude → openai → bedrock → ollama (last, opt-in).
Ollama is intentionally checked LAST so a paid API key (Anthropic/OpenAI/etc.)
is never silently shadowed by an incidental OLLAMA_BASE_URL in the environment
— see security finding F-002/F-029. Setting OLLAMA_BASE_URL alongside a paid
key now keeps you on the paid backend; remove the paid key (or pass
--backend ollama explicitly) to route to the local model.
"""
for backend in ("gemini", "kimi", "claude", "openai", "deepseek"):
if _get_backend_api_key(backend):
return backend
if os.environ.get("AWS_PROFILE") or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION"):
return "bedrock"
ollama_url = os.environ.get("OLLAMA_BASE_URL")
if ollama_url:
_validate_ollama_base_url(ollama_url)
return "ollama"
for name in BACKENDS:
if name not in ("gemini", "kimi", "claude", "openai", "deepseek", "bedrock", "ollama", "claude-cli"):
if _get_backend_api_key(name):
return name
return None
# ── Community labeling ────────────────────────────────────────────────────────
# When graphify runs inside an orchestrating agent (Claude Code / Gemini CLI),
# the agent names communities itself per skill.md Step 5 - it reads the analysis
# file and writes 2-5 word names with its own reasoning, no API call. When
# graphify is run as a bare CLI (``graphify extract . --backend X``), there is no
# agent to do that step, so community labels stay ``Community 0/1/2...``. These
# helpers fill that gap: ask the configured backend to name communities in ONE
# batched call and return a complete ``{cid: name}`` map (#1097).
_LABEL_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*|\s*```\s*$", re.IGNORECASE)
_LABEL_MAX_COMMUNITIES = 200 # cap LLM-named communities; tail stays placeholder
_LABEL_TOP_K = 12 # node labels sampled per community for the prompt
_LABEL_MAXLEN = 60 # truncate individual labels to keep the prompt small
def _placeholder_community_labels(communities) -> dict[int, str]:
return {int(cid): f"Community {cid}" for cid in communities}
def _community_label_lines(G, communities, gods, max_communities, top_k):
"""One prompt line per community (largest first), sampling up to ``top_k``
representative node labels (god nodes first). Returns (lines, labeled_cids);
skips communities with no resolvable nodes."""
# gods may be node-id strings or god_nodes() dicts ({"id": ..., "label": ...}).
god_set = {g["id"] if isinstance(g, dict) else g for g in (gods or [])}
ordered = sorted(communities.items(), key=lambda kv: -len(kv[1]))
lines: list[str] = []
labeled_cids: list[int] = []
for cid, members in ordered[:max_communities]:
ranked = [m for m in members if m in god_set] + [m for m in members if m not in god_set]
names: list[str] = []
seen: set[str] = set()
for nid in ranked:
label = str(G.nodes[nid].get("label", nid)) if nid in G.nodes else str(nid)
label = label.strip().strip("()")[:_LABEL_MAXLEN]
if label and label.lower() not in seen:
seen.add(label.lower())
names.append(label)
if len(names) >= top_k:
break
if names:
lines.append(f"Community {cid}: {', '.join(names)}")
labeled_cids.append(int(cid))
return lines, labeled_cids
def _parse_label_response(text: str, labeled_cids: list[int]) -> dict[int, str]:
"""Parse the backend's JSON ``{cid: name}`` reply. Raises on non-JSON or a
non-object payload; silently ignores cids it didn't name."""
cleaned = _LABEL_FENCE_RE.sub("", text.strip())
if not cleaned.startswith("{"):
start, end = cleaned.find("{"), cleaned.rfind("}")
if start != -1 and end > start:
cleaned = cleaned[start:end + 1]
data = json.loads(cleaned)
if not isinstance(data, dict):
raise ValueError("label response is not a JSON object")
out: dict[int, str] = {}
for cid in labeled_cids:
name = data.get(str(cid))
if name is None:
name = data.get(cid)
if isinstance(name, str) and name.strip():
out[cid] = name.strip()
return out
def label_communities(
G,
communities,
*,
backend: str,
gods=None,
max_communities: int = _LABEL_MAX_COMMUNITIES,
top_k: int = _LABEL_TOP_K,
) -> dict[int, str]:
"""Return a complete ``{cid: name}`` map using ``backend`` for naming.
Placeholders (``Community N``) are used for any community the backend did not
name. Raises on backend/parse failure - callers that want graceful
degradation should use :func:`generate_community_labels`.
"""
labels = _placeholder_community_labels(communities)
lines, labeled_cids = _community_label_lines(G, communities, gods, max_communities, top_k)
if not lines:
return labels
prompt = (
"You are naming clusters in a knowledge graph. For each community below, "
"return a concise 2-5 word plain-language name describing what it is about "
"(e.g. \"Order Management\", \"Payment Flow\", \"Auth Middleware\"). "
"Respond ONLY with a JSON object mapping the community id (as a string) to "
"its name - no prose, no markdown fences.\n\n" + "\n".join(lines)
)
max_tokens = min(40 + 16 * len(labeled_cids), 4096)
text = _call_llm(prompt, backend=backend, max_tokens=max_tokens)
labels.update(_parse_label_response(text, labeled_cids))
return labels
def generate_community_labels(
G,
communities,
*,
backend: str | None = None,
gods=None,
quiet: bool = False,
) -> tuple[dict[int, str], str]:
"""CLI entry point: resolve a backend, name communities, and degrade to
``Community N`` placeholders on any failure (no backend, API error, malformed
reply). Returns ``(labels, source)`` where source is ``"llm"`` or
``"placeholder"``. Never raises."""
if backend is None:
try:
backend = detect_backend()
except Exception:
backend = None
if not backend:
if not quiet:
print(
"[graphify label] no LLM backend configured; keeping Community N "
"placeholders. Set an API key (e.g. GOOGLE_API_KEY) or pass --backend.",
file=sys.stderr,
)
return _placeholder_community_labels(communities), "placeholder"
try:
labels = label_communities(G, communities, backend=backend, gods=gods)
return labels, "llm"
except Exception as exc:
if not quiet:
print(
f"[graphify label] warning: community labeling failed ({exc}); "
"using Community N placeholders.",
file=sys.stderr,
)
return _placeholder_community_labels(communities), "placeholder"