graphite-code 0.3.0__py3-none-any.whl

graphite/io.py

@@ -0,0 +1,32 @@
+"""Safe file-writing helpers for Graphite artifacts."""
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+
+def atomic_write_text(path: Path, text: str) -> None:
+    """Atomically replace a text file after fsyncing the temporary file."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent))
+    tmp_path = Path(tmp_name)
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8", newline="") as f:
+            f.write(text)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp_path, path)
+    except Exception:
+        try:
+            tmp_path.unlink(missing_ok=True)
+        finally:
+            raise
+
+
+def atomic_write_json(path: Path, data: Any, *, indent: int | None = 2) -> None:
+    """Atomically write JSON with deterministic UTF-8 encoding."""
+    text = json.dumps(data, ensure_ascii=False, indent=indent)
+    atomic_write_text(path, text)

graphite/listing.py

@@ -0,0 +1,51 @@
+"""Render a bounded list so that emptiness and truncation are always visible.
+
+Every human-output listing goes through :func:`listing_lines`. Two guarantees
+hold by construction: a header is never printed over an empty body, and a
+capped list always names how many items it dropped.
+
+See docs/superpowers/specs/2026-07-26-listing-marker-contract-design.md.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+__all__ = ["listing_lines"]
+
+
+def listing_lines(
+    values: Sequence[Any],
+    render: Callable[[Any], str] = str,
+    *,
+    header: str | None = None,
+    cap: int | None = None,
+    indent: str = "  ",
+    empty: str | None = "none found",
+    more_hint: str | None = None,
+) -> list[str]:
+    """Lines for one bounded listing: header, body, and a truncation marker.
+
+    ``cap`` of None (or <= 0) means no cap. ``empty=None`` means the whole
+    listing is omitted when there is nothing to show -- header included, since
+    a header over an empty body is the defect this module exists to prevent.
+    ``more_hint`` is appended verbatim; callers supply their own separator.
+    """
+    if not values and empty is None:
+        return []
+
+    lines: list[str] = []
+    if header is not None:
+        lines.append(header)
+
+    if not values:
+        lines.append(f"{indent}- {empty}")
+        return lines
+
+    shown = list(values) if cap is None or cap <= 0 else list(values[:cap])
+    lines.extend(f"{indent}- {render(value)}" for value in shown)
+    dropped = len(values) - len(shown)
+    if dropped > 0:
+        lines.append(f"{indent}... {dropped} more{more_hint or ''}")
+    return lines

graphite/llm.py

