polycodegraph 0.1.0__py3-none-any.whl

codegraph/config.py

@@ -0,0 +1,64 @@
+"""Codegraph configuration model and helpers."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, Field
+
+
+class DeadCodeConfig(BaseModel):
+    """User-supplied dead-code analysis tweaks.
+
+    Extends the built-in entry-point catalog. All fields are optional;
+    user patterns are merged with the built-ins at parse time.
+    """
+
+    entry_point_decorators: list[str] = Field(default_factory=list)
+    """Extra decorator strings (e.g. ``"@my.handler"``) treated as entry
+    points. Matched as substring of the raw decorator text."""
+
+    entry_point_names: list[str] = Field(default_factory=list)
+    """Extra function/method/class name globs treated as entry points
+    (fnmatch syntax). Reserved for future use."""
+
+    entry_point_files: list[str] = Field(default_factory=list)
+    """File-path globs whose definitions are all treated as entry points.
+    Reserved for future use."""
+
+
+class CodegraphConfig(BaseModel):
+    version: int = 1
+    languages: list[str] = Field(
+        default_factory=lambda: ["python", "typescript", "javascript"]
+    )
+    default_branch: str = "main"
+    ignore: list[str] = Field(default_factory=list)
+    baseline: dict[str, Any] = Field(
+        default_factory=lambda: {"backend": "local"}
+    )
+    critical_paths: list[dict[str, Any]] = Field(default_factory=list)
+    mcp: dict[str, Any] = Field(default_factory=lambda: {"enabled": False})
+    install_hook: bool = False
+    register_mcp: bool = False
+    dead_code: DeadCodeConfig = Field(default_factory=DeadCodeConfig)
+
+
+def load_config(repo_root: Path) -> CodegraphConfig:
+    cfg_path = repo_root / ".codegraph.yml"
+    if not cfg_path.exists():
+        return CodegraphConfig()
+    with cfg_path.open() as f:
+        data = yaml.safe_load(f) or {}
+    return CodegraphConfig.model_validate(data)
+
+
+def save_config(repo_root: Path, cfg: CodegraphConfig) -> None:
+    cfg_path = repo_root / ".codegraph.yml"
+    with cfg_path.open("w") as f:
+        yaml.dump(cfg.model_dump(), f, default_flow_style=False, sort_keys=True)
+
+
+def default_data_dir(repo_root: Path) -> Path:
+    return repo_root / ".codegraph"

codegraph/embed/__init__.py

@@ -0,0 +1,35 @@
+"""Local, open-weight embedding layer for codegraph (v0.3).
+
+Public surface:
+
+* :class:`Chunk` — a single embeddable code unit (function/method/class).
+* :class:`Hit` — a search result.
+* :func:`chunk_repo` — turn graph nodes into chunks.
+* :func:`build_index` — chunk + embed + write index to ``.codegraph/embeddings.lance``.
+* :func:`query` — convenience wrapper around :func:`semantic_query`.
+
+The heavy dependencies (``sentence-transformers``, ``lancedb``) are imported
+lazily.  Install with ``pip install -e ".[embed]"``.
+"""
+from __future__ import annotations
+
+from codegraph.embed.chunker import Chunk, chunk_repo
+from codegraph.embed.query import Hit, hybrid_query, semantic_query
+from codegraph.embed.store import EmbeddingStore, build_index
+
+__all__ = [
+    "Chunk",
+    "EmbeddingStore",
+    "Hit",
+    "build_index",
+    "chunk_repo",
+    "hybrid_query",
+    "query",
+    "semantic_query",
+]
+
+
+# pragma: codegraph-public-api
+def query(text: str, *, k: int = 5) -> list[Hit]:
+    """Shortcut: run a semantic query against the cwd ``.codegraph`` index."""
+    return semantic_query(text, k=k)

codegraph/embed/chunker.py

@@ -0,0 +1,120 @@
+"""Chunk a codegraph SQLite store into embeddable units."""
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from codegraph.graph.schema import Node, NodeKind
+from codegraph.graph.store_sqlite import SQLiteGraphStore
+
+# Node kinds that produce a chunk by default.  MODULE / FILE / TEST / IMPORT /
+# VARIABLE / PARAMETER are skipped: they're either coarse, generated, or
+# duplicate the chunks of the symbols they contain.
+_DEFAULT_KINDS: frozenset[NodeKind] = frozenset(
+    {NodeKind.FUNCTION, NodeKind.METHOD, NodeKind.CLASS}
+)
+
+
+@dataclass(frozen=True)
+class Chunk:
+    """A single embeddable code chunk pulled from the graph + source."""
+
+    qualname: str
+    file: str
+    line_start: int
+    line_end: int
+    kind: str
+    text: str
+    params: list[str] = field(default_factory=list)
+    returns: str | None = None
+    role: str | None = None
+
+    @property
+    def id(self) -> str:
+        """Stable id for upsert / dedupe."""
+        return f"{self.file}::{self.qualname}::{self.line_start}"
+
+
+def _read_lines(repo_root: Path, file: str) -> list[str]:
+    """Best-effort read; returns ``[]`` on any IO error."""
+    try:
+        return (repo_root / file).read_text(encoding="utf-8").splitlines()
+    except (OSError, UnicodeDecodeError):
+        return []
+
+
+def _slice(lines: list[str], start: int, end: int) -> str:
+    if not lines:
+        return ""
+    s = max(0, start - 1)
+    e = max(s, min(len(lines), end))
+    return "\n".join(lines[s:e])
+
+
+def _md_list(value: Any) -> list[str]:
+    if isinstance(value, list):
+        return [str(v) for v in value]
+    return []
+
+
+def _node_to_chunk(node: Node, repo_root: Path, _cache: dict[str, list[str]]) -> Chunk | None:
+    lines = _cache.get(node.file)
+    if lines is None:
+        lines = _read_lines(repo_root, node.file)
+        _cache[node.file] = lines
+
+    body = _slice(lines, node.line_start, node.line_end)
+    if not body.strip():
+        # Fall back to signature / docstring so we still index something useful.
+        body = "\n".join(filter(None, [node.signature or "", node.docstring or ""]))
+    if not body.strip():
+        return None
+
+    md = node.metadata or {}
+    role_val = md.get("role") if isinstance(md, dict) else None
+    returns_val = md.get("returns") if isinstance(md, dict) else None
+
+    return Chunk(
+        qualname=node.qualname,
+        file=node.file,
+        line_start=node.line_start,
+        line_end=node.line_end,
+        kind=node.kind.value,
+        text=body,
+        params=_md_list(md.get("params")) if isinstance(md, dict) else [],
+        returns=str(returns_val) if returns_val is not None else None,
+        role=str(role_val) if role_val is not None else None,
+    )
+
+
+def chunk_repo(
+    repo_root: Path,
+    *,
+    db_path: Path | None = None,
+    kinds: Iterable[NodeKind] | None = None,
+) -> Iterator[Chunk]:
+    """Yield one :class:`Chunk` per matching graph node.
+
+    ``kinds`` defaults to FUNCTION / METHOD / CLASS.
+    """
+    db = db_path or (repo_root / ".codegraph" / "graph.db")
+    if not db.exists():
+        raise FileNotFoundError(
+            f"No graph database at {db}. Run `codegraph build` first."
+        )
+
+    selected = frozenset(kinds) if kinds is not None else _DEFAULT_KINDS
+    line_cache: dict[str, list[str]] = {}
+
+    store = SQLiteGraphStore(db)
+    try:
+        for node in store.iter_nodes():
+            if node.kind not in selected:
+                continue
+            chunk = _node_to_chunk(node, repo_root, line_cache)
+            if chunk is not None:
+                yield chunk
+    finally:
+        store.close()

codegraph/embed/embedder.py

@@ -0,0 +1,113 @@
+"""Sentence-transformers wrapper.
+
+The model is lazy-loaded so unit tests can substitute a deterministic fake
+without pulling the real dependency tree.
+"""
+from __future__ import annotations
+
+import os
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, Protocol
+
+DEFAULT_MODEL = "nomic-ai/CodeRankEmbed"
+DEFAULT_DIM = 768
+
+
+class _EncoderLike(Protocol):
+    def encode(
+        self,
+        sentences: Sequence[str],
+        *,
+        batch_size: int = ...,
+        show_progress_bar: bool = ...,
+        convert_to_numpy: bool = ...,
+        normalize_embeddings: bool = ...,
+    ) -> Any:
+        ...
+
+
+def _cache_dir() -> Path:
+    """Where downloaded models live.  Honours ``XDG_CACHE_HOME`` if set."""
+    xdg = os.environ.get("XDG_CACHE_HOME")
+    base = Path(xdg) if xdg else Path.home() / ".cache"
+    return base / "codegraph" / "models"
+
+
+class MissingDependencyError(RuntimeError):
+    """Raised when the optional ``embed`` extra is not installed."""
+
+
+def _load_sentence_transformer(model: str) -> _EncoderLike:
+    try:
+        from sentence_transformers import SentenceTransformer
+    except ImportError as exc:  # pragma: no cover — exercised via mock test
+        raise MissingDependencyError(
+            "sentence-transformers is not installed.\n"
+            "Run: pip install -e \".[embed]\""
+        ) from exc
+
+    cache = _cache_dir()
+    cache.mkdir(parents=True, exist_ok=True)
+    encoder: _EncoderLike = SentenceTransformer(
+        model, cache_folder=str(cache), trust_remote_code=True
+    )
+    return encoder
+
+
+class Embedder:
+    """Tiny wrapper around a sentence-transformer model.
+
+    The actual encoder is loaded lazily on first :meth:`embed` call so that:
+
+    * Construction is cheap and side-effect free (good for tests).
+    * Multiple embedders can co-exist without re-downloading.
+
+    Pass ``encoder`` to inject a fake / mock for tests.
+    """
+
+    def __init__(
+        self,
+        model: str = DEFAULT_MODEL,
+        *,
+        dim: int | None = None,
+        encoder: _EncoderLike | None = None,
+    ) -> None:
+        self.model = model
+        self._encoder: _EncoderLike | None = encoder
+        self._dim: int | None = dim
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+    def _ensure_loaded(self) -> _EncoderLike:
+        if self._encoder is None:
+            self._encoder = _load_sentence_transformer(self.model)
+        return self._encoder
+
+    @property
+    def dim(self) -> int:
+        if self._dim is not None:
+            return self._dim
+        # Probe the encoder with a single token so we don't have to assume.
+        vecs = self.embed(["probe"])
+        self._dim = len(vecs[0])
+        return self._dim
+
+    # ------------------------------------------------------------------
+    # Encoding
+    # ------------------------------------------------------------------
+    def embed(self, texts: Sequence[str], *, batch_size: int = 32) -> list[list[float]]:
+        """Return one row of floats per input string."""
+        if not texts:
+            return []
+        enc = self._ensure_loaded()
+        out = enc.encode(
+            list(texts),
+            batch_size=batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+        )
+        # Accept numpy arrays, lists, or anything that iterates rows.
+        return [list(map(float, row)) for row in out]

codegraph/embed/query.py

@@ -0,0 +1,181 @@
+"""Query helpers for the embeddings index."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from codegraph.embed.embedder import Embedder
+from codegraph.embed.store import EmbeddingStore, StoredChunk
+
+
+@dataclass(frozen=True)
+class Hit:
+    """A single search result."""
+
+    qualname: str
+    file: str
+    line: int
+    kind: str
+    role: str | None
+    score: float
+    text_snippet: str
+
+    # pragma: codegraph-public-api
+    def as_dict(self, *, score_field: str = "score") -> dict[str, Any]:
+        return {
+            "qualname": self.qualname,
+            "file": self.file,
+            "line": self.line,
+            "kind": self.kind,
+            "role": self.role,
+            score_field: round(float(self.score), 6),
+            "text_snippet": self.text_snippet,
+        }
+
+
+def _snippet(text: str, max_lines: int = 6, max_chars: int = 400) -> str:
+    lines = text.splitlines()[:max_lines]
+    snippet = "\n".join(lines)
+    if len(snippet) > max_chars:
+        snippet = snippet[:max_chars] + "…"
+    return snippet
+
+
+def _index_dir(repo_root: Path | None = None) -> Path:
+    return (repo_root or Path.cwd()) / ".codegraph"
+
+
+class IndexMissingError(RuntimeError):
+    """Raised when the embeddings index does not exist on disk."""
+
+
+def _open_store(repo_root: Path | None = None) -> EmbeddingStore:
+    base = _index_dir(repo_root)
+    lance = base / "embeddings.lance"
+    json_fb = base / "embeddings.json"
+    if not lance.exists() and not json_fb.exists():
+        raise IndexMissingError(
+            "no embedding index — run `codegraph embed` first"
+        )
+    backend = "auto" if lance.exists() else "json"
+    return EmbeddingStore(base, backend=backend)
+
+
+def _to_hit(row: StoredChunk, score: float) -> Hit:
+    return Hit(
+        qualname=row.qualname,
+        file=row.file,
+        line=row.line_start,
+        kind=row.kind,
+        role=row.role,
+        score=score,
+        text_snippet=_snippet(row.text),
+    )
+
+
+def semantic_query(
+    text: str,
+    *,
+    k: int = 5,
+    repo_root: Path | None = None,
+    embedder: Embedder | None = None,
+    store: EmbeddingStore | None = None,
+) -> list[Hit]:
+    """Pure cosine-similarity ranking against the index."""
+    s = store or _open_store(repo_root)
+    emb = embedder or Embedder()
+    vector = emb.embed([text])[0]
+    hits = s.query(vector, k=k)
+    return [_to_hit(row, score) for row, score in hits]
+
+
+def hybrid_query(
+    text: str,
+    *,
+    k: int = 5,
+    role: str | None = None,
+    focus_qn: str | None = None,
+    repo_root: Path | None = None,
+    embedder: Embedder | None = None,
+    store: EmbeddingStore | None = None,
+    graph: Any | None = None,
+) -> list[Hit]:
+    """Blend semantic similarity with graph distance from a focus node.
+
+    ``final_score = 0.6 * cosine + 0.4 * (1 / (1 + graph_hops))``
+    """
+    s = store or _open_store(repo_root)
+    emb = embedder or Embedder()
+    vector = emb.embed([text])[0]
+
+    pool_size = max(k * 4, 20)
+    raw = s.query(vector, k=pool_size)
+
+    if role is not None:
+        raw = [(row, score) for row, score in raw if row.role == role]
+
+    if focus_qn is None:
+        return [_to_hit(row, score) for row, score in raw[:k]]
+
+    g = graph if graph is not None else _load_graph(repo_root)
+    focus_id = _find_node_by_qualname(g, focus_qn) if g is not None else None
+
+    rescored: list[tuple[StoredChunk, float, float]] = []
+    for row, semantic in raw:
+        target_id = _find_node_by_qualname(g, row.qualname) if g is not None else None
+        hops = _graph_distance(g, focus_id, target_id) if g is not None else None
+        graph_score = 0.0 if hops is None else 1.0 / (1.0 + float(hops))
+        final = 0.6 * float(semantic) + 0.4 * graph_score
+        rescored.append((row, final, semantic))
+
+    rescored.sort(key=lambda triple: triple[1], reverse=True)
+    return [_to_hit(row, final) for row, final, _ in rescored[:k]]
+
+
+# ---------------------------------------------------------------------------
+# Graph helpers (lazy nx import; tolerate missing graph)
+# ---------------------------------------------------------------------------
+
+
+def _load_graph(repo_root: Path | None) -> Any | None:
+    try:
+        from codegraph.graph.store_networkx import to_digraph
+        from codegraph.graph.store_sqlite import SQLiteGraphStore
+    except Exception:  # pragma: no cover
+        return None
+
+    db = (repo_root or Path.cwd()) / ".codegraph" / "graph.db"
+    if not db.exists():
+        return None
+    store = SQLiteGraphStore(db)
+    try:
+        return to_digraph(store)
+    finally:
+        store.close()
+
+
+def _find_node_by_qualname(graph: Any, qualname: str) -> str | None:
+    if graph is None:
+        return None
+    if qualname in graph:
+        return qualname
+    q = qualname.lower()
+    for nid, attrs in graph.nodes(data=True):
+        if str(attrs.get("qualname") or "").lower() == q:
+            return str(nid)
+    return None
+
+
+def _graph_distance(graph: Any, src: str | None, dst: str | None) -> int | None:
+    if graph is None or src is None or dst is None:
+        return None
+    if src == dst:
+        return 0
+    try:
+        import networkx as nx
+
+        ug = graph.to_undirected(as_view=True) if hasattr(graph, "to_undirected") else graph
+        return int(nx.shortest_path_length(ug, src, dst))
+    except Exception:
+        return None