polycodegraph 0.1.0__py3-none-any.whl

codegraph/embed/store.py

@@ -0,0 +1,360 @@
+"""On-disk vector store.
+
+Tries LanceDB first (the production backend) and falls back to a tiny JSON
+file when the optional ``embed`` extra isn't installed.  The fallback is good
+enough for unit tests and for repos that just want a quick local index without
+pulling the full Arrow / LanceDB stack.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import math
+from collections.abc import Iterable, Sequence
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+from codegraph.embed.chunker import Chunk  # noqa: F401  (re-export friendly)
+from codegraph.embed.embedder import DEFAULT_DIM, DEFAULT_MODEL, Embedder
+
+_STORE_FILENAME = "embeddings.lance"
+_FALLBACK_FILENAME = "embeddings.json"
+
+
+@dataclass
+class StoredChunk:
+    id: str
+    qualname: str
+    file: str
+    line_start: int
+    line_end: int
+    kind: str
+    role: str | None
+    text: str
+    vector: list[float] = field(default_factory=list)
+
+    # pragma: codegraph-public-api
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "qualname": self.qualname,
+            "file": self.file,
+            "line_start": self.line_start,
+            "line_end": self.line_end,
+            "kind": self.kind,
+            "role": self.role,
+            "text": self.text,
+            "vector": list(self.vector),
+        }
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> StoredChunk:
+        return cls(
+            id=str(data["id"]),
+            qualname=str(data["qualname"]),
+            file=str(data["file"]),
+            line_start=int(data["line_start"]),
+            line_end=int(data["line_end"]),
+            kind=str(data["kind"]),
+            role=(str(data["role"]) if data.get("role") else None),
+            text=str(data["text"]),
+            vector=[float(v) for v in data.get("vector") or []],
+        )
+
+
+def _cosine(a: Sequence[float], b: Sequence[float]) -> float:
+    if not a or not b:
+        return 0.0
+    dot = 0.0
+    na = 0.0
+    nb = 0.0
+    for x, y in zip(a, b, strict=False):
+        dot += x * y
+        na += x * x
+        nb += y * y
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    return dot / (math.sqrt(na) * math.sqrt(nb))
+
+
+# ---------------------------------------------------------------------------
+# Backend abstraction
+# ---------------------------------------------------------------------------
+
+class _JsonBackend:
+    """JSON-backed backend.  Used in tests and as the no-deps fallback."""
+
+    def __init__(self, path: Path) -> None:
+        self.path = path
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._rows: list[StoredChunk] = []
+        if self.path.exists():
+            try:
+                raw = json.loads(self.path.read_text(encoding="utf-8"))
+                self._rows = [StoredChunk.from_json(r) for r in raw]
+            except (OSError, json.JSONDecodeError, KeyError, TypeError):
+                self._rows = []
+
+    def upsert(self, rows: Iterable[StoredChunk]) -> None:
+        new = list(rows)
+        new_ids = {r.id for r in new}
+        kept = [r for r in self._rows if r.id not in new_ids]
+        self._rows = kept + new
+        self._flush()
+
+    def replace_all(self, rows: Iterable[StoredChunk]) -> None:
+        self._rows = list(rows)
+        self._flush()
+
+    def _flush(self) -> None:
+        payload = [r.to_json() for r in self._rows]
+        self.path.write_text(json.dumps(payload), encoding="utf-8")
+
+    def all(self) -> list[StoredChunk]:
+        return list(self._rows)
+
+    def query(self, vector: Sequence[float], k: int) -> list[tuple[StoredChunk, float]]:
+        scored = [(row, _cosine(vector, row.vector)) for row in self._rows]
+        scored.sort(key=lambda pair: pair[1], reverse=True)
+        return scored[:k]
+
+    def size_bytes(self) -> int:
+        return self.path.stat().st_size if self.path.exists() else 0
+
+
+class _LanceBackend:
+    """LanceDB backend.  Schema mirrors :class:`StoredChunk`."""
+
+    def __init__(self, path: Path, dim: int) -> None:
+        import lancedb
+        import pyarrow as pa
+
+        self.path = path
+        self._dim = dim
+        self._pa = pa
+        self._db = lancedb.connect(str(path))
+        self._schema = self._make_schema(dim)
+        if "chunks" in self._db.table_names():
+            self._table = self._db.open_table("chunks")
+        else:
+            self._table = self._db.create_table("chunks", schema=self._schema, mode="create")
+
+    def _make_schema(self, dim: int) -> Any:
+        pa = self._pa
+        return pa.schema(
+            [
+                ("id", pa.string()),
+                ("qualname", pa.string()),
+                ("file", pa.string()),
+                ("line_start", pa.int64()),
+                ("line_end", pa.int64()),
+                ("kind", pa.string()),
+                ("role", pa.string()),
+                ("text", pa.string()),
+                ("vector", pa.list_(pa.float32(), dim)),
+            ]
+        )
+
+    def _to_dict(self, row: StoredChunk) -> dict[str, Any]:
+        return {
+            "id": row.id,
+            "qualname": row.qualname,
+            "file": row.file,
+            "line_start": row.line_start,
+            "line_end": row.line_end,
+            "kind": row.kind,
+            "role": row.role or "",
+            "text": row.text,
+            "vector": row.vector,
+        }
+
+    def upsert(self, rows: Iterable[StoredChunk]) -> None:
+        batch = [self._to_dict(r) for r in rows]
+        if not batch:
+            return
+        ids = ", ".join(f"'{r['id']}'" for r in batch)
+        with contextlib.suppress(Exception):
+            self._table.delete(f"id IN ({ids})")
+        self._table.add(batch)
+
+    def replace_all(self, rows: Iterable[StoredChunk]) -> None:
+        batch = [self._to_dict(r) for r in rows]
+        with contextlib.suppress(Exception):
+            self._db.drop_table("chunks", ignore_missing=True)
+        self._table = self._db.create_table("chunks", schema=self._schema, mode="create")
+        if batch:
+            self._table.add(batch)
+
+    def _row_from_record(self, r: dict[str, Any]) -> StoredChunk:
+        return StoredChunk(
+            id=str(r["id"]),
+            qualname=str(r["qualname"]),
+            file=str(r["file"]),
+            line_start=int(r["line_start"]),
+            line_end=int(r["line_end"]),
+            kind=str(r["kind"]),
+            role=str(r["role"]) or None,
+            text=str(r["text"]),
+            vector=list(r["vector"]),
+        )
+
+    def all(self) -> list[StoredChunk]:
+        rows = self._table.to_pandas().to_dict(orient="records")
+        return [self._row_from_record(r) for r in rows]
+
+    def query(self, vector: Sequence[float], k: int) -> list[tuple[StoredChunk, float]]:
+        results = self._table.search(list(vector)).limit(k).to_pandas()
+        out: list[tuple[StoredChunk, float]] = []
+        for r in results.to_dict(orient="records"):
+            chunk = self._row_from_record(r)
+            distance = float(r.get("_distance", 0.0))
+            similarity = 1.0 / (1.0 + distance)
+            out.append((chunk, similarity))
+        return out
+
+    def size_bytes(self) -> int:
+        total = 0
+        for p in self.path.rglob("*"):
+            if p.is_file():
+                total += p.stat().st_size
+        return total
+
+
+# ---------------------------------------------------------------------------
+# Public store
+# ---------------------------------------------------------------------------
+
+
+class EmbeddingStore:
+    """High-level interface that auto-selects a backend.
+
+    ``backend='auto'`` (default) tries LanceDB and falls back to JSON.
+    ``backend='json'`` forces the lightweight backend (used in tests).
+    """
+
+    def __init__(
+        self,
+        data_dir: Path,
+        *,
+        dim: int = DEFAULT_DIM,
+        backend: str = "auto",
+    ) -> None:
+        self.data_dir = data_dir
+        self.dim = dim
+        self.backend_name: str
+        self._backend: _LanceBackend | _JsonBackend
+        data_dir.mkdir(parents=True, exist_ok=True)
+
+        if backend == "json":
+            self._backend = _JsonBackend(data_dir / _FALLBACK_FILENAME)
+            self.backend_name = "json"
+            return
+
+        try:
+            self._backend = _LanceBackend(data_dir / _STORE_FILENAME, dim=dim)
+            self.backend_name = "lancedb"
+        except ImportError:
+            if backend == "lancedb":
+                raise
+            self._backend = _JsonBackend(data_dir / _FALLBACK_FILENAME)
+            self.backend_name = "json"
+
+    # ------------------------------------------------------------------
+    # pragma: codegraph-public-api
+    def upsert(self, rows: Iterable[StoredChunk]) -> None:
+        self._backend.upsert(rows)
+
+    # pragma: codegraph-public-api
+    def replace_all(self, rows: Iterable[StoredChunk]) -> None:
+        self._backend.replace_all(rows)
+
+    # pragma: codegraph-public-api
+    def all(self) -> list[StoredChunk]:
+        return self._backend.all()
+
+    # pragma: codegraph-public-api
+    def query(self, vector: Sequence[float], k: int = 5) -> list[tuple[StoredChunk, float]]:
+        return self._backend.query(vector, k)
+
+    # pragma: codegraph-public-api
+    def size_bytes(self) -> int:
+        return self._backend.size_bytes()
+
+
+# ---------------------------------------------------------------------------
+# build_index — orchestrator wired up to chunker + embedder
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class IndexStats:
+    chunks_indexed: int
+    model: str
+    dim: int
+    backend: str
+    on_disk_bytes: int
+
+    # pragma: codegraph-public-api
+    def as_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def build_index(
+    repo_root: Path,
+    *,
+    db_path: Path | None = None,
+    embeddings_dir: Path | None = None,
+    embedder: Embedder | None = None,
+    model: str = DEFAULT_MODEL,
+    force: bool = False,
+    progress: Any | None = None,
+    backend: str = "auto",
+) -> IndexStats:
+    """Chunk + embed + persist.
+
+    ``progress`` (optional) is anything with an ``advance(step: int)`` method
+    — typically a ``rich.progress.Progress`` task.  Pass ``None`` to disable.
+    """
+    from codegraph.embed.chunker import chunk_repo
+
+    chunks = list(chunk_repo(repo_root, db_path=db_path))
+    emb = embedder or Embedder(model)
+
+    rows: list[StoredChunk] = []
+    dim = DEFAULT_DIM
+    if chunks:
+        vectors = emb.embed([c.text for c in chunks], batch_size=32)
+        dim = len(vectors[0]) if vectors else DEFAULT_DIM
+        for c, v in zip(chunks, vectors, strict=False):
+            rows.append(
+                StoredChunk(
+                    id=c.id,
+                    qualname=c.qualname,
+                    file=c.file,
+                    line_start=c.line_start,
+                    line_end=c.line_end,
+                    kind=c.kind,
+                    role=c.role,
+                    text=c.text,
+                    vector=v,
+                )
+            )
+            if progress is not None:
+                with contextlib.suppress(Exception):
+                    progress.advance(1)
+
+    out_dir = embeddings_dir or (repo_root / ".codegraph")
+    store = EmbeddingStore(out_dir, dim=dim, backend=backend)
+    if force:
+        store.replace_all(rows)
+    else:
+        store.upsert(rows)
+
+    return IndexStats(
+        chunks_indexed=len(rows),
+        model=emb.model,
+        dim=dim,
+        backend=store.backend_name,
+        on_disk_bytes=store.size_bytes(),
+    )

codegraph/graph/builder.py

@@ -0,0 +1,212 @@
+"""Repo walker and incremental graph builder."""
+from __future__ import annotations
+
+import hashlib
+import logging
+import subprocess
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import pathspec
+
+# Ensure extractors register themselves.
+import codegraph.parsers.go
+import codegraph.parsers.python
+import codegraph.parsers.typescript  # noqa: F401
+from codegraph.config import CodegraphConfig
+from codegraph.graph.schema import Node, NodeKind, make_node_id
+from codegraph.graph.store_sqlite import SQLiteGraphStore
+from codegraph.parsers.base import get_extractor_for
+from codegraph.parsers.python import PythonExtractor
+
+logger = logging.getLogger(__name__)
+
+_BUILTIN_IGNORES = [
+    ".git", ".venv", "venv", "node_modules", ".codegraph",
+    "dist", "build", "__pycache__", ".next", ".pytest_cache",
+    ".mypy_cache", ".ruff_cache", ".tox", ".eggs", "*.egg-info",
+    ".DS_Store", "*.pyc", "*.pyo",
+]
+
+_IGNORE_DIRS: set[str] = {
+    ".git", ".venv", "venv", "node_modules", ".codegraph",
+    "dist", "build", "__pycache__", ".next", ".pytest_cache",
+    ".mypy_cache", ".ruff_cache", ".tox",
+}
+
+
+@dataclass
+class BuildStats:
+    files_scanned: int = 0
+    files_parsed: int = 0
+    nodes_added: int = 0
+    edges_added: int = 0
+    files_skipped: int = 0
+    errors: list[str] = field(default_factory=list)
+
+
+def _sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(65536), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _get_git_sha(repo_root: Path) -> str | None:
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "--short", "HEAD"],
+            cwd=repo_root,
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0:
+            return result.stdout.strip()
+    except Exception:
+        pass
+    return None
+
+
+class GraphBuilder:
+    def __init__(
+        self,
+        repo_root: Path,
+        store: SQLiteGraphStore,
+        ignore: list[str] | None = None,
+        config: CodegraphConfig | None = None,
+    ) -> None:
+        self._repo_root = repo_root
+        self._store = store
+        self._ignore = ignore or []
+        self._config = config or CodegraphConfig()
+        self._apply_config_to_extractors()
+
+    def _apply_config_to_extractors(self) -> None:
+        """Forward user dead-code patterns onto the singleton extractors."""
+        extra = tuple(self._config.dead_code.entry_point_decorators)
+        # PythonExtractor is registered as a singleton in the registry; we
+        # mutate its class attribute so subsequent parse_file calls pick up
+        # the user patterns.
+        PythonExtractor.extra_entry_point_decorators = extra
+
+    def build(self, incremental: bool = True) -> BuildStats:
+        stats = BuildStats()
+        patterns = _BUILTIN_IGNORES + self._ignore
+        spec = pathspec.PathSpec.from_lines("gitwildmatch", patterns)
+
+        files = list(self._walk_repo(spec))
+        stats.files_scanned = len(files)
+
+        for file_path in files:
+            rel = file_path.relative_to(self._repo_root).as_posix()
+            try:
+                content_hash = _sha256(file_path)
+
+                extractor = get_extractor_for(file_path)
+                language = extractor.language if extractor else "unknown"
+
+                file_node_id = make_node_id(NodeKind.FILE, rel, rel)
+                if incremental:
+                    existing = self._store.get_node(file_node_id)
+                    if existing and existing.content_hash == content_hash:
+                        stats.files_skipped += 1
+                        continue
+
+                self._store.delete_file(rel)
+
+                file_node = Node(
+                    id=file_node_id,
+                    kind=NodeKind.FILE,
+                    name=file_path.name,
+                    qualname=rel,
+                    file=rel,
+                    line_start=1,
+                    line_end=0,
+                    content_hash=content_hash,
+                    language=language,
+                    metadata={"size": file_path.stat().st_size},
+                )
+                self._store.upsert_node(file_node)
+                stats.nodes_added += 1
+
+                if extractor is not None:
+                    nodes, edges = extractor.parse_file(
+                        file_path, self._repo_root
+                    )
+                    self._store.upsert_nodes(nodes)
+                    self._store.upsert_edges(edges)
+                    stats.nodes_added += len(nodes)
+                    stats.edges_added += len(edges)
+                    stats.files_parsed += 1
+
+            except Exception as exc:
+                logger.warning("Error parsing %s: %s", rel, exc)
+                stats.errors.append(f"{rel}: {exc}")
+
+        now = datetime.now(tz=timezone.utc).isoformat()
+        self._store.set_meta("last_build_time", now)
+        git_sha = _get_git_sha(self._repo_root)
+        if git_sha:
+            self._store.set_meta("last_git_sha", git_sha)
+
+        # Best-effort cross-file resolution of unresolved CALLS/IMPORTS edges.
+        try:
+            from codegraph.resolve import resolve_unresolved_edges
+            rstats = resolve_unresolved_edges(self._store)
+            self._store.set_meta(
+                "last_resolve",
+                f"{rstats.resolved}/{rstats.inspected} resolved",
+            )
+        except Exception as exc:  # pragma: no cover - defensive
+            logger.warning("resolver failed: %s", exc)
+            stats.errors.append(f"resolver: {exc}")
+
+        # Architectural role classification (DF1.5): stamp HANDLER/SERVICE/
+        # COMPONENT/REPO onto FUNCTION/METHOD/CLASS nodes.
+        try:
+            from codegraph.analysis.roles import classify_roles
+            from codegraph.graph.store_networkx import to_digraph
+
+            graph = to_digraph(self._store)
+            count = classify_roles(graph)
+            if count:
+                updated: list[Node] = []
+                for nid, attrs in graph.nodes(data=True):
+                    metadata = attrs.get("metadata") or {}
+                    if not metadata.get("role"):
+                        continue
+                    existing = self._store.get_node(nid)
+                    if existing is None:
+                        continue
+                    existing.metadata["role"] = metadata["role"]
+                    updated.append(existing)
+                if updated:
+                    self._store.upsert_nodes(updated)
+            self._store.set_meta("last_roles", str(count))
+            logger.info("roles: %d nodes classified", count)
+        except Exception as exc:  # pragma: no cover - defensive
+            logger.warning("role classifier failed: %s", exc)
+            stats.errors.append(f"roles: {exc}")
+
+        return stats
+
+    def _walk_repo(self, spec: Any) -> list[Path]:
+        result: list[Path] = []
+        for file_path in sorted(self._repo_root.rglob("*")):
+            if not file_path.is_file():
+                continue
+            try:
+                rel = file_path.relative_to(self._repo_root).as_posix()
+            except ValueError:
+                continue
+            if spec.match_file(rel):
+                continue
+            parts = Path(rel).parts
+            if any(part in _IGNORE_DIRS for part in parts[:-1]):
+                continue
+            result.append(file_path)
+        return result

codegraph/graph/schema.py

@@ -0,0 +1,69 @@
+"""Graph schema: Node, Edge, and ID generation."""
+from __future__ import annotations
+
+import hashlib
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class NodeKind(str, Enum):
+    FILE = "FILE"
+    MODULE = "MODULE"
+    CLASS = "CLASS"
+    FUNCTION = "FUNCTION"
+    METHOD = "METHOD"
+    VARIABLE = "VARIABLE"
+    PARAMETER = "PARAMETER"
+    IMPORT = "IMPORT"
+    TEST = "TEST"
+
+
+class EdgeKind(str, Enum):
+    DEFINED_IN = "DEFINED_IN"
+    IMPORTS = "IMPORTS"
+    CALLS = "CALLS"
+    INHERITS = "INHERITS"
+    IMPLEMENTS = "IMPLEMENTS"
+    READS = "READS"
+    WRITES = "WRITES"
+    RETURNS = "RETURNS"
+    PARAM_OF = "PARAM_OF"
+    TESTED_BY = "TESTED_BY"
+    # v0.2 cross-stack data-flow edges (populated by DF1 / DF2 extractors).
+    # Reserved here so DF1/DF2 agents don't both edit this enum in parallel.
+    ROUTE = "ROUTE"               # HANDLER → URL pattern (DF1, FastAPI/Flask)
+    READS_FROM = "READS_FROM"     # function → SQLAlchemy model on read (DF1)
+    WRITES_TO = "WRITES_TO"       # function → SQLAlchemy model on write (DF1)
+    FETCH_CALL = "FETCH_CALL"     # frontend call site → URL string (DF2, fetch/axios)
+
+
+class Node(BaseModel):
+    id: str
+    kind: NodeKind
+    name: str
+    qualname: str
+    file: str
+    line_start: int
+    line_end: int
+    signature: str | None = None
+    docstring: str | None = None
+    content_hash: str | None = None
+    language: str
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class Edge(BaseModel):
+    src: str
+    dst: str
+    kind: EdgeKind
+    file: str | None = None
+    line: int | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+def make_node_id(kind: NodeKind, qualname: str, file: str) -> str:
+    """Stable BLAKE2b-128 hex hash of (kind, qualname, file)."""
+    data = f"{kind.value}:{qualname}:{file}".encode()
+    return hashlib.blake2b(data, digest_size=16).hexdigest()

codegraph/graph/store_networkx.py

@@ -0,0 +1,55 @@
+"""NetworkX adapter for the SQLiteGraphStore."""
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import cast
+
+import networkx as nx
+
+from codegraph.graph.schema import EdgeKind
+from codegraph.graph.store_sqlite import SQLiteGraphStore
+
+
+def to_digraph(store: SQLiteGraphStore) -> nx.MultiDiGraph:
+    g: nx.MultiDiGraph = nx.MultiDiGraph()
+    for node in store.iter_nodes():
+        g.add_node(node.id, **node.model_dump(mode="json"))
+    for edge in store.iter_edges():
+        g.add_edge(edge.src, edge.dst, key=edge.kind.value, **edge.model_dump(mode="json"))
+    return g
+
+
+def subgraph_around(
+    g: nx.MultiDiGraph,
+    node_id: str,
+    depth: int,
+    direction: str = "both",
+    edge_kinds: Iterable[EdgeKind] | None = None,
+) -> nx.MultiDiGraph:
+    """Return a MultiDiGraph of nodes within `depth` BFS hops from node_id."""
+    allowed_kinds: set[str] | None = (
+        {k.value for k in edge_kinds} if edge_kinds is not None else None
+    )
+    visited: set[str] = set()
+    frontier: set[str] = {node_id}
+    for _ in range(depth):
+        next_frontier: set[str] = set()
+        for n in frontier:
+            if n not in g:
+                continue
+            neighbors: list[str] = []
+            if direction in ("out", "both"):
+                for _src, dst, data in g.out_edges(n, data=True):
+                    if allowed_kinds is None or data.get("kind") in allowed_kinds:
+                        neighbors.append(dst)
+            if direction in ("in", "both"):
+                for src, _dst, data in g.in_edges(n, data=True):
+                    if allowed_kinds is None or data.get("kind") in allowed_kinds:
+                        neighbors.append(src)
+            for nb in neighbors:
+                if nb not in visited and nb not in frontier:
+                    next_frontier.add(nb)
+        visited.update(frontier)
+        frontier = next_frontier - visited
+    visited.update(frontier)
+    return cast(nx.MultiDiGraph, g.subgraph(visited).copy())