@@ -0,0 +1,518 @@
+"""Provider-agnostic LLM enrichment primitives for explicit overlays only."""
+from __future__ import annotations
+
+import ipaddress
+import json
+from dataclasses import replace
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+from typing import Any, Literal, Protocol
+from urllib.parse import urlsplit
+
+from .config import Config
+
+CANONICAL_ENRICHMENT_MIGRATION_MESSAGE = (
+    "[graphite] canonical graph commands do not accept LLM enrichment; "
+    "use 'graphite overlay build' after building the canonical graph"
+)
+
+MAX_RESPONSE_BYTES = 64 * 1024
+DEFAULT_MAX_OUTPUT_TOKENS = 512
+PROBE_MAX_OUTPUT_TOKENS = 16
+MAX_OUTPUT_TOKENS = 4096
+ProviderErrorCategory = Literal[
+    "configuration",
+    "authentication",
+    "timeout",
+    "connection",
+    "provider_error",
+]
+_PROVIDER_ERROR_CATEGORIES = frozenset(
+    {"configuration", "authentication", "timeout", "connection", "provider_error"}
+)
+
+
+@dataclass(frozen=True)
+class CompletionResult:
+    text: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+
+
+class CompletionProvider(Protocol):
+    name: str
+
+    def complete(self, system: str, user: str) -> CompletionResult:
+        ...
+
+
+class LLMConfigurationError(ValueError):
+    """Raised when optional LLM enrichment is enabled but misconfigured."""
+
+
+class LLMProviderError(RuntimeError):
+    """Fixed-category provider failure that never contains remote response data."""
+
+    def __init__(self, category: ProviderErrorCategory) -> None:
+        self.category: ProviderErrorCategory = (
+            category if category in _PROVIDER_ERROR_CATEGORIES else "provider_error"
+        )
+        super().__init__(f"LLM provider failure: {self.category}")
+
+
+def _bounded_output_tokens(value: object) -> int:
+    if not isinstance(value, int) or isinstance(value, bool):
+        return DEFAULT_MAX_OUTPUT_TOKENS
+    return min(max(value, 1), MAX_OUTPUT_TOKENS)
+
+
+class _NoRedirectHandler(urllib.request.HTTPRedirectHandler):
+    """Turn every redirect into an HTTPError without forwarding credentials."""
+
+    def redirect_request(
+        self,
+        req: urllib.request.Request,
+        fp: Any,
+        code: int,
+        msg: str,
+        headers: Any,
+        newurl: str,
+    ) -> None:
+        del req, fp, code, msg, headers, newurl
+        return None
+
+
+def _open_no_redirect(request: urllib.request.Request, timeout: float) -> Any:
+    opener = urllib.request.build_opener(_NoRedirectHandler())
+    return opener.open(request, timeout=timeout)
+
+
+class OpenAICompatibleProvider:
+    """Chat-completions adapter for OpenAI-compatible HTTP APIs.
+
+    This intentionally avoids a vendor SDK. It works with OpenAI-compatible gateways such as
+    OpenAI, OpenRouter, LM Studio, vLLM, Groq-style endpoints, and local servers that expose
+    `/v1/chat/completions`.
+    """
+
+    name = "openai-compatible"
+
+    def __init__(self, cfg: Config) -> None:
+        base_url = cfg.llm_base_url or _default_openai_base_url(cfg.llm_provider)
+        if not base_url:
+            raise LLMConfigurationError(
+                "GRAPHITE_LLM_BASE_URL is required for openai-compatible providers"
+            )
+        self.base_url = base_url.rstrip("/")
+        _validate_llm_base_url(self.base_url, provider=cfg.llm_provider)
+        self.model = cfg.llm_model or _default_model(cfg.llm_provider)
+        self.api_key = cfg.llm_api_key
+        self.timeout = cfg.llm_timeout_seconds
+        self.seed = cfg.seed
+        self.max_output_tokens = _bounded_output_tokens(cfg.llm_max_output_tokens)
+
+    def complete(self, system: str, user: str) -> CompletionResult:
+        payload = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ],
+            "temperature": 0.1,
+            "stream": False,
+            "max_tokens": self.max_output_tokens,
+        }
+        if self.seed is not None:
+            payload["seed"] = self.seed
+        data = json.dumps(payload).encode("utf-8")
+        headers = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        request = urllib.request.Request(
+            f"{self.base_url}/chat/completions",
+            data=data,
+            headers=headers,
+            method="POST",
+        )
+        body = _urlopen_json(request, self.timeout)
+        choices = body.get("choices") or []
+        text = ""
+        if choices:
+            message = choices[0].get("message") or {}
+            text = str(message.get("content") or "").strip()
+        usage = body.get("usage") or {}
+        return CompletionResult(
+            text=text,
+            input_tokens=int(usage.get("prompt_tokens") or 0),
+            output_tokens=int(usage.get("completion_tokens") or 0),
+            total_tokens=int(usage.get("total_tokens") or 0),
+        )
+
+
+class OllamaProvider:
+    """Native Ollama chat adapter using `/api/chat`."""
+
+    name = "ollama"
+
+    def __init__(self, cfg: Config) -> None:
+        self.base_url = (cfg.llm_base_url or "http://localhost:11434").rstrip("/")
+        _validate_llm_base_url(self.base_url, provider=cfg.llm_provider)
+        self.model = cfg.llm_model or "llama3.1"
+        self.timeout = cfg.llm_timeout_seconds
+        self.seed = cfg.seed
+        self.max_output_tokens = _bounded_output_tokens(cfg.llm_max_output_tokens)
+
+    def complete(self, system: str, user: str) -> CompletionResult:
+        payload = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ],
+            "stream": False,
+            "options": {
+                "temperature": 0.1,
+                "seed": self.seed,
+                "num_predict": self.max_output_tokens,
+            },
+        }
+        request = urllib.request.Request(
+            f"{self.base_url}/api/chat",
+            data=json.dumps(payload).encode("utf-8"),
+            headers={"Content-Type": "application/json"},
+            method="POST",
+        )
+        body = _urlopen_json(request, self.timeout)
+        message = body.get("message") or {}
+        input_tokens = int(body.get("prompt_eval_count") or 0)
+        output_tokens = int(body.get("eval_count") or 0)
+        return CompletionResult(
+            text=str(message.get("content") or "").strip(),
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=input_tokens + output_tokens,
+        )
+
+
+def make_provider(cfg: Config) -> CompletionProvider:
+    provider = cfg.llm_provider.strip().lower().replace("_", "-")
+    if provider in {"ollama", "local"}:
+        return OllamaProvider(cfg)
+    if provider in {"openai", "openai-compatible", "compatible", "lmstudio", "lm-studio", "vllm", "openrouter", "groq"}:
+        return OpenAICompatibleProvider(cfg)
+    raise LLMConfigurationError(
+        f"unsupported Graphite LLM provider '{cfg.llm_provider}'. Use 'ollama' or 'openai-compatible'."
+    )
+
+
+def canonical_provider_name(provider: str) -> Literal[
+    "ollama", "openai-compatible", "custom/unknown"
+]:
+    """Return a fixed public identifier without reflecting configuration input."""
+    normalized = provider.strip().lower().replace("_", "-")
+    if normalized in {"ollama", "local"}:
+        return "ollama"
+    if normalized in {
+        "openai",
+        "openai-compatible",
+        "compatible",
+        "lmstudio",
+        "lm-studio",
+        "vllm",
+        "openrouter",
+        "groq",
+    }:
+        return "openai-compatible"
+    return "custom/unknown"
+
+
+def _provider_requires_api_key(provider: str) -> bool:
+    normalized = provider.strip().lower().replace("_", "-")
+    return normalized in {"openai", "openrouter", "groq"}
+
+
+def _provider_requires_secure_egress(provider: str) -> bool:
+    normalized = provider.strip().lower().replace("_", "-")
+    return normalized in {"openai", "openrouter", "groq"}
+
+
+# RFC 6598 carrier-grade NAT / shared address space. Not classified as private by
+# is_private on every CPython patch version, so reject it explicitly rather than relying
+# on the stdlib's version-dependent classification.
+_CGNAT_SHARED_IPV4 = ipaddress.ip_network("100.64.0.0/10")
+
+
+def _validate_llm_base_url(base_url: str, *, provider: str) -> None:
+    """Apply the provider-class egress policy to an operator-supplied base URL.
+
+    Every provider rejects non-HTTP(S) schemes. The keyed cloud providers
+    (openai/openrouter/groq) additionally require HTTPS and reject IP-literal
+    loopback/private/link-local/reserved hosts and the name 'localhost' -- those
+    providers have no legitimate internal target, so an internal target there is a
+    misconfiguration that would leak the bearer token. Local and generic
+    openai-compatible providers keep http + loopback/private, which are intended.
+    DNS-name hosts are accepted; name-based private targets are out of scope
+    because base_url is operator-controlled (no DNS-rebinding adversary).
+    """
+    parts = urlsplit(base_url)
+    scheme = parts.scheme.lower()
+    if scheme not in {"http", "https"}:
+        raise LLMConfigurationError("GRAPHITE_LLM_BASE_URL must use http or https")
+
+    if not _provider_requires_secure_egress(provider):
+        return
+
+    if scheme != "https":
+        raise LLMConfigurationError("this LLM provider requires an https GRAPHITE_LLM_BASE_URL")
+    host = parts.hostname or ""
+    if host.lower() == "localhost":
+        raise LLMConfigurationError("this LLM provider may not target a loopback or private host")
+    try:
+        address = ipaddress.ip_address(host)
+    except ValueError:
+        return  # A DNS name; name-based private targets are out of scope.
+    if (
+        address.is_loopback
+        or address.is_private
+        or address.is_link_local
+        or address.is_reserved
+        or address.is_multicast
+        or address.is_unspecified
+        or (address.version == 4 and address in _CGNAT_SHARED_IPV4)
+    ):
+        raise LLMConfigurationError("this LLM provider may not target a loopback or private host")
+
+
+def _provider_has_default_base_url(provider: str) -> bool:
+    normalized = provider.strip().lower().replace("_", "-")
+    return normalized in {"openai", "openrouter", "groq", "lmstudio", "lm-studio", "vllm"}
+
+
+def _auto_effective_mode(provider: str) -> str:
+    normalized = provider.strip().lower().replace("_", "-")
+    if normalized in {"ollama", "local", "lmstudio", "lm-studio", "vllm"}:
+        return "local"
+    return "cloud"
+
+
+def decide_auto_llm(
+    graph_data: dict[str, Any],
+    clusters: dict[str, Any],
+    analysis: dict[str, Any],
+    cfg: Config,
+) -> dict[str, Any]:
+    """Decide whether optional LLM enrichment is worth running for this graph."""
+    metadata = graph_data.get("metadata") or {}
+    node_count = int(metadata.get("node_count") or len(graph_data.get("nodes") or []))
+    edge_count = int(metadata.get("edge_count") or len(graph_data.get("edges") or []))
+    cluster_count = int(clusters.get("count") or len(clusters.get("clusters") or []))
+    god_node_count = len(analysis.get("god_nodes") or [])
+    surprising_count = len(analysis.get("surprising_connections") or [])
+    top_link_count = len(analysis.get("top_files_by_links") or [])
+
+    signals: list[str] = []
+    if node_count >= 250:
+        signals.append("node_count>=250")
+    if edge_count >= 350:
+        signals.append("edge_count>=350")
+    if cluster_count >= 8:
+        signals.append("cluster_count>=8")
+    if god_node_count >= 2:
+        signals.append("god_nodes>=2")
+    if surprising_count >= 2:
+        signals.append("surprising_connections>=2")
+    if top_link_count >= 10 and edge_count >= 200:
+        signals.append("high_linked_files")
+
+    provider = cfg.llm_provider.strip().lower().replace("_", "-")
+    base_url_ready = bool(cfg.llm_base_url) or _provider_has_default_base_url(provider) or provider in {"ollama", "local"}
+    api_key_ready = bool(cfg.llm_api_key) or not _provider_requires_api_key(provider)
+    provider_ready = base_url_ready and api_key_ready
+
+    decision = {
+        "enabled": bool(signals) and provider_ready,
+        "signals": signals,
+        "provider": cfg.llm_provider,
+        "model": cfg.llm_model or _default_model(cfg.llm_provider),
+        "effective_mode": _auto_effective_mode(cfg.llm_provider),
+        "node_count": node_count,
+        "edge_count": edge_count,
+        "cluster_count": cluster_count,
+        "god_node_count": god_node_count,
+        "surprising_connection_count": surprising_count,
+        "provider_ready": provider_ready,
+    }
+    if not signals:
+        decision["reason"] = "graph_below_auto_threshold"
+    elif not base_url_ready:
+        decision["reason"] = "missing_llm_base_url"
+    elif not api_key_ready:
+        decision["reason"] = "missing_llm_api_key"
+    else:
+        decision["reason"] = "auto_threshold_met"
+    return decision
+
+def enrich_report(
+    graph_data: dict[str, Any],
+    clusters: dict[str, Any],
+    analysis: dict[str, Any],
+    cfg: Config,
+) -> dict[str, Any]:
+    """Return an optional annotation payload for the explicit overlay pipeline."""
+    mode = cfg.llm_mode.strip().lower()
+    if mode in {"", "0", "false", "none", "off", "disabled"}:
+        return {"enabled": False, "mode": "none", "tokens": 0}
+
+    auto_decision: dict[str, Any] | None = None
+    effective_cfg = cfg
+    if mode == "auto":
+        auto_decision = decide_auto_llm(graph_data, clusters, analysis, cfg)
+        if not auto_decision["enabled"]:
+            return {
+                "enabled": False,
+                "status": "skipped",
+                "mode": "auto",
+                "tokens": 0,
+                "provider": cfg.llm_provider,
+                "model": auto_decision["model"],
+                "reason": auto_decision["reason"],
+                "auto": auto_decision,
+            }
+        effective_cfg = replace(cfg, llm_mode=str(auto_decision["effective_mode"]))
+
+    try:
+        provider = make_provider(effective_cfg)
+        system, user = build_report_prompt(graph_data, clusters, analysis, effective_cfg)
+        completion = provider.complete(system, user)
+        return {
+            "enabled": True,
+            "status": "ok",
+            "mode": mode,
+            "effective_mode": effective_cfg.llm_mode,
+            "provider": provider.name,
+            "model": cfg.llm_model or _default_model(cfg.llm_provider),
+            "summary": completion.text,
+            "auto": auto_decision,
+            "input_tokens": completion.input_tokens,
+            "output_tokens": completion.output_tokens,
+            "tokens": completion.total_tokens or completion.input_tokens + completion.output_tokens,
+        }
+    except Exception as exc:  # Overlay failures remain isolated from canonical graphs.
+        if isinstance(exc, LLMConfigurationError):
+            category: ProviderErrorCategory = "configuration"
+        elif isinstance(exc, LLMProviderError):
+            category = exc.category
+        elif isinstance(exc, TimeoutError):
+            category = "timeout"
+        else:
+            category = "provider_error"
+        return {
+            "enabled": True,
+            "status": "error",
+            "mode": mode,
+            "effective_mode": effective_cfg.llm_mode,
+            "provider": cfg.llm_provider,
+            "model": cfg.llm_model or _default_model(cfg.llm_provider),
+            "auto": auto_decision,
+            "tokens": 0,
+            "error_category": category,
+        }
+
+
+def build_report_prompt(
+    graph_data: dict[str, Any],
+    clusters: dict[str, Any],
+    analysis: dict[str, Any],
+    cfg: Config,
+) -> tuple[str, str]:
+    """Build a bounded prompt from graph metadata only, not source code."""
+    payload = {
+        "metadata": graph_data.get("metadata", {}),
+        "top_files_by_links": analysis.get("top_files_by_links", [])[:12],
+        "god_nodes": analysis.get("god_nodes", [])[:12],
+        "entry_points": analysis.get("entry_points", [])[:12],
+        "surprising_connections": analysis.get("surprising_connections", [])[:12],
+        "clusters": [
+            {
+                "id": c.get("id"),
+                "size": c.get("size"),
+                "file_count": c.get("file_count"),
+                "function_count": c.get("function_count"),
+                "class_count": c.get("class_count"),
+                "labels": c.get("labels", []),
+            }
+            for c in clusters.get("clusters", [])[:12]
+        ],
+    }
+    data = json.dumps(payload, ensure_ascii=False, indent=2)
+    if len(data) > cfg.llm_max_input_chars:
+        data = data[: cfg.llm_max_input_chars] + "\n... truncated by GRAPHITE_LLM_MAX_INPUT_CHARS"
+    system = (
+        "You are Graphite's optional code-graph analyst. Use only the provided graph metrics. "
+        "Do not claim to have inspected source code. Focus on architecture risks, likely hotspots, "
+        "and practical next checks. Keep the answer concise."
+    )
+    user = (
+        "Analyze this code graph summary. Return concise bullets under: Hotspots, Coupling risks, "
+        "Suggested next checks.\n\n"
+        f"{data}"
+    )
+    return system, user
+
+
+def _urlopen_json(request: urllib.request.Request, timeout: float) -> dict[str, Any]:
+    try:
+        with _open_no_redirect(request, timeout) as response:  # nosec: opt-in user URL
+            raw = response.read(MAX_RESPONSE_BYTES + 1)
+    except urllib.error.HTTPError as exc:
+        try:
+            exc.read(MAX_RESPONSE_BYTES + 1)
+        except Exception:
+            pass
+        category: ProviderErrorCategory = (
+            "authentication" if exc.code in {401, 403} else "provider_error"
+        )
+        raise LLMProviderError(category) from None
+    except (TimeoutError, urllib.error.URLError) as exc:
+        reason = getattr(exc, "reason", None)
+        category = "timeout" if isinstance(exc, TimeoutError) or isinstance(reason, TimeoutError) else "connection"
+        raise LLMProviderError(category) from None
+    except OSError:
+        raise LLMProviderError("connection") from None
+    if len(raw) > MAX_RESPONSE_BYTES:
+        raise LLMProviderError("provider_error")
+    try:
+        payload = json.loads(raw.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        raise LLMProviderError("provider_error") from None
+    if not isinstance(payload, dict):
+        raise LLMProviderError("provider_error")
+    return payload
+
+
+def _default_openai_base_url(provider: str) -> str | None:
+    normalized = provider.strip().lower().replace("_", "-")
+    if normalized == "openai":
+        return "https://api.openai.com/v1"
+    if normalized in {"lmstudio", "lm-studio"}:
+        return "http://localhost:1234/v1"
+    if normalized == "vllm":
+        return "http://localhost:8000/v1"
+    if normalized == "openrouter":
+        return "https://openrouter.ai/api/v1"
+    if normalized == "groq":
+        return "https://api.groq.com/openai/v1"
+    return None
+
+
+def _default_model(provider: str) -> str:
+    normalized = provider.strip().lower().replace("_", "-")
+    if normalized in {"ollama", "local"}:
+        return "llama3.1"
+    if normalized == "openrouter":
+        return "moonshotai/kimi-k2.7-code"
+    if normalized == "groq":
+        return "llama-3.1-8b-instant"
+    return "gpt-4o-mini